#!/usr/bin/env nomsu -V7.0.0
###
    This file contains compile-time actions that define basic control flow structures
    like "if" statements and loops.
    
use "core/metaprogramming"
use "core/operators"

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

### No-Op
test:
    do nothing
(do nothing) compiles to ""

### Conditionals
test:
    if (no):
        fail "conditional fail"

(if $condition $if_body) compiles to ("
    if \($condition as lua expr) then
        \($if_body as lua)
    end
")

test:
    unless (yes):
        fail "conditional fail"

(unless $condition $unless_body) parses as (if (not $condition) $unless_body)
[
    if $condition $if_body else $else_body, unless $condition $else_body else $if_body
] all compile to ("
    if \($condition as lua expr) then
        \($if_body as lua)
    else
        \($else_body as lua)
    end
")

(else $) compiles to:
    at (this tree) fail ("
        Compile error: This 'else' is not connected to any 'if' or 'unless' condition.
        Hint: You should probably have a ".." in front of the "else", to indicate that it's attached \
        ..to the previous condition.
    ")

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

### Conditional expression (ternary operator)
###   Note: this uses a function instead of "(condition and if_expr or else_expr)"
    because that breaks if $if_expr is falsey, e.g. "x < 5 and false or 99"
test:
    assume ((1 if (yes) else 2) == 1)
    assume ((1 if (no) else 2) == 2)

[
    $when_true_expr if $condition else $when_false_expr
    $when_true_expr if $condition otherwise $when_false_expr
    $when_false_expr unless $condition else $when_true_expr
    $when_false_expr unless $condition then $when_true_expr
] all compile to:
    ###   If $when_true_expr is guaranteed to be truthy, we can use Lua's idiomatic
        equivalent of a conditional expression: (cond and if_true or if_false)
    if {.Text, .List, .Dict, .Number}.($when_true_expr.type):
        return Lua ("
            (\($condition as lua expr) and \($when_true_expr as lua expr) or \
            ..\($when_false_expr as lua expr))
        ")
    ..else:
        ###   Otherwise, need to do an anonymous inline function (yuck, too bad lua 
            doesn't have a proper ternary operator!)
            To see why this is necessary consider: (random()<.5 and false or 99)
        return Lua ("
            ((function()
                if \($condition as lua expr) then
                    return \($when_true_expr as lua expr)
                else
                    return \($when_false_expr as lua expr)
                end
            end)())
        ")

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

### GOTOs
test:
    $i = 0
    --- $loop ---
    $i += 1
    unless ($i == 10):
        go to $loop
    assume ($i == 10)
    --- (Loop) ---
    $i -= 1
    unless ($i == 0):
        go to (Loop)
    assume ($i == 0)

(--- $label ---) compiles to ("
    ::label_\(
        ($label.stub, as lua id) if ($label.type == "Action") else
            $label as lua identifier
    )::
")

(go to $label) compiles to ("
    goto label_\(
        ($label.stub, as lua id) if ($label.type == "Action") else
            $label as lua identifier
    )
")

### Basic loop control
(stop $var) compiles to:
    if $var:
        return Lua "goto stop_\($var as lua identifier)"
    ..else:
        return Lua "break"

(do next $var) compiles to:
    if $var:
        return Lua "goto continue_\($var as lua identifier)"
    ..else:
        return Lua "goto continue"

(---stop $var ---) compiles to "::stop_\($var as lua identifier)::"
(---next $var ---) compiles to "::continue_\($var as lua identifier)::"

### While loops
test:
    $x = 0
    repeat while ($x < 10): $x += 1
    assume ($x == 10)
    repeat while ($x < 20): stop
    assume ($x == 10)
    repeat while ($x < 20):
        $x += 1
        if (yes):
            do next
        fail "Failed to 'do next'"
    assume ($x == 20)

(repeat while $condition $body) compiles to:
    $lua =
        Lua ("
            while \($condition as lua expr) do
                \($body as lua)
        ")
    
    if ($body, contains \(do next)):
        $lua, add "\n    ::continue::"
    
    $lua, add "\nend --while-loop"
    return $lua

(repeat $body) parses as (repeat while (yes) $body)
(repeat until $condition $body) parses as (repeat while (not $condition) $body)

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

### For-each loop (lua's "ipairs()")
(for $var in $iterable $body) compiles to:
    unless $var:
        at (this tree) fail "No var here"
    
    ### This uses Lua's approach of only allowing loop-scoped variables in a loop
    if (($var.type == "Action") and ($var.stub == "1 =")):
        [$key, $value] = [$var.1, $var.3]
    ..else:
        [$key, $value] = [nil, $var]
    
    unless $value:
        at (this tree) fail "No value here"
    
    ### Numeric loop:
    if (($iterable.type == "Action") and (($iterable, get stub) == "1 to")):
        [$start, $stop] = [$iterable.1, $iterable.3]
        $loop =
            Lua ("
                local _start = \($start as lua expr)
                for \($value as lua identifier)=_start,\($stop as lua expr) do
            ")
        
        if $key:
            $loop, add ("
                
                    local \($key as lua identifier) = \($value as lua identifier) - _start + 1;
            ")
        
        go to (loop set)
    
    ### Numeric loop with step:
    if (($iterable.type == "Action") and (($iterable, get stub) == "1 to 2 by")):
        [$start, $stop, $step] = [$iterable.1, $iterable.3, $iterable.5]
        $loop =
            Lua ("
                local _start, _step = \($start as lua expr), \($step as lua expr);
                for \($value as lua identifier)=_start,\($stop as lua expr),_step do
            ")
        
        if $key:
            $loop, add ("
                
                    local \($key as lua identifier) = (\($value as lua identifier) - _start)/_step + 1
            ")
        
        go to (loop set)
    
    ### for $ in (...):
    if $key:
        $loop =
            Lua ("
                for \($key as lua identifier),\($value as lua identifier) in pairs(\
                ..\($iterable as lua expr)) do
            ")
    ..else:
        $loop =
            Lua "for _i,\($value as lua identifier) in _ipairs(\($iterable as lua expr)) do"
    
    --- (loop set) ---
    
    ### TODO: don't always wrap in block
    $lua =
        Lua ("
            do -- for-loop
                \$loop
                    \;
        ")
    $lua, add ($body as lua)
    if ($body, contains \(do next)):
        $lua, add "\n        ::continue::"
    
    if ($key and ($body, contains \(do next \$key))):
        $lua, add "\n        " (\(---next \$key ---) as lua)
    
    if ($body, contains \(do next \$value)):
        $lua, add "\n        " (\(---next \$value ---) as lua)
    
    $lua, add "\n    end"
    if ($key and ($body, contains \(stop \$key))):
        $lua, add "\n    " (\(---stop \$key ---) as lua)
    
    if ($body, contains \(stop \$value)):
        $lua, add "\n    " (\(---stop \$value ---) as lua)
    
    $lua, add "\nend -- for-loop"
    $lua, remove free vars
        [($value as lua identifier, text), $key and ($key as lua identifier, text)]
    return $lua

test:
    $d = {.a = 10, .b = 20, .c = 30, .d = 40, .e = 50}
    $result = []
    for ($k = $v) in $d:
        if ($k == "a"):
            do next $k
        
        if ($v == 20):
            do next $v
        
        $result, add "\$k = \$v"
    assume (($result sorted) == ["c = 30", "d = 40", "e = 50"])

### Numeric range for loops
test:
    assume ([: for $ in (1 to 5): add $] == [1, 2, 3, 4, 5])
    assume ([: for $ in (1 to 5 by 2): add $] == [1, 3, 5])
    assume ([: for $ in (5 to 1): add $] == [])
    $nums = []
    for $outer in (1 to 100):
        for $inner in ($outer to ($outer + 2)):
            if ($inner == 2):
                $nums, add -2
                do next $inner
            $nums, add $inner
            if ($inner == 5):
                stop $outer
    
    assume ($nums == [1, -2, 3, -2, 3, 4, 3, 4, 5])

### repeat $n times is a shorthand:
test:
    $x = 0
    repeat 5 times:
        $x += 1
    assume $x == 5
(repeat $n times $body) parses as (for (=lua "_i") in (1 to $n by 1) $body)

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

test:
    when:
        (1 == 2) (100 < 0):
            fail "bad conditional"
        (1 == 0) (1 == 1) $not_a_variable.x: do nothing
        (1 == 1):
            fail "bad conditional"
        
        (1 == 2):
            fail "bad conditional"
        
        else:
            fail "bad conditional"

### Multi-branch conditional (if..elseif..else)
(when $body) compiles to:
    $code = (Lua "")
    $clause = "if"
    $else_allowed = (yes)
    unless ($body.type == "Block"):
        at $body fail ("
            Compile error: 'if' expected a Block, but got a \($body.type).
            Hint: Perhaps you forgot to put a ':' after 'if'?
        ")
    
    for $line in $body:
        unless
            (($line.type == "Action") and (#$line >= 2)) and
                $line.(#$line) is "Block" syntax tree
        ..:
            at $line fail ("
                Compile error: Invalid line for the body of an 'if' block.
                Hint: Each line should contain one or more conditional expressions followed by a block, \
                ..or "else" followed by a block.
            ")
        $action = $line.(#$line)
        if (($line.1 == "else") and (#$line == 2)):
            unless $else_allowed:
                at $line fail ("
                    Compile error: You can't have two 'else' blocks.
                    Hint: Merge all of the 'else' blocks together.
                ")
            
            unless (#"\$code" > 0):
                at $line fail ("
                    Compile error: You can't have an 'else' block without a preceding condition.
                    Hint: If you want the code in this block to always execute, you don't need a conditional \
                    ..block around it. Otherwise, make sure the 'else' block comes last.
                ")
            
            $code, add "\nelse\n    " ($action as lua)
            $else_allowed = (no)
        ..else:
            $code, add $clause " "
            for $i in (1 to (#$line - 1)):
                if ($i > 1):
                    $code, add " or "
                $code, add ($line.$i as lua expr)
            $code, add " then\n    " ($action as lua)
            $clause = "\nelseif"
    
    if (#"\$code" == 0):
        at $body fail ("
            Compile error: 'if' block has an empty body.
            Hint: This means nothing would happen, so the 'if' block should be deleted.
        ")
    
    $code, add "\nend --when"
    return $code

test:
    if 5 is:
        1 2 3:
            fail "bad switch statement"
        
        4 5:
            do nothing
        
        5 6:
            fail "bad switch statement"
        
        else:
            fail "bad switch statement"

### Switch statement
[if $branch_value is $body, when $branch_value is $body] all compile to:
    $code = (Lua "")
    $clause = "if"
    $else_allowed = (yes)
    define mangler
    unless ($body.type == "Block"):
        at $body fail ("
            Compile error: 'if' expected a Block, but got a \($body.type).
            Hint: Perhaps you forgot to put a ':' after the 'is'?
        ")
    
    for $line in $body:
        unless
            (($line.type == "Action") and (#$line >= 2)) and
                $line.(#$line) is "Block" syntax tree
        ..:
            at $line fail ("
                Compile error: Invalid line for 'if' block.
                Hint: Each line should contain expressions followed by a block, or "else" followed by a block.
            ")
        $action = $line.(#$line)
        if (($line.1 == "else") and (#$line == 2)):
            unless $else_allowed:
                at $line fail ("
                    Compile error: You can't have two 'else' blocks.
                    Hint: Merge all of the 'else' blocks together.
                ")
            
            unless (#"\$code" > 0):
                at $line fail ("
                    Compile error: You can't have an 'else' block without a preceding condition.
                    Hint: If you want the code in this block to always execute, you don't need a conditional \
                    ..block around it. Otherwise, make sure the 'else' block comes last.
                ")
            
            $code, add "\nelse\n    " ($action as lua)
            $else_allowed = (no)
        ..else:
            $code, add $clause " "
            for $i in (1 to (#$line - 1)):
                if ($i > 1):
                    $code, add " or "
                $code, add "\(mangle "branch value") == " ($line.$i as lua expr)
            $code, add " then\n    " ($action as lua)
            $clause = "\nelseif"
    
    if (#"\$code" == 0):
        at $body fail ("
            Compile error: 'if' block has an empty body.
            Hint: This means nothing would happen, so the 'if' block should be deleted.
        ")
    
    $code, add "\nend --when"
    return Lua ("
        do --if $ is...
            local \(mangle "branch value") = \($branch_value as lua expr)
            \$code
        end -- if $ is...
    ")

### Do/finally
(do $action) compiles to ("
    do
        \($action as lua)
    end -- do
")

test:
    assume ((result of: return 99) == 99)

### Inline thunk:
(result of $body) compiles to "\(\(->(\$body)) as lua)()"
test:
    $t = [1, [2, [[3], 4], 5, [[[6]]]]]
    $flat = []
    for $ in recursive $t:
        if ((lua type of $) == "table"):
            for $2 in $:
                recurse $ on $2
        ..else:
            $flat, add $
    
    assume (sorted $flat) == [1, 2, 3, 4, 5, 6]

### Recurion control flow
(recurse $v on $x) compiles to
    Lua "table.insert(_stack_\($v as lua expr), \($x as lua expr))"

(for $var in recursive $structure $body) compiles to:
    $lua =
        Lua ("
            do
                local _stack_\($var as lua expr) = a_List{\($structure as lua expr)}
                while #_stack_\($var as lua expr) > 0 do
                    \($var as lua expr) = table.remove(_stack_\($var as lua expr), 1)
                    \($body as lua)
        ")
    
    if ($body, contains \(do next)):
        $lua, add "\n        ::continue::"
    
    if ($body, contains \(do next \$var)):
        $lua, add "\n        \(\(---next \$var ---) as lua)"
    
    $lua, add "\n    end -- Recursive loop"
    if ($body, contains \(stop \$var)):
        $lua, add "\n        \(\(---stop \$var ---) as lua)"
    $lua, add "\nend -- Recursive scope"
    return $lua