#..
    This file contains compile-time actions that define basic control flow structures
    like "if" statements and loops.

use "lib/metaprogramming.nom"
use "lib/text.nom"
use "lib/operators.nom"

# No-Op
immediately
    compile [do nothing] to code ""

# Return
immediately
    compile [return] to code "do return; end"
    compile [return %return_value] to code "do return \(%return_value as lua); end"

# Conditionals
immediately
    compile [if %condition %if_body] to code ".."
        if \(%condition as lua) then
        \(%if_body as lua statements)
        end --end if
    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 code ".."
        if \(%condition as lua) then
        \(%if_body as lua statements)
        else
        \(%else_body as lua statements)
        end --end if

# 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"
immediately
    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: (%when_true_expr's "type") in {Text:yes, List:yes, Dict:yes, Number:yes}
            return "(\(%condition as lua) and \(%when_true_expr as lua) or \(%when_false_expr as lua))"
        ..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 ".."
                (function()
                    if \(%condition as lua) then
                        return \(%when_true_expr as lua);
                    else
                        return \(%when_false_expr as lua);
                    end
                end)()


# GOTOs
immediately
    compile [=== %label ===, --- %label ---] to code ".."
        ::label_\(nomsu "var_to_lua_identifier" [%label])::;
    compile [go to %label] to code ".."
        goto label_\(nomsu "var_to_lua_identifier" [%label]);

# Basic loop control
immediately
    compile [do next] to code "continue;"
    compile [stop] to code "break;"

# Helper function
immediately
    compile [for subtree %subtree where %condition in %tree %body] to code ".."
        for \(%subtree as lua) in coroutine.wrap(function() nomsu:walk_tree(\(%tree as lua)) end) do
            if type(\(%subtree as lua)) == 'table' and \(%subtree as lua).type then
                if \(%condition as lua) then
                    \(%body as lua statements)
                end
            end
        end

# While loops
immediately
    compile [do next repetition] to code "goto continue_repeat;"
    compile [stop repeating] to code "goto stop_repeat;"
    compile [repeat while %condition %body] to code
        %continue_labels <- ""
        for subtree % where
            ((%'s "type") = "FunctionCall") and ((%'s "stub") == "do next repetition")
        ..in %body
            %continue_labels <- "\n::continue_repeat::;"
            stop
        %code <- ".."
            while \(%condition as lua) do
            \(%body as lua statements)\
            ..\%continue_labels
            end --while-loop
        for subtree % where
            ((%'s "type") = "FunctionCall") and ((%'s "stub") == "stop repeating")
        ..in %body
            %code <- ".."
                do -- scope of "stop repeating" label
                \%code
                ::stop_repeat::;
                end -- end of "stop repeating" label scope
            %continue_labels <- "\n::continue_repeat::;"
            stop
        return %code
    parse [repeat %body] as: repeat while (yes) %body
    parse [repeat until %condition %body] as: repeat while (not %condition) %body

    compile [..]
        repeat %n times %body
    ..to code
        %continue_labels <- ""
        for subtree % where
            ((%'s "type") = "FunctionCall") and ((%'s "stub") == "do next repetition")
        ..in %body
            %continue_labels <- "\n::continue_repeat::;"
            stop
        # This trashes the loop variables, just like in Python.
        %code <- ".."
            for i=1,\(%n as lua) do
            \(%body as lua statements)\
            ..\%continue_labels
            end --numeric for-loop
        %stop_labels <- ""
        for subtree % where
            ((%'s "type") = "FunctionCall") and ((%'s "stub") == "stop repeating")
        ..in %body
            %code <- ".."
                do -- scope of "stop repeating" label
                \%code
                ::stop_repeat::;
                end -- end of "stop repeating" label scope
            %continue_labels <- "\n::continue_repeat::;"
            stop
        return %code

# For loop control flow:
immediately
    compile [stop %var] to code ".."
        goto stop_\(nomsu "var_to_lua_identifier" [%var]);
    compile [do next %var] to code ".."
        goto continue_\(nomsu "var_to_lua_identifier" [%var]);

# Numeric range for loops
immediately
    compile [..]
        for %var from %start to %stop by %step %body
        for %var from %start to %stop via %step %body
    ..to code
        %continue_labels <- ""
        for subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") == "do next %") and 
                    ((3rd in (%'s "value"))'s "src") == (%var's "src")
        ..in %body
            %continue_labels <- "\n::continue_\(nomsu "var_to_lua_identifier" [%var])::;"
            stop

        if: (%var's "type") is "Var"
            %loop_var <- (%var as lua)
            %loop_var_shim <- ""
        ..else
            %loop_var <- "i"
            %loop_var_shim <- "\n\(%var as lua) = i;"
        # This trashes the loop variables, just like in Python.
        %code <- ".."
            for \(%loop_var)=\(%start as lua),\(%stop as lua),\(%step as lua) do\
            ..\%loop_var_shim
            \(%body as lua statements)\
            ..\%continue_labels
            end --numeric for-loop

        for subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") == "stop %") and
                    ((2nd in (%'s "value"))'s "src") == (%var's "src")
        ..in %body
            %code <- ".."
                do -- scope for stopping for-loop
                    \%code
                    ::stop_\(nomsu "var_to_lua_identifier" [%var])::;
                end -- end of scope for stopping for-loop
            stop

        return %code

immediately
    parse [for %var from %start to %stop %body] as: for %var from %start to %stop via 1 %body
    parse [..]
        for all %start to %stop by %step %body
        for all %start to %stop via %step %body
    ..as: for % from %start to %stop via %step %body
    parse [for all %start to %stop %body] as: for all %start to %stop via 1 %body

# For-each loop
immediately
    compile [for %var in %iterable %body] to code
        %continue_labels <- ""
        for subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") == "do next %") and
                    ((3rd in (%'s "value"))'s "src") == (%var's "src")
        ..in %body
            %continue_labels <- "\n::continue_\(nomsu "var_to_lua_identifier" [%var])::;"
            stop
        if: (%var's "type") is "Var"
            %loop_var <- (%var as lua)
            %loop_var_shim <- ""
        ..else
            %loop_var <- "value"
            %loop_var_shim <- "\n\(%var as lua) = value;"
        # This trashes the loop variables, just like in Python.
        %code <- ".."
            for i,\%loop_var in ipairs(\(%iterable as lua)) do\
            ..\%loop_var_shim
            \(%body as lua statements)\
            ..\%continue_labels
            end --foreach-loop
        for subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") == "stop %") and
                    ((2nd in (%'s "value"))'s "src") == (%var's "src")
        ..in %body
            %code <- ".."
                do -- scope for stopping for-loop
                    \%code
                    ::stop_\(nomsu "var_to_lua_identifier" [%var])::;
                end -- end of scope for stopping for-loop
            stop
        return %code
    parse [for all %iterable %body] as: for % in %iterable %body

# Dict iteration (lua's "pairs()")
immediately
    compile [for %key = %value in %iterable %body] to code
        %continue_labels <- ""
        for subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") == "do next %") and
                    ((3rd in (%'s "value"))'s "src") == (%key's "src")
        ..in %body
            <- %continue_labels + "\n::continue_\(nomsu "var_to_lua_identifier" [%key])::;"
            stop

        for subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") == "do next %") and
                    ((3rd in (%'s "value"))'s "src") == (%value's "src")
        ..in %body
            <- %continue_labels + "\n::continue_\(nomsu "var_to_lua_identifier" [%value])::;"
            stop

        if: (%key's "type") is "Var"
            %key_loop_var <- (%key as lua)
            %key_loop_var_shim <- ""
        ..else
            %key_loop_var <- "key"
            %key_loop_var_shim <- "\n\(%key as lua) = key;"
        if: (%value's "type") is "Var"
            %value_loop_var <- (%value as lua)
            %value_loop_var_shim <- ""
        ..else
            %value_loop_var <- "value"
            %value_loop_var_shim <- "\n\(%value as lua) = value;"
        # This trashes the loop variables, just like in Python.
        %code <- ".."
            for \%key_loop_var,\%value_loop_var in pairs(\(%iterable as lua)) do\
            ..\%key_loop_var_shim\%value_loop_var_shim
            \(%body as lua statements)\
            ..\%continue_labels
            end --foreach-loop
        %stop_labels <- ""
        for subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") == "stop %") and
                    ((2nd in (%'s "value"))'s "src") == (%key's "src")
        ..in %body
            <- %stop_labels + "\n::stop_\(nomsu "var_to_lua_identifier" [%key])::;"
            stop

        for subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") == "stop %") and
                    ((2nd in (%'s "value"))'s "src") == (%value's "src")
        ..in %body
            <- %stop_labels + "\n::stop_\(nomsu "var_to_lua_identifier" [%value])::;"
            stop

        if: %stop_labels is not ""
            %code <- ".."
                do -- scope for stopping for % = % loop
                    \%code\%stop_labels
                end
        return %code

# Switch statement/multi-branch if
immediately
    compile [when %body] to code
        %result <- ""
        %fallthroughs <- []
        %is_first <- (yes)
        for %func_call in (%body's "value")
            assume ((%func_call's "type") is "FunctionCall") or barf ".."
                Invalid format for 'when' statement. Only '*' blocks are allowed.
            with [..]
                %tokens <- (%func_call's "value")
                %star <- (1st in %tokens)
                %condition <- (2nd in %tokens)
                %action <- (3rd in %tokens)
            ..
                assume (=lua "\%star and \%star.type == 'Word' and \%star.value == '*'") or barf ".."
                    Invalid format for 'when' statement. Lines must begin with '*'
                assume %condition or barf ".."
                    Invalid format for 'when' statement. Lines must begin with '*' and have a condition or the word "else"
                if: %action is (nil)
                    lua do> "table.insert(\%fallthroughs, \%condition)"
                    do next %func_call

                if: =lua "\%condition.type == 'Word' and \%condition.value == 'else'"

                    <- %result + ".."
                        
                        else
                        \(%action as lua statements)
                    stop
                ..else
                    %condition <- (%condition as lua)
                    for all %fallthroughs
                        <- %condition + " or \(% as lua)"
                    <- %result + ".."
                        
                        \("if" if %is_first else "elseif") \%condition then
                        \(%action as lua statements)

                %fallthroughs <- []
                %is_first <- (no)

        if: %result is not ""
            <- %result + "\nend"
        return %result

# Switch statement
immediately
    compile [when %branch_value = ? %body, when %branch_value is ? %body] to code
        %result <- ""
        %fallthroughs <- []
        %is_first <- (yes)
        %seen_else <- (no)
        for %func_call in (%body's "value")
            assume ((%func_call's "type") is "FunctionCall") or barf ".."
                Invalid format for 'when' statement. Only '*' blocks are allowed.
            %tokens <- (%func_call's "value")
            with [%star<-(1st in %tokens), %condition<-(2nd in %tokens), %action<-(3rd in %tokens)]
                assume (=lua "\%star and \%star.type == 'Word' and \%star.value == '*'") or barf ".."
                    Invalid format for 'when' statement. Lines must begin with '*'
                assume %condition or barf ".."
                    Invalid format for 'when' statement. Lines must begin with '*' and have a condition or the word "else"
                if: %action is (nil)
                    lua> "table.insert(\%fallthroughs, \%condition)"
                    do next %func_call

                if: =lua "\%condition.type == 'Word' and \%condition.value == 'else'"
                    assume (not %is_first) or barf "'else' clause cannot be first in 'when % = ?' block"
                    <- %result + ".."

                        else
                            \(%action as lua statements)
                        end
                    %seen_else <- (yes)
                ..else
                    assume (not %seen_else) or barf "'else' clause needs to be last in 'when % = ?' block"
                    %clause <- ""
                    if: ((%condition's "type") is "Text") or ((%condition's "type") is "Number")
                        %clause <- "branch_value == (\(%condition as lua))"
                    ..else 
                        %clause <- "utils.equivalent(branch_value, \(%condition as lua))"
                    for all %fallthroughs
                        if: ((%'s "type") is "Text") or ((%'s "type") is "Number")
                            <- %clause + " or branch_value == (\(%condition as lua))"
                        ..else 
                            <- %clause + " or utils.equivalent(branch_value, \(%condition as lua))"
                    <- %result + ".."
                        
                        \("if" if %is_first else "elseif") \%clause then
                        \(%action as lua statements)

            %fallthroughs <- []
            %is_first <- (no)
        
        assume (%result is not "") or barf "No body for 'when % = ?' block!"
        unless %seen_else
            <- %result + "\nend"
        %result <- ".."
            do --when % = ?
            local branch_value = \(%branch_value as lua);\
            ..\%result
            end --when % = ?
        return %result

# Try/except
immediately
    compile [..]
        try %action and if it succeeds %success or if it fails %fallback
        try %action and if it fails %fallback or if it succeeds %success
    ..to code ".."
        do
            local fell_through = false;
            local ok, ret = pcall(function()
                \(%action as lua statements)
                fell_through = true;
            end);
            if ok then
                \(%success as lua statements)
            end
            if not ok then
                \(%fallback as lua statements)
            elseif not fell_through then
                return ret;
            end
        end
    parse [try %action] as
        try %action and if it succeeds: do nothing
        ..or if it fails: do nothing
    parse [try %action and if it fails %fallback] as
        try %action and if it succeeds: do nothing
        ..or if it fails %fallback
    parse [try %action and if it succeeds %success] as
        try %action and if it succeeds %success or if it fails: do nothing

# Do/finally:
immediately
    compile [do %action] to code
        "do\n    \(%action as lua statements)\nend" if ((%action's "type") is "Block")
        ..else "(\(%action as lua))(nomsu);"

    compile [do %action then always %final_action] to code ".."
        do
            local fell_through = false;
            local ok, ret1 = pcall(function()
                \(%action as lua statements)
                fell_through = true;
            end);
            local ok2, ret2 = pcall(function()
                \(%final_action as lua statements)
            end);
            if not ok then error(ret1); end
            if not ok2 then error(ret2); end
            if not fell_through then
                return ret1;
            end
        end