#..
    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 {statements:""}

# Conditionals
immediately
    compile [if %condition %if_body] to
        %if_body <- (%if_body as lua)
        return {..}
            locals: %if_body's "locals"
            statements:".."
                if \(%condition as lua expr) then
                \((%if_body's "statements") or "\(%if_body's "expr");")
                end
    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
        %if_body <- (%if_body as lua)
        %else_body <- (%else_body as lua)
        lua> ".."
            local \%locals = {unpack(\%if_body.locals or {})};
            for i,loc in ipairs(\%else_body.locals or {}) do table.insert(\%locals, loc); end
            utils.deduplicate(\%locals);
        return {..}
            locals:%locals
            statements:".."
                if \(%condition as lua expr) then
                \((%if_body's "statements") or "\(%if_body's "expr");")
                else
                \((%else_body's "statements") or "\(%else_body's "expr");")
                end

# 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 {..}
                expr:"(\(%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 {..}
                expr: ".."
                    (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
immediately
    compile [=== %label ===, --- %label ---, *** %label ***] to {..}
        statements:"::label_\(nomsu "var_to_lua_identifier" [%label])::;"
    compile [go to %label] to {..}
        statements:"goto label_\(nomsu "var_to_lua_identifier" [%label]);"

# Basic loop control
immediately
    compile [do next] to {statements:"continue;"}
    compile [stop] to {statements:"break;"}

# Helper function
immediately
    compile [if %tree has subtree %subtree where %condition %body] to
        %body_lua <- (%body as lua)
        %body_statements <- ((%body_lua's "statements") or "\(%body_lua's "expr");")
        return {..}
            locals: %body_lua's "locals"
            statements:".."
                for \(%subtree as lua expr) in coroutine.wrap(function() nomsu:walk_tree(\(%tree as lua expr)) end) do
                    if type(\(%subtree as lua expr)) == 'table' and \(%subtree as lua expr).type then
                        if \(%condition as lua expr) then
                            \%body_statements
                            break;
                        end
                    end
                end

# While loops
immediately
    compile [do next repetition] to {statements:"goto continue_repeat;"}
    compile [stop repeating] to {statements:"goto stop_repeat;"}
    compile [repeat while %condition %body] to
        %body_lua <- (%body as lua)
        %body_statements <- ((%body_lua's "statements") or "\(%body_lua's "expr");")
        if %body has subtree % where 
            ((%'s "type") = "FunctionCall") and ((%'s "stub") is "do next repetition")
        ..: <- %body_statments + "\n::continue_repeat::;"
        %code <- ".."
            while \(%condition as lua expr) do
            \%body_statements
            end --while-loop
        if %body has subtree % where
            ((%'s "type") = "FunctionCall") and ((%'s "stub") is "stop repeating")
        ..
            %code <- ".."
                do -- scope of "stop repeating" label
                \%code
                ::stop_repeat::;
                end -- end of "stop repeating" label scope
        return {statements:%code, locals:%body_lua's "locals"}
    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
        %body_lua <- (%body as lua)
        %body_statements <- ((%body_lua's "statements") or "\(%body_lua's "expr");")
        if %body has subtree % where
            ((%'s "type") = "FunctionCall") and ((%'s "stub") is "do next repetition")
        ..: <- %body_statements + "\n::continue_repeat::;"
        %code <- ".."
            for i=1,\(%n as lua expr) do
            \%body_statements
            end --numeric for-loop
        if %body has subtree % where
            ((%'s "type") = "FunctionCall") and ((%'s "stub") is "stop repeating")
        ..
            %code <- ".."
                do -- scope of "stop repeating" label
                \%code
                ::stop_repeat::;
                end -- end of "stop repeating" label scope
        return {statements:%code, locals:%body_lua's "locals"}

# For loop control flow:
immediately
    compile [stop %var] to {..}
        statements:"goto stop_\(nomsu "var_to_lua_identifier" [%var]);"
    compile [do next %var] to {..}
        statements:"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
        %body_lua <- (%body as lua)
        %body_statements <- ((%body_lua's "statements") or "\(%body_lua's "expr");")
        if %body has subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") is "do next %") and 
                    ((3rd in (%'s "value"))'s "value") is (%var's "value")
        ..: <- %body_statements + "\n::continue_\(nomsu "var_to_lua_identifier" [%var])::;"

        # This uses Lua's approach of only allowing loop-scoped variables in a loop
        assume ((%var's "type") is "Var") or barf "Loop expected variable, not: \(%var's source code)"
        %code <- ".."
            for \(%var as lua expr)=\(%start as lua expr),\(%stop as lua expr),\(%step as lua expr) do
            \%body_statements
            end --numeric for-loop

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

        return {statements:%code, locals:%body_lua's "locals"}

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 (lua's "ipairs()")
immediately
    compile [for %var in %iterable %body] to
        %body_lua <- (%body as lua)
        %body_statements <- ((%body_lua's "statements") or "\(%body_lua's "expr");")
        if %body has subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") is "do next %") and
                    ((3rd in (%'s "value"))'s "value") is (%var's "value")
        ..: <- %body_statements + "\n::continue_\(nomsu "var_to_lua_identifier" [%var])::;"
        # This uses Lua's approach of only allowing loop-scoped variables in a loop
        assume ((%var's "type") is "Var") or barf "Loop expected variable, not: \(%var's source code)"
        %code <- ".."
            for i,\(%var as lua expr) in ipairs(\(%iterable as lua expr)) do
            \%body_statements
            end --foreach-loop
        if %body has subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") is "stop %") and
                    ((2nd in (%'s "value"))'s "value") is (%var's "value")
        ..
            %code <- ".."
                do -- scope for stopping for-loop
                    \%code
                    ::stop_\(nomsu "var_to_lua_identifier" [%var])::;
                end -- end of scope for stopping for-loop
        return {statements:%code, locals:%body_lua's "locals"}

    parse [for all %iterable %body] as: for % in %iterable %body

# Dict iteration (lua's "pairs()")
immediately
    compile [for %key = %value in %iterable %body] to
        %body_lua <- (%body as lua)
        %body_statements <- ((%body_lua's "statements") or "\(%body_lua's "expr");")
        if %body has subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") is "do next %") and
                    ((3rd in (%'s "value"))'s "value") is (%key's "value")
        ..: <- %body_statements + "\n::continue_\(nomsu "var_to_lua_identifier" [%key])::;"

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

        # This uses Lua's approach of only allowing loop-scoped variables in a loop
        assume ((%key's "type") is "Var") or barf "Loop expected variable, not: \(%key's source code)"
        assume ((%value's "type") is "Var") or barf "Loop expected variable, not: \(%value's source code)"
        %code <- ".."
            for \(%key as lua expr),\(%value as lua expr) in pairs(\(%iterable as lua expr)) do
            \%body_statements
            end --foreach-loop
        
        %stop_labels <- ""
        if %body has subtree % where
            ((%'s "type") = "FunctionCall") and
                ((%'s "stub") is "stop %") and
                    ((2nd in (%'s "value"))'s "value") is (%key's "value")
        ..: <- %stop_labels + "\n::stop_\(nomsu "var_to_lua_identifier" [%key])::;"

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

        if: %stop_labels is not ""
            %code <- ".."
                do -- scope for stopping for % = % loop
                    \%code\%stop_labels
                end
        return {statements:%code, locals:%body_lua's "locals"}

# Switch statement/multi-branch if
immediately
    compile [when %body] to
        %code <- ""
        %fallthroughs <- []
        %locals <- []
        %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.
            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> "table.insert(\%fallthroughs, \%condition);"
                    do next %func_call
                %action <- (%action as lua)
                %action_statements <- ((%action's "statements") or "\(%action's "expr");")
                for %local in ((%action's "locals") or [])
                    lua> "table.insert(\%locals, \%local);"

                if: =lua "\%condition.type == 'Word' and \%condition.value == 'else'"
                    <- %code + ".."
                        
                        else
                        \%action_statements
                    %seen_else <- (yes)
                ..else
                    assume (not %seen_else) or barf "'else' clause needs to be last in 'when' block"
                    %condition <- (%condition as lua expr)
                    for all %fallthroughs
                        <- %condition + " or \(% as lua)"
                    <- %code + ".."
                        
                        \("if" if %is_first else "elseif") \%condition then
                        \%action_statements

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

        assume (%fallthroughs = []) or barf "Unfinished fallthrough conditions in 'when' block"
        if: %code is not ""
            <- %code + "\nend"
        lua> "utils.deduplicate(\%locals);"
        return {statements:%code, locals:%locals}

# Switch statement
immediately
    compile [when %branch_value = ? %body, when %branch_value is ? %body] to
        %code <- ""
        %fallthroughs <- []
        %locals <- []
        %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

                %action <- (%action as lua)
                %action_statements <- ((%action's "statements") or "\(%action's "expr");")
                for %local in ((%action's "locals") or [])
                    lua> "table.insert(\%locals, \%local);"

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

                        else
                            \%action_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 expr))"
                    ..else 
                        %clause <- "utils.equivalent(branch_value, \(%condition as lua expr))"
                    for all %fallthroughs
                        if: ((%'s "type") is "Text") or ((%'s "type") is "Number")
                            <- %clause + " or branch_value == (\(%condition as lua expr))"
                        ..else 
                            <- %clause + " or utils.equivalent(branch_value, \(%condition as lua expr))"
                    <- %code + ".."
                        
                        \("if" if %is_first else "elseif") \%clause then
                        \%action_statements

            %fallthroughs <- []
            %is_first <- (no)
        
        assume (%fallthroughs = []) or barf "Unfinished fallthrough conditions in 'when' block"
        assume (%code is not "") or barf "No body for 'when % = ?' block!"
        unless %seen_else
            <- %code + "\nend"
        %code <- ".."
            do --when % = ?
            local branch_value = \(%branch_value as lua expr);\
            ..\%code
            end --when % = ?
        lua> "utils.deduplicate(\%locals);"
        return {statements:%code, locals:%locals}

# 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
        %locals <- []
        %action <- (%action as lua)
        %fallback <- (%fallback as lua)
        for %sub_locals in [%action's "locals", %success's "locals", %fallback's "locals"]
            for %local in %sub_locals
                lua> "table.insert(\%locals, \%local);"
        lua> "utils.deduplicate(\%locals);"
        return {..}
            locals: %locals
            statements: ".."
                do
                    local fell_through = false;
                    local ok, ret = pcall(function()
                        \((%action's "statements") or "\(%action's "expr");")
                        fell_through = true;
                    end);
                    if ok then
                        \((%success's "statements") or "\(%success's "expr");")
                    end
                    if not ok then
                        \((%fallback's "statements") or "\(%fallback's "expr");")
                    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
        %action <- (%action as lua)
        return {..}
            locals: %action's "locals"
            statements: ".."
                do
                    \((%action's "statements") or "\(%action's "expr");")
                end

    compile [do %action then always %final_action] to
        %action <- (%action as lua)
        %final_action <- (%action as lua)
        %locals <- []
        for %sub_locals in [%action's "locals", %final_action's "locals"]
            for %local in %sub_locals
                lua> "table.insert(\%locals, \%local);"
        lua> "utils.deduplicate(\%locals);"
        return {..}
            locals: %locals
            statements: ".."
                do
                    local fell_through = false;
                    local ok, ret1 = pcall(function()
                        \((%action's "statements") or "\(%action's "expr");")
                        fell_through = true;
                    end);
                    local ok2, ret2 = pcall(function()
                        \((%final_action's "statements") or "\(%final_action's "expr");")
                    end);
                    if not ok then error(ret1); end
                    if not ok2 then error(ret2); end
                    if not fell_through then
                        return ret1;
                    end
                end