#!/usr/bin/env nomsu -V4.8.8.6
#
    This file contains compile-time actions that define basic control flow structures
    like "if" statements and loops.

use "core/metaprogramming.nom"
use "core/text.nom"
use "core/operators.nom"
use "core/errors.nom"

# No-Op
test: do nothing
compile [do nothing] to (Lua "")

# Conditionals
test:
    if (no):
        barf "conditional fail"
compile [if %condition %if_body] to:
    %lua = (Lua "if ")
    %lua::append (%condition as lua expr)
    %lua::append " then\n    "
    %lua::append (%if_body as lua statements)
    %lua::append "\nend"
    return %lua

test:
    unless (yes):
        barf "conditional fail"
parse [unless %condition %unless_body] as (if (not %condition) %unless_body)
compile [..]
    if %condition %if_body else %else_body, unless %condition %else_body else %if_body
..to:
    %lua = (Lua "if ")
    %lua::append (%condition as lua expr)
    %lua::append " then\n    "
    %lua::append (%if_body as lua statements)
    %lua::append "\nelse\n    "
    %lua::append (%else_body as lua statements)
    %lua::append "\nend"
    return %lua

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

# 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)
compile [..]
    %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
..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:yes, List:yes, Dict:yes, Number:yes}.(%when_true_expr.type):
        return (..)
            Lua value "\
                ..(\(%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 value "\
                ..((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)
compile [=== %label ===, --- %label ---, *** %label ***] to (..)
    Lua "::label_\((%label.stub if (%label.type == "Action") else %label) as lua identifier)::"

compile [go to %label] to (..)
    Lua "goto label_\((%label.stub if (%label.type == "Action") else %label) as lua identifier)"

# Basic loop control
compile [do next] to (Lua "goto continue")
compile [stop] to (Lua "break")

# While loops
test:
    %x = 0
    repeat while (%x < 10): %x += 1
    assume (%x == 10)
    repeat while (%x < 20): stop
    repeat while (%x < 20): stop repeating
    assume (%x == 10)
    repeat while (%x < 20):
        %x += 1
        if (yes): do next
        barf "Failed to 'do next'"
    
    assume (%x == 20)
    repeat while (%x < 30):
        %x += 1
        if (yes): do next repeat
        barf "Failed to 'do next repeat'"
    
    assume (%x == 30)
compile [do next repeat] to (Lua "goto continue_repeat")
compile [stop repeating] to (Lua "goto stop_repeat")
compile [repeat while %condition %body] to:
    %lua = (..)
        Lua "\
            ..while \(%condition as lua expr) do
                \(%body as lua statements)"
    
    if (%body has subtree \(do next)):
        %lua::append "\n    ::continue::"
    if (%body has subtree \(do next repeat)):
        %lua::append "\n    ::continue_repeat::"
    %lua::append "\nend --while-loop"
    if (%body has subtree \(stop repeating)):
        %inner_lua = %lua
        %lua = (Lua "do -- scope of 'stop repeating' label\n    ")
        %lua::append %inner_lua
        %lua::append "\
                ..
                    ::stop_repeat::
                end -- end of 'stop repeating' label scope"
    
    return %lua

parse [repeat %body] as (repeat while (yes) %body)
parse [repeat until %condition %body] as (repeat while (not %condition) %body)

test:
    %x = 0
    repeat 10 times: %x += 1
    assume (%x == 10)
compile [repeat %n times %body] to:
    define mangler
    %lua = (..)
        Lua "for \(mangle "i")=1,\(%n as lua expr) do\n    "
    %lua::append (%body as lua statements)
    
    if (%body has subtree \(do next)):
        %lua::append "\n    ::continue::"
    if (%body has subtree \(do next repeat)):
        %lua::append "\n    ::continue_repeat::"
    %lua::append "\nend --numeric for-loop"
    if (%body has subtree \(stop repeating)):
        %inner_lua = %lua
        %lua = (Lua "do -- scope of 'stop repeating' label\n    ")
        %lua::append %inner_lua
        %lua::append "\
                ..
                    ::stop_repeat::
                end -- end of 'stop repeating' label scope"
    
    return %lua

# For loop control flow
compile [stop %var] to (..)
    Lua "goto stop_\((%var.stub if (%var.type == "Action") else %var) as lua identifier)"
compile [do next %var] to (..)
    Lua "goto continue_\((%var.stub if (%var.type == "Action") else %var) as lua identifier)"
compile [===stop %var ===, ---stop %var ---, ***stop %var ***] to (..)
    Lua "::stop_\((%var.stub if (%var.type == "Action") else %var) as lua identifier)::"

compile [===next %var ===, ---next %var ---, ***next %var ***] to (..)
    Lua "::continue_\((%var.stub if (%var.type == "Action") else %var) as lua identifier)::"

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

test:
    %nums = []
    for %x in 1 to 5: %nums::add %x
    assume (%nums == [1, 2, 3, 4, 5])
    %nums = []
    for %x in 1 to 5 via 2: %nums::add %x
    assume (%nums == [1, 3, 5])
    %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])

