path: root/lib/control_flow.nom

require "lib/metaprogramming.nom"
require "lib/operators.nom"
require "lib/utils.nom"

# 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
    parse [if %x == %y %if_body] as: if (%x == %y) %if_body
    parse [if %x != %y %if_body] as: if (%x != %y) %if_body
    parse [unless %x == %y %if_body] as: if (%x != %y) %if_body
    parse [unless %x != %y %if_body] as: if (%x == %y) %if_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
    parse [if %x == %y %if_body else %else_body] as: if (%x == %y) %if_body else %else_body
    parse [if %x != %y %if_body else %else_body] as: if (%x != %y) %if_body else %else_body
    parse [unless %x == %y %if_body else %else_body] as: if (%x != %y) %if_body else %else_body
    parse [unless %x != %y %if_body else %else_body] as: if (%x == %y) %if_body else %else_body

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

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

# Helper function
immediately:
    rule [tree %tree has function call %call] =:
        lua> ".."
            local target = (\%call).stub;
            for subtree,_ in coroutine.wrap(function() nomsu:walk_tree(\%tree); end) do
                if type(subtree) == 'table' and subtree.type == "FunctionCall"
                and subtree.stub == target then
                    return true;
                end
            end
            return false;

# While loops
immediately:
    compile [do next repeat-loop] to code: "goto continue_repeat;"
    compile [stop repeat-loop] to code: "goto stop_repeat;"
    compile [repeat while %condition %body] to code:
        %continue_labels = (..)
            "\n::continue_repeat::;" if (tree %body has function call \(do next repeat-loop)) else ""
        %code = ".."
            while \(%condition as lua) do
            \(%body as lua statements)\
            ..\%continue_labels
            end --while-loop
        if (tree %body has function call \(stop repeat-loop)):
            return ".."
                do --while-loop label scope
                \%code
                ::stop_repeat::;
                end --while-loop label scope
        return %code
    parse [repeat %body] as: repeat while (true) %body
    parse [repeat until %condition %body] as: repeat while (not %condition) %body
    parse [repeat while %x == %y %body] as: repeat while (%x == %y) %body
    parse [repeat while %x != %y %body] as: repeat while (%x != %y) %body
    parse [repeat until %x == %y %body] as: repeat while (%x != %y) %body
    parse [repeat until %x != %y %body] as: repeat while (%x == %y) %body

# For loop control flow:
immediately:
    compile [stop for-loop] to code: "goto stop_for;"
    compile [stop %var] to code: ".."
        goto stop_\(nomsu "var_to_lua_identifier" [%var]);
    compile [do next for-loop] to code: "goto continue_for;"
    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:
        lua> "local \%continue_labels, \%code, \%stop_labels"
        %continue_labels = ""
        if (tree %body has function call \(do next for-loop)):
            %continue_labels join= "\n::continue_for::;"
        if (tree %body has function call (tree \(do next %) with {""=%var})):
            %continue_labels join= "\n::continue_\(nomsu "var_to_lua_identifier" [%var])::;"
        # This trashes the loop variables, just like in Python.
        %code = ".."
            for i=\(%start as lua),\(%stop as lua),\(%step as lua) do
            \(%var as lua) = i;
            \(%body as lua statements)\
            ..\%continue_labels
            end --numeric for-loop
        %stop_labels = ""
        if (tree %body has function call \(stop for-loop)):
            %stop_labels join= "\n::stop_for::;"
        if (tree %body has function call (tree \(stop %) with {""=%var})):
            %stop_labels join= "\n::stop_\(nomsu "var_to_lua_identifier" [%var])::;"
        return (..)
            ".."
                do --for-loop label scope
                \%code\
                ..\%stop_labels
                end --for-loop label scope
            ..if %stop_labels != "" else %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

