#!/usr/bin/env nomsu -V2.3.4.3
#
    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 "if \(%condition as lua expr) then\n    \(%if_body as lua statements)\nend"

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 ".."
        if \(%condition as lua expr) then
            \(%if_body as lua statements)
        else
            \(%else_body as lua statements)
        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"
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
compile [=== %label ===, --- %label ---, *** %label ***] to (..)
    Lua "::label_\(%label as lua identifier)::"

compile [go to %label] to (Lua "goto label_\(%label as lua identifier)")

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

# Helper function
#TODO: do "using % compile %" instead so it's actually a helper function
compile [%tree has subtree %subtree where %condition] to (..)
    Lua value ".."
        (function()
            for \(%subtree as lua expr) in coroutine.wrap(function() \(%tree as lua expr):map(coroutine.yield) end) do
                if \(%condition as lua expr) then
                    return true
                end
            end
            return false
        end)()

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

# While loops
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\n    \(%body as lua statements)"
    
    if (..)
        %body has subtree % where ((%.type == "Action") and (%.stub is "do next repeat"))
    ..:
        to %lua write "\n    ::continue_repeat::"
    
    to %lua write "\nend --while-loop"
    if (..)
        %body has subtree % where ((%.type == "Action") and (%.stub is "stop repeating"))
    ..:
        %lua = (..)
            Lua ".."
                do -- scope of "stop repeating" label
                    \%lua
                    ::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)
compile [repeat %n times %body] to:
    %lua = (..)
        Lua "for i=1,\(%n as lua expr) do\n    \(%body as lua statements)"
    
    if (..)
        %body has subtree % where ((%.type == "Action") and (%.stub is "do next repeat"))
    ..:
        to %lua write "\n    ::continue_repeat::"
    
    to %lua write "\nend --numeric for-loop"
    if (..)
        %body has subtree % where ((%.type == "Action") and (%.stub is "stop repeating"))
    ..:
        %lua = (..)
            Lua ".."
                do -- scope of "stop repeating" label
                    \%lua
                    ::stop_repeat::
                end -- end of "stop repeating" label scope
    
    return %lua


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

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

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

# 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.source "Loop expected variable, not: %s"
    %lua = (..)
        Lua ".."
            for \(%var as lua expr)=\(%start as lua expr),\(%stop as lua expr),\(%step as lua expr) do
                \(%body as lua statements)
    
    if (..)
        %body has subtree % where (..)
            (%.type == "Action") and ((%.stub is "do next %") and (%.(3).1 == %var.1))
    ..:
        to %lua write "\n    \(compile as (===next %var ===))"
    
    to %lua write "\nend --numeric for-loop"
    if (..)
        %body has subtree % where (..)
            (%.type == "Action") and ((%.stub is "stop %") and (%.(2).1 == %var.1))
    ..:
        %lua = (..)
            Lua ".."
                do -- scope for stopping for-loop
                    \%lua
                    \(compile as (===stop %var ===))
                end -- 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