# Numeric range for loops
compile [..]
    for %var in %start to %stop by %step %body
    for %var in %start to %stop via %step %body
..to:
    # This uses Lua's approach of only allowing loop-scoped variables in a loop
    unless (%var.type is "Var"):
        compile error at %var "Expected a variable here, not a \(%var.type)"
    %lua = (..)
        Lua "\
            ..for \(%var as lua expr)=\(%start as lua expr),\(%stop as lua expr),\(..)
                %step as lua expr
            .. do"
    %lua::append "\n    "
    %lua::append (%body as lua statements)
    
    if (%body has subtree \(do next)):
        %lua::append "\n    ::continue::"
    if (%body has subtree \(do next %var)):
        %lua::append "\n    "
        %lua::append (compile as (===next %var ===))
    %lua::append "\nend --numeric for-loop"
    if (%body has subtree \(stop %var)):
        %inner_lua = %lua
        %lua = (Lua "do -- scope for stopping for-loop\n    ")
        %lua::append %inner_lua
        %lua::append "\n    "
        %lua::append (compile as (===stop %var ===))
        %lua::append "\nend -- end of scope for stopping for-loop"
    
    return %lua

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

parse [for %var in %start to %stop %body] as (..)
    for %var in %start to %stop via 1 %body

test:
    %a = [10, 20, 30, 40, 50]
    %b = []
    for %x in %a: %b::add %x
    assume (%a == %b)
    %b = []
    for %x in %a:
        if (%x == 10): do next %x
        if (%x == 50): stop %x
        %b::add %x
    
    assume (%b == [20, 30, 40])

# For-each loop (lua's "ipairs()")
compile [for %var in %iterable %body] to:
    define mangler
    # This uses Lua's approach of only allowing loop-scoped variables in a loop
    %lua = (..)
        Lua "for \(mangle "i"),\(%var as lua identifier) in ipairs(\(%iterable as lua expr)) do\n    "
    %lua::append (%body as lua statements)
    
    if (%body has subtree \(do next)):
        %lua::append "\n    ::continue::"
    if (%body has subtree \(do next %var)):
        %lua::append "\n    "
        %lua::append (compile as (===next %var ===))
    %lua::append "\nend --foreach-loop"
    if (%body has subtree \(stop %var)):
        %inner_lua = %lua
        %lua = (Lua "do -- scope for stopping for-loop\n    ")
        %lua::append %inner_lua
        %lua::append "\n    "
        %lua::append (compile as (===stop %var ===))
        %lua::append "\nend -- end of scope for stopping for-loop"
    
    return %lua