immediately:
    compile [for %var in %iterable %body] to code:
        lua> "local \%continue_labels, \%stop_labels, \%code, \%stop_labels"
        %continue_labels = ""
        if (tree %body has function call \(do next for-loop)):
            %continue_labels join= "\n::continue_for::;"
        if (tree %body has function call (tree \(do next %) with {""=%var})):
            %continue_labels join= "\n::continue_\(nomsu "var_to_lua_identifier" [%var])::;"
        # This trashes the loop variables, just like in Python.
        %code = ".."
            for i,value in ipairs(\(%iterable as lua)) do
            \(%var as lua) = value;
            \(%body as lua statements)\
            ..\%continue_labels
            end --foreach-loop
        %stop_labels = ""
        if (tree %body has function call \(stop for-loop)):
            %stop_labels join= "\n::stop_for::;"
        if (tree %body has function call (tree \(stop %) with {""=%var})):
            %stop_labels join= "\n::stop_\(nomsu "var_to_lua_identifier" [%var])::;"
        if %stop_labels != "":
            %code = ".."
                do --for-loop label scope
                \%code\%stop_labels
                end --for-loop label scope
        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 = ""
        if (tree %body has function call \(do next for-loop)):
            %continue_labels join= "\n::continue_for::;"
        if (tree %body has function call (tree \(do next %x) with {x=%key})):
            %continue_labels join= "\n::continue_\(nomsu "var_to_lua_identifier" [%key])::;"
        if (tree %body has function call (tree \(do next %) with {""=%value})):
            %continue_labels join= "\n::continue_\(nomsu "var_to_lua_identifier" [%value])::;"
        # This trashes the loop variables, just like in Python.
        %code = ".."
            for key,value in ipairs(\(%iterable as lua)) do
            \(%key as lua), \(%value as lua) = key, value;
            \(%body as lua statements)\
            ..\%continue_labels
            end --foreach-loop
        %stop_labels = ""
        if (tree %body has function call \(stop for-loop)):
            %stop_labels join= "\n::stop_for::;"
        if (tree %body has function call (tree \(stop %) with {""=%key})):
            %stop_labels join= "\n::stop_\(nomsu "var_to_lua_identifier" [%key])::;"
        if (tree %body has function call (tree \(stop %) with {""=%value})):
            %stop_labels join= "\n::stop_\(nomsu "var_to_lua_identifier" [%value])::;"
        return (..)
            ".."
                do --for-loop label scope
                \%code\
                ..\%stop_labels
                end --for-loop label scope
            ..if %stop_labels != "" else %code

# Switch statement/multi-branch if
immediately:
    compile [when %body] to code:
        lua> "local \%result, \%fallthroughs, \%first"
        %result = ""
        %fallthroughs = []
        %first = (yes)
        for %func_call in (%body's "value"):
            lua> "local \%tokens, \%star, \%condition, \%action"
            assert ((%func_call's "type") == "FunctionCall") ".."
                Invalid format for 'when' statement. Only '*' blocks are allowed.
            %tokens = (%func_call's "value")
            %star = (%tokens -> 1)
            assert (=lua "\%star and \%star.type == 'Word' and \%star.value == '*'") ".."
                Invalid format for 'when' statement. Lines must begin with '*'

            %condition = (%tokens -> 2)
            assert %condition ".."
                Invalid format for 'when' statement. Lines must begin with '*' and have a condition or the word "else"

            %action = (%tokens -> 3)
            if (%action == (nil)):
                lua do> "table.insert(\%fallthroughs, \%condition)"
                do next %func_call

            if (=lua "\%condition.type == 'Word' and \%condition.value == 'else'"):
                %result join= ".."
                    
                    else
                    \(%action as lua statements)
                stop for-loop
            ..else:
                %condition = (%condition as lua)
                for all %fallthroughs:
                    %condition join= " or \(% as lua)"
                %result join= ".."
                    
                    \("if" if %first else "elseif") \%condition then
                    \(%action as lua statements)

            %fallthroughs = []
            %first = (no)

        if (%result != ""):
            %result join= "\nend"
        return %result

# Switch statement
immediately:
    compile [when %branch_value == ? %body] to code:
        %result = ""
        %fallthroughs = []
        %first = (yes)
        for %func_call in (%body's "value"):
            assert ((%func_call's "type") == "FunctionCall") ".."
                Invalid format for 'when' statement. Only '*' blocks are allowed.
            %tokens = (%func_call's "value")
            %star = (%tokens -> 1)
            assert (=lua "\%star and \%star.type == 'Word' and \%star.value == '*'") ".."
                Invalid format for 'when' statement. Lines must begin with '*'

            %condition = (%tokens -> 2)
            assert %condition ".."
                Invalid format for 'when' statement. Lines must begin with '*' and have a condition or the word "else"

            %action = (%tokens -> 3)
            if (%action == (nil)):
                lua> "table.insert(\%fallthroughs, \%condition)"
                do next %func_call

            if (=lua "\%condition.type == 'Word' and \%condition.value == 'else'"):
                %result join= ".."
                    
                    else
                    \(%action as lua statements)
                stop for-loop
            ..else:
                %condition = "branch_value == (\(%condition as lua))"
                for all %fallthroughs:
                    %condition join= " or (branch_value == \(% as lua))"
                %result join= ".."
                    
                    \("if" if %first else "elseif") \%condition then
                    \(%action as lua statements)

            %fallthroughs = []
            %first = (no)

        if (%result != ""):
            %result = ".."
                do --when == ?
                local branch_value = \(%branch_value as lua);\
                ..\%result
                end
                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, ret1, ret2 = pcall(function(nomsu)
                \(%action as lua statements)
                fell_through = true;
            end, nomsu);
            if ok then
                \(%success as lua statements)
            end
            if not ok then
                \(%fallback as lua statements)
            elseif not fell_through then
                return ret1, ret2;
            end
        end
    parse [try %action] as:
        try %action and if it succeeds: pass
        ..or if it fails: pass
    parse [try %action and if it fails %fallback] as:
        try %action and if it succeeds: pass
        ..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: pass

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