# For-each loop (lua's "ipairs()")
compile [for %var in %iterable %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.source "Loop expected variable, not: %s"
    %lua = (..)
        Lua ".."
            for i,\(%var as lua identifier) in ipairs(\(%iterable as lua expr)) do
                \(%body as lua statements)
    
    if (..)
        %body has subtree % where (..)
            (%.type == "Action") and ((%.stub is "do next %") and (%.(3).1 == %var.1))
    ..:
        to %lua write (..)
            Lua "\n\(compile as (===next %var ===))"
    
    to %lua write "\nend --foreach-loop"
    if (..)
        %body has subtree % where (..)
            (%.type == "Action") and ((%.stub is "stop %") and (%.(2).1 == %var.1))
    ..:
        %lua = (..)
            Lua ".."
                do -- scope for stopping for-loop
                    \%lua
                    \(compile as (===stop %var ===))
                end -- end of scope for stopping for-loop
    
    return %lua


# 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.source "Loop expected variable, not: %s"
    unless (%value.type is "Var"): compile error at %value.source "Loop expected variable, not: %s"
    %lua = (..)
        Lua ".."
            for \(%key as lua identifier),\(%value as lua identifier) in pairs(\(%iterable as lua expr)) do
                \(%body as lua statements)
    
    if (..)
        %body has subtree % where (..)
            (%.type == "Action") and ((%.stub is "do next %") and (%.(3).1 == %key.1))
    ..:
        to %lua write (..)
            Lua "\n\(compile as (===next %key ===))"
    
    if (..)
        %body has subtree % where (..)
            (%.type == "Action") and ((%.stub is "do next %") and (%.(3).1 == %value.1))
    ..:
        to %lua write (..)
            Lua "\n\(compile as (===next %value ===))"
    
    to %lua write "\nend --foreach-loop"
    %stop_labels = (Lua "")
    if (..)
        %body has subtree % where (..)
            (%.type == "Action") and ((%.stub is "stop %") and (%.(2).1 == %key.1))
    ..:
        to %stop_labels write "\n\(compile as (===stop %key ===))"
    
    if (..)
        %body has subtree % where (..)
            (%.type == "Action") and ((%.stub is "stop %") and (%.(2).1 == %value.1))
    ..:
        to %stop_labels write "\n\(compile as (===stop %value ===))"
    
    if ((length of "\%stop_labels") > 0):
        %lua = (..)
            Lua "do -- scope for stopping for % = % loop\n    \%lua\%stop_labels\nend"
    
    return %lua

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

# 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.source "'if' expected a Block, but got: %s"
    for %line in %body:
        unless (..)
            ((%line.type is "Action") and ((length of %line) >= 2))
            ..and (%line.(length of %line) is "Block" syntax tree)
        ..:
            compile error at %line.source ".."
                Invalid line for 'if', each line should contain conditional expressions \
                ..followed by a block, or "else" followed by a block:
                %s
        %action = %line.(length of %line)
        if ((%line.1 is "else") and ((length of %line) == 2)):
            unless %else_allowed: compile error at %line.source "Can't have two 'else' blocks"
            unless ((length of "\%code") > 0):
                compile error at %line.source "Can't have an 'else' block without a preceeding condition"
            to %code write "\nelse\n    \(%action as lua statements)"
            %else_allowed = (no)
        ..else:
            to %code write "\%clause "
            for %i in 1 to ((length of %line) - 1):
                unless (%line.%i is syntax tree):
                    compile error at %line.source ".."
                        Invalid condition for 'if' statement:
                        %s
                if (%i > 1): to %code write " or "
                to %code write (%line.%i as lua expr)
            to %code write " then\n    \(%action as lua statements)"
            %clause = "\nelseif"
    
    if ((length of "\%code") == 0):
        compile error at %body.source "'if' block has an empty body"
    to %code write "\nend --when"
    return %code

# Switch statement
compile [if %branch_value is %body, when %branch_value is %body] to:
    %code = (Lua "")
    %clause = "if"
    %else_allowed = (yes)
    unless (%body.type is "Block"): compile error at %body.source "'if' expected a Block, but got: %s"
    for %line in %body:
        unless (..)
            ((%line.type is "Action") and ((length of %line) >= 2))
            ..and (%line.(length of %line) is "Block" syntax tree)
        ..:
            compile error at %line.source ".."
                Invalid line for 'if % is % %', each line should contain expressions \
                ..followed by a block, or "else" followed by a block:
                %s
        %action = %line.(length of %line)
        if ((%line.1 is "else") and ((length of %line) == 2)):
            unless %else_allowed: compile error at %line.source "Can't have two 'else' blocks"
            unless ((length of "\%code") > 0):
                compile error at %line.source "Can't have an 'else' block without a preceeding condition"
            to %code write "\nelse\n    \(%action as lua statements)"
            %else_allowed = (no)
        ..else:
            to %code write "\%clause "
            for %i in 1 to ((length of %line) - 1):
                unless (%line.%i is syntax tree):
                    compile error at %line.source ".."
                        Invalid condition for 'if' statement:
                        %s
                if (%i > 1): to %code write " or "
                to %code write "branch_value == \(%line.%i as lua expr)"
            to %code write " then\n    \(%action as lua statements)"
            %clause = "\nelseif"
    
    if ((length of "\%code") == 0):
        compile error at %body.source "'if % is % %' block has an empty body"
    to %code write "\nend --when"
    return (..)
        Lua ".."
            do --if % is
                local branch_value = \(%branch_value as lua expr)
                \%code
            end --if % is

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

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


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

# Recurion control flow
compile [for %var in recursive %structure %body] to (..)
    with local compile actions:
        compile [recurse %v on %x] to (..)
            Lua "table.insert(stack\(%v as lua id), \(%x as lua expr))"
        
        return (..)
            Lua ".."
                do
                    local stack\(%var as lua id) = list{\(%structure as lua expr)}
                    while #stack\(%var as lua id) > 0 do
                        \(%var as lua expr) = table.remove(stack\(%var as lua id), 1)
                        \(%body as lua statements)
                        \(compile as (===next %var ===))
                    end
                    \(compile as (===stop %var ===))
                end