# TODO: reduce code duplication
compile [for %var in %iterable at %i %body] to:
    # This uses Lua's approach of only allowing loop-scoped variables in a loop
    %lua = (..)
        Lua "for \(%i as lua identifier),\(%var as lua identifier) in ipairs(\(%iterable as lua expr)) do\n    "
    %lua::append (%body as lua statements)
    
    if (%body has subtree \(do next)):
        %lua::append "\n    ::continue::"
    if (%body has subtree \(do next %var)):
        %lua::append "\n    "
        %lua::append (compile as (===next %var ===))
    %lua::append "\nend --foreach-loop"
    if (%body has subtree \(stop %var)):
        %inner_lua = %lua
        %lua = (Lua "do -- scope for stopping for-loop\n    ")
        %lua::append %inner_lua
        %lua::append "\n    "
        %lua::append (compile as (===stop %var ===))
        %lua::append "\nend -- end of scope for stopping for-loop"
    
    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"])

# Dict iteration (lua's "pairs()")
compile [..]
    for %key = %value in %iterable %body, for %key %value in %iterable %body
..to:
    # This uses Lua's approach of only allowing loop-scoped variables in a loop
    unless (%key.type is "Var"):
        compile error at %key "Expected a variable here, not a \(%key.type)"
    unless (%value.type is "Var"):
        compile error at %value "Expected a variable here, not a \(%value.type)"
    %lua = (..)
        Lua "\
            ..for \(%key as lua identifier),\(%value as lua identifier) in pairs(\(..)
                %iterable as lua expr
            ..) do"
    %lua::append "\n    "
    %lua::append (%body as lua statements)
    
    if (%body has subtree \(do next)):
        %lua::append "\n    ::continue::"
    if (%body has subtree \(do next %key)):
        %lua::append "\n    "
        %lua::append (compile as (===next %key ===))
    if (%body has subtree \(do next %value)):
        %lua::append "\n    "
        %lua::append (compile as (===next %value ===))
    %lua::append "\nend --foreach-loop"
    %stop_labels = (Lua "")
    if (%body has subtree \(stop %key)):
        %stop_labels::append "\n"
        %stop_labels::append (compile as (===stop %key ===))
    if (%body has subtree \(stop %value)):
        %stop_labels::append "\n"
        %stop_labels::append (compile as (===stop %value ===))
    if ((size of "\%stop_labels") > 0):
        %inner_lua = %lua
        %lua = (Lua "do -- scope for stopping for % = % loop\n    ")
        %lua::append %inner_lua
        %lua::append %stop_labels
        %lua::append "\nend"
    
    return %lua

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

test:
    if:
        (1 == 2) (100 < 0):
            barf "bad conditional"
        (1 == 0) (1 == 1) %not_a_variable.x: do nothing
        (1 == 1):
            barf "bad conditional"
        (1 == 2):
            barf "bad conditional"
        else:
            barf "bad conditional"

# Multi-branch conditional (if..elseif..else)
compile [if %body, when %body] to:
    %code = (Lua "")
    %clause = "if"
    %else_allowed = (yes)
    unless (%body.type is "Block"):
        compile error at %body "'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 is "Action") and ((size of %line) >= 2)) and (..)
                %line.(size of %line) is "Block" syntax tree
        ..:
            compile error at %line "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.(size of %line)
        if ((%line.1 is "else") and ((size of %line) == 2)):
            unless %else_allowed:
                compile error at %line "You can't have two 'else' blocks."
                ..hint "Merge all of the 'else' blocks together."
            unless ((size of "\%code") > 0):
                compile error at %line "\
                    ..You can't have an 'else' block without a preceeding 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::append "\nelse\n    "
            %code::append (%action as lua statements)
            %else_allowed = (no)
        ..else:
            %code::append %clause
            %code::append " "
            for %i in 1 to ((size of %line) - 1):
                if (%i > 1):
                    %code::append " or "
                %code::append (%line.%i as lua expr)
            
            %code::append " then\n    "
            %code::append (%action as lua statements)
            %clause = "\nelseif"
    
    if ((size of "\%code") == 0):
        compile error at %body "'if' block has an empty body."
        ..hint "This means nothing would happen, so the 'if' block should be deleted."
    %code::append "\nend --when"
    return %code

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

# Switch statement
compile [if %branch_value is %body, when %branch_value is %body] to:
    %code = (Lua "")
    %clause = "if"
    %else_allowed = (yes)
    define mangler
    unless (%body.type is "Block"):
        compile error at %body "'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 is "Action") and ((size of %line) >= 2)) and (..)
                %line.(size of %line) is "Block" syntax tree
        ..:
            compile error at %line "Invalid line for 'if' block."
            ..hint "Each line should contain expressions \
                ..followed by a block, or "else" followed by a block"
        
        %action = %line.(size of %line)
        if ((%line.1 is "else") and ((size of %line) == 2)):
            unless %else_allowed:
                compile error at %line "You can't have two 'else' blocks."
                ..hint "Merge all of the 'else' blocks together."
            unless ((size of "\%code") > 0):
                compile error at %line "\
                    ..You can't have an 'else' block without a preceeding 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::append "\nelse\n    "
            %code::append (%action as lua statements)
            %else_allowed = (no)
        ..else:
            %code::append %clause
            %code::append " "
            for %i in 1 to ((size of %line) - 1):
                if (%i > 1):
                    %code::append " or "
                %code::append "\(mangle "branch value") == "
                %code::append (%line.%i as lua expr)
            
            %code::append " then\n    "
            %code::append (%action as lua statements)
            %clause = "\nelseif"
    
    if ((size of "\%code") == 0):
        compile error at %body "'if' block has an empty body."
        ..hint "This means nothing would happen, so the 'if' block should be deleted."
    %code::append "\nend --when"
    %lua = (..)
        Lua "\
            ..do --if % is...
                local \(mangle "branch value") = "
    %lua::append (%branch_value as lua expr)
    %lua::append "\n    "
    %lua::append %code
    %lua::append "\nend --if % is..."
    return %lua

# Do/finally
compile [do %action] to:
    %lua = (Lua "do\n    ")
    %lua::append (%action as lua statements)
    %lua::append "\nend -- do"
    return %lua

test:
    %d = {}
    try:
        do:
            %d.x = "bad"
            barf
        ..then always: %d.x = "good"
    ..and if it barfs: do nothing
    
    assume (%d.x == "good")
compile [do %action then always %final_action] to:
    define mangler
    %lua = (..)
        Lua "\
            ..do
                local \(mangle "fell_through") = false
                local \(mangle "ok"), \(mangle "ret") = pcall(function()"
    %lua::append "\n        "
    %lua::append (%action as lua statements)
    %lua::append "\
        ..        \(mangle "fell_through") = true
            end)"
    %lua::append "\n    "
    %lua::append (%final_action as lua statements)
    %lua::append "\
        ..    if not \(mangle "ok") then error(ret, 0) end
            if not \(mangle "fell_through") then return ret end
        end"
    return %lua

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

# Inline thunk:
compile [result of %body] to (Lua value "\(compile as ([] -> %body))()")

test:
    %t = [1, [2, [[3], 4], 5, [[[6]]]]]
    %flat = []
    for % in recursive %t:
        if ((lua type of %) is "table"):
            for %2 in %: recurse % on %2
        ..else: %flat::add %
    
    assume (sorted %flat) == [1, 2, 3, 4, 5, 6]

# Recurion control flow
compile [for %var in recursive %structure %body] to:
    with local compile actions:
        define mangler
        compile [recurse %v on %x] to (..)
            Lua "table.insert(\(mangle "stack \(%v.1)"), \(%x as lua expr))"
        %lua = (..)
            Lua "\
                ..do
                    local \(mangle "stack \(%var.1)") = _List{\(%structure as lua expr)}
                    while #\(mangle "stack \(%var.1)") > 0 do
                        \(%var as lua expr) = table.remove(\(mangle "stack \(%var.1)"), 1)"
        %lua::append "\n        "
        %lua::append (%body as lua statements)
        
        if (%body has subtree \(do next)):
            %lua::append "\n        ::continue::"
        if (%body has subtree \(do next %var)):
            %lua::append "\n        \(compile as (===next %var ===))"
        %lua::append "\n    end -- Recursive loop"
        if (%body has subtree \(stop %var)):
            %lua::append "\n        \(compile as (===stop %var ===))"
        %lua::append "\nend -- Recursive scope"
        return %lua