{-
-- EPITECH PROJECT, 2025
-- Glados
-- File description:
-- CodeGeneration
-}

{-|
Module : Compiler.ASM.Compiler
Description : High-level code generation (Flow control, Definitions, Calls).
Stability : experimental
-}
module Compiler.ASM.Compiler 
    ( compileAst
    , compileIf
    , compileFor
    , compileWhile
    , compileSetVar
    , compileSetStruct
    , compileDefineFun
    , compileDefineLambda
    , compileDefineStruct
    , compileTail
    , compileLoop
    , compileAccessStruct
    , getLambdaFreeVariables
    , inferType
    ) where

import Data.Char (chr)
import Data.Text (Text, pack, unpack)
import Data.Set (Set)
import qualified Data.Set as Set
import qualified Data.Map.Strict as Map
import Control.Monad (zipWithM_, forM_)
import Control.Monad.State (lift, get, modify)

import Compiler.ASM.CompilerMonad
import Compiler.ASM.AstToAsm
    ( builtinMap
    , astSymbolToAsm
    , astIntToAsm
    , astBoolToAsm
    , astListToAsm
    , astCallToAsm
    )
import Compiler.Instruction (Instruction(..), Immediate(..))
import Compiler.PsInstruction (PsInstruction(..))
import Compiler.CompilerState (CompilerState(..), ScopeType(..))
import Compiler.ASM.Builtins (getBuiltinReturnType)
import Common.Type.Integer (IntValue(..))
import Common.Utils.List (zipWith3M_)
import AST.Ast (Ast(..))

-- | Analyzes an AST node to find free variables (variables used but not defined locally).
--
-- @args
--   - node: The AST node to analyze.
--
-- @details
--   Recursively traverses the AST to compute the set of free variables.
--   For a Lambda, free vars are (body_free_vars - parameters).
--   For a Define, the defined name is excluded from the body's free vars.
--   This function is critical for determining closure captures.
--   Built-in operators (like +, -, *) are excluded from free variables.
--
-- @return
--   A Set of Text representing the names of free variables.
--
getLambdaFreeVariables :: Ast -> Set Text
getLambdaFreeVariables :: Ast -> Set Text
getLambdaFreeVariables (ASymbol Text
s)
    | Text -> Map Text Instruction -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member Text
s Map Text Instruction
builtinMap = Set Text
forall a. Set a
Set.empty
    | Bool
otherwise = Text -> Set Text
forall a. a -> Set a
Set.singleton Text
s
getLambdaFreeVariables (AInteger IntValue
_) = Set Text
forall a. Set a
Set.empty
getLambdaFreeVariables (ABool Bool
_) = Set Text
forall a. Set a
Set.empty
getLambdaFreeVariables (ADefineLambda [Text]
params Ast
body) =
    Set Text -> Set Text -> Set Text
forall a. Ord a => Set a -> Set a -> Set a
Set.difference (Ast -> Set Text
getLambdaFreeVariables Ast
body) ([Text] -> Set Text
forall a. Ord a => [a] -> Set a
Set.fromList [Text]
params)
getLambdaFreeVariables (ASetVar Text
name Text
_ Ast
body) =
    Text -> Set Text -> Set Text
forall a. Ord a => a -> Set a -> Set a
Set.delete Text
name (Ast -> Set Text
getLambdaFreeVariables Ast
body)
getLambdaFreeVariables (ACall Ast
func [Ast]
args) =
    Set Text -> Set Text -> Set Text
forall a. Ord a => Set a -> Set a -> Set a
Set.union (Ast -> Set Text
getLambdaFreeVariables Ast
func) (
        [Set Text] -> Set Text
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ((Ast -> Set Text) -> [Ast] -> [Set Text]
forall a b. (a -> b) -> [a] -> [b]
map Ast -> Set Text
getLambdaFreeVariables [Ast]
args))
getLambdaFreeVariables (AIf Ast
cond Ast
t Ast
e) =
    [Set Text] -> Set Text
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions [Ast -> Set Text
getLambdaFreeVariables Ast
cond, Ast -> Set Text
getLambdaFreeVariables Ast
t,
        Ast -> Set Text
getLambdaFreeVariables Ast
e]
getLambdaFreeVariables (AList [Ast]
e) = [Set Text] -> Set Text
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ((Ast -> Set Text) -> [Ast] -> [Set Text]
forall a b. (a -> b) -> [a] -> [b]
map Ast -> Set Text
getLambdaFreeVariables [Ast]
e)
getLambdaFreeVariables (APos Int
_ Int
_ Ast
ast) = Ast -> Set Text
getLambdaFreeVariables Ast
ast
getLambdaFreeVariables Ast
_ = Set Text
forall a. Set a
Set.empty

-- | Helper function for Tail Call Optimization
--
-- @args
--   - compileFn: The standard compilation function for non-tail expressions.
--   - ast: The AST node to compile in tail position.
--
-- @details
--   Detects if the expression is a function call. If so, emits 'TailCallLabel'
--   (TCO). If it is a control structure (If, List), it propagates the tail
--   context recursively. Otherwise, compiles normally and emits 'Ret'.
--
compileTail :: (Ast -> CompilerMonad ()) -> Ast -> CompilerMonad ()
compileTail :: (Ast -> CompilerMonad ()) -> Ast -> CompilerMonad ()
compileTail Ast -> CompilerMonad ()
compileFn (ACall (ASymbol Text
name) [Ast]
args) =
    case Text -> Map Text Instruction -> Maybe Instruction
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
name Map Text Instruction
builtinMap of
        Just Instruction
instr -> do
            (Ast -> CompilerMonad ()) -> [Ast] -> CompilerMonad ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Ast -> CompilerMonad ()
compileFn [Ast]
args
            Instruction -> CompilerMonad ()
emitInstruction Instruction
instr
            s <- StateT CompilerState (Either Text) CompilerState
forall s (m :: * -> *). MonadState s m => m s
get
            emitInstruction (Ret (csCurrentArgCount s))
        Maybe Instruction
Nothing -> (Ast -> CompilerMonad ()) -> [Ast] -> CompilerMonad ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Ast -> CompilerMonad ()
compileFn [Ast]
args CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
            PsInstruction -> CompilerMonad ()
appendPseudoInstruction (Text -> PsInstruction
TailCallLabel Text
name)
compileTail Ast -> CompilerMonad ()
compileFn (ACall Ast
func [Ast]
args) = do
    Ast -> CompilerMonad ()
compileFn Ast
func
    (Ast -> CompilerMonad ()) -> [Ast] -> CompilerMonad ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Ast -> CompilerMonad ()
compileFn [Ast]
args
    Instruction -> CompilerMonad ()
emitInstruction Instruction
CallIndirect
    s <- StateT CompilerState (Either Text) CompilerState
forall s (m :: * -> *). MonadState s m => m s
get
    emitInstruction (Ret (csCurrentArgCount s))
compileTail Ast -> CompilerMonad ()
compileFn (AIf Ast
cond Ast
t Ast
e) = do
    lElse <- Text -> CompilerMonad Text
generateUniqueLabel (String -> Text
pack String
"else")
    lEnd  <- generateUniqueLabel (pack "endif")
    compileFn cond
    emitJumpIfFalseToLabel lElse
    compileTail compileFn t
    emitJumpToLabel lEnd
    emitLabelDefinition lElse
    compileTail compileFn e
    emitLabelDefinition lEnd
compileTail Ast -> CompilerMonad ()
compileFn (AList [Ast]
exprs)
    | [Ast] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Ast]
exprs = do
        s <- StateT CompilerState (Either Text) CompilerState
forall s (m :: * -> *). MonadState s m => m s
get
        emitInstruction (Ret (csCurrentArgCount s))
    | Bool
otherwise = (Ast -> CompilerMonad ()) -> [Ast] -> CompilerMonad ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Ast -> CompilerMonad ()
compileFn ([Ast] -> [Ast]
forall a. HasCallStack => [a] -> [a]
init [Ast]
exprs) CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
        (Ast -> CompilerMonad ()) -> Ast -> CompilerMonad ()
compileTail Ast -> CompilerMonad ()
compileFn ([Ast] -> Ast
forall a. HasCallStack => [a] -> a
last [Ast]
exprs)
compileTail Ast -> CompilerMonad ()
compileFn Ast
other = do
    Ast -> CompilerMonad ()
compileFn Ast
other
    s <- StateT CompilerState (Either Text) CompilerState
forall s (m :: * -> *). MonadState s m => m s
get
    emitInstruction (Ret (csCurrentArgCount s))

-- | Compiles a conditional expression (If-Then-Else).
--
-- @args
--   - compileFn: The recursive compilation function.
--   - cond: The condition AST.
--   - thenBranch: The AST for the True branch.
--   - elseBranch: The AST for the False branch.
--
-- @details
--   Generates labels for the Else block and the End. Emits JumpIfFalse to
--   skip the Then block if appropriate, and an unconditional Jump to skip
--   the Else block after executing Then.
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
compileIf :: (Ast -> CompilerMonad ()) -> Ast -> Ast -> Ast -> CompilerMonad ()
compileIf :: (Ast -> CompilerMonad ()) -> Ast -> Ast -> Ast -> CompilerMonad ()
compileIf Ast -> CompilerMonad ()
compileFn Ast
cond Ast
thenBranch Ast
elseBranch = do
    lElse <- Text -> CompilerMonad Text
generateUniqueLabel (String -> Text
pack String
"else")
    lEnd  <- generateUniqueLabel (pack "endif")
    compileFn cond
    emitJumpIfFalseToLabel lElse
    compileFn thenBranch
    emitJumpToLabel lEnd
    emitLabelDefinition lElse
    compileFn elseBranch
    emitLabelDefinition lEnd

-- | Helper function to compile the common logic of loops.
--
-- @args
--   - compileFn: The recursive compilation function.
--   - cond: The condition AST.
--   - body: The body AST.
--   - lEnd: The label to jump to if the condition is false.
--
-- @details
--   Compiles the condition, emits the jump-if-false check, and compiles the body.
--   This sequence is shared between While and For loops.
--
compileLoop :: (Ast -> CompilerMonad ()) -> Ast -> Ast -> Text -> CompilerMonad ()
compileLoop :: (Ast -> CompilerMonad ()) -> Ast -> Ast -> Text -> CompilerMonad ()
compileLoop Ast -> CompilerMonad ()
compileFn Ast
cond Ast
body Text
lEnd =
    Ast -> CompilerMonad ()
compileFn Ast
cond CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
    Text -> CompilerMonad ()
emitJumpIfFalseToLabel Text
lEnd CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
    Text -> CompilerMonad ()
pushLoopExit Text
lEnd CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
    Ast -> CompilerMonad ()
compileFn Ast
body CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
    CompilerMonad ()
popLoopExit

-- | Compiles a While loop.
--
-- @args
--   - compileFn: The recursive compilation function.
--   - cond: The loop condition AST.
--   - body: The loop body AST.
--
-- @details
--   Generates start and end labels. Uses 'compileLoop' for the core logic
--   and handles the looping jump back to start.
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
compileWhile :: (Ast -> CompilerMonad ()) -> Ast -> Ast -> CompilerMonad ()
compileWhile :: (Ast -> CompilerMonad ()) -> Ast -> Ast -> CompilerMonad ()
compileWhile Ast -> CompilerMonad ()
compileFn Ast
cond Ast
body = do
    lStart <- Text -> CompilerMonad Text
generateUniqueLabel (String -> Text
pack String
"while_start")
    lEnd   <- generateUniqueLabel (pack "while_end")
    emitLabelDefinition lStart
    compileLoop compileFn cond body lEnd
    emitJumpToLabel lStart
    emitLabelDefinition lEnd

-- | Compiles a For loop.
--
-- @args
--   - compileFn: The recursive compilation function.
--   - initAst: The initialization AST (executed once).
--   - cond: The loop condition AST.
--   - updateAst: The update/increment AST (executed after each iteration).
--   - body: The loop body AST.
--
-- @details
--   Compiles the initialization step first. Then uses 'compileLoop' for the
--   condition and body. Finally, compiles the update step before jumping back.
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
compileFor :: (Ast -> CompilerMonad ()) -> Ast -> Ast -> Ast -> Ast ->
    CompilerMonad ()
compileFor :: (Ast -> CompilerMonad ())
-> Ast -> Ast -> Ast -> Ast -> CompilerMonad ()
compileFor Ast -> CompilerMonad ()
compileFn Ast
initAst Ast
cond Ast
updateAst Ast
body = do
    lStart <- Text -> CompilerMonad Text
generateUniqueLabel (String -> Text
pack String
"for_start")
    lEnd   <- generateUniqueLabel (pack "for_end")
    compileFn initAst
    emitLabelDefinition lStart
    compileLoop compileFn cond body lEnd
    compileFn updateAst
    emitJumpToLabel lStart
    emitLabelDefinition lEnd

-- | Compiles a variable definition.
--
-- @args
--   - compileFn: The recursive compilation function.
--   - name: The name of the variable to define.
--   - body: The expression to evaluate.
--
-- @details
--   Compiles the body to push the value onto the stack, allocates a new
--   global symbol index, and emits a StoreGlobal instruction.
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
compileSetVar :: (Ast -> CompilerMonad ()) -> Text -> Text -> Ast ->
    Bool -> CompilerMonad ()
compileSetVar :: (Ast -> CompilerMonad ())
-> Text -> Text -> Ast -> Bool -> CompilerMonad ()
compileSetVar Ast -> CompilerMonad ()
compileFn Text
name Text
typeName Ast
body Bool
isStatement = do
    Ast -> CompilerMonad ()
compileFn Ast
body
    (scope, idx) <- Text -> Text -> CompilerMonad (ScopeType, Int)
defineSymbol Text
name Text
typeName
    case not isStatement of
        Bool
True -> Instruction -> CompilerMonad ()
emitInstruction Instruction
Dup
        Bool
False -> () -> CompilerMonad ()
forall a. a -> StateT CompilerState (Either Text) a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    case scope of
        ScopeType
ScopeGlobal  -> Instruction -> CompilerMonad ()
emitInstruction (Int -> Instruction
StoreGlobal Int
idx)
        ScopeType
ScopeLocal   -> Instruction -> CompilerMonad ()
emitInstruction (Int -> Instruction
StoreLocal Int
idx)
        ScopeType
ScopeCapture -> Instruction -> CompilerMonad ()
emitInstruction (Int -> Instruction
StoreCapture Int
idx)

-- | Compiles a structure instantiation.
--
-- @args
--   - compileFn: The recursive compilation function.
--   - name: The name of the structure to instantiate.
--   - assignedFields: The list of (Field Name, Value AST) pairs.
--
-- @details
--   Retrieves the structure definition from 'CompilerState' using 'getStructDefinition'.
--   Reorders the provided values to match the definition order, compiles them
--   (pushing to stack), and emits a 'BuildStruct' instruction.
--   Returns an error if the struct is undefined or fields mismatch.
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
compileSetStruct :: (Ast -> CompilerMonad ()) -> Text -> [(Text, Ast)] ->
    CompilerMonad ()
compileSetStruct :: (Ast -> CompilerMonad ())
-> Text -> [(Text, Ast)] -> CompilerMonad ()
compileSetStruct Ast -> CompilerMonad ()
compileFn Text
name [(Text, Ast)]
assignedFields = do
    defFields <- Text -> CompilerMonad [(Text, Text)]
getStructDefinition Text
name
    forM_ defFields $ \(Text
fName, Text
_fType) ->
        case Text -> Map Text Ast -> Maybe Ast
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
fName ([(Text, Ast)] -> Map Text Ast
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Text, Ast)]
assignedFields) of
            Just Ast
valAst -> Ast -> CompilerMonad ()
compileFn Ast
valAst
            Maybe Ast
Nothing -> Either Text () -> CompilerMonad ()
forall (m :: * -> *) a. Monad m => m a -> StateT CompilerState m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either Text () -> CompilerMonad ())
-> Either Text () -> CompilerMonad ()
forall a b. (a -> b) -> a -> b
$ Text -> Either Text ()
forall a b. a -> Either a b
Left (String -> Text
pack (
                String
"Missing field in struct instantiation: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
forall a. Show a => a -> String
show Text
fName))
    emitInstruction (BuildStruct (length defFields))

-- | Compiles a named function definition.
--
-- @args
--   - compileFn: The recursive compilation function.
--   - name: The function name.
--   - args: The list of parameter names.
--   - body: The function body AST.
--
-- @details
--   Generates a label for the function and compiles the body within an
--   isolated function scope. Arguments are registered as 'ScopeLocal'
--   symbols (indices 0 to N-1). The compiled code is stored in 'csFuncs'.
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
compileDefineFun :: (Ast -> CompilerMonad ()) -> Text -> [(Text, Text)] ->
    Ast -> CompilerMonad ()
compileDefineFun :: (Ast -> CompilerMonad ())
-> Text -> [(Text, Text)] -> Ast -> CompilerMonad ()
compileDefineFun Ast -> CompilerMonad ()
compileFn Text
name [(Text, Text)]
args Ast
body = let nArgs :: Int
nArgs = [(Text, Text)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Text, Text)]
args in
    CompilerMonad () -> CompilerMonad ()
compileInIsolatedFunctionScope (
        (CompilerState -> CompilerState) -> CompilerMonad ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\CompilerState
s -> CompilerState
s { csCurrentArgCount = nArgs }) CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
        Text -> CompilerMonad ()
emitLabelDefinition (String -> Text
pack String
"fun_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name) CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
        ((Text, Text) -> Int -> CompilerMonad ())
-> [(Text, Text)] -> [Int] -> CompilerMonad ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (\(Text
aName, Text
aType) Int
i ->
            Text -> Text -> ScopeType -> Int -> CompilerMonad ()
registerSymbol Text
aName Text
aType ScopeType
ScopeLocal (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
nArgs)) [(Text, Text)]
args [Int
0..] CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
        Ast -> CompilerMonad ()
compileFn Ast
body CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
        Instruction -> CompilerMonad ()
emitInstruction (Immediate -> Instruction
Push (IntValue -> Immediate
ImmInt (Int64 -> IntValue
I64 Int64
0))) CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
        Instruction -> CompilerMonad ()
emitInstruction (Int -> Instruction
Ret Int
nArgs)
    ) CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> CompilerMonad ()
forall a. a -> StateT CompilerState (Either Text) a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Registers lambda parameters with negative offsets.
--
-- @args
--   - params: The list of parameter names.
--
-- @details
--   Calculates indices relative to the Frame Pointer.
--   Arg 0 is at (0 - N), Last Arg is at -1.
--
registerLambdaParams :: [Text] -> CompilerMonad ()
registerLambdaParams :: [Text] -> CompilerMonad ()
registerLambdaParams [Text]
params =
    let count :: Int
count = [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
params in
    (Text -> Int -> CompilerMonad ())
-> [Text] -> [Int] -> CompilerMonad ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (\Text
p Int
i ->
        Text -> Text -> ScopeType -> Int -> CompilerMonad ()
registerSymbol Text
p (String -> Text
pack String
"auto") ScopeType
ScopeLocal (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
count)) [Text]
params [Int
0..]

-- | Internal helper to compile the isolated scope of a lambda.
--
-- @args
--   - compileFn: Recursion callback.
--   - params: Lambda parameters.
--   - body: Lambda body.
--   - label: The unique label for this function.
--   - fvars: List of captured variable names.
--   - types: List of captured variable types.
--
compileLambdaScope :: (Ast -> CompilerMonad ()) -> [Text] -> Ast -> Text ->
    [Text] -> [Text] -> CompilerMonad ()
compileLambdaScope :: (Ast -> CompilerMonad ())
-> [Text] -> Ast -> Text -> [Text] -> [Text] -> CompilerMonad ()
compileLambdaScope Ast -> CompilerMonad ()
compileFn [Text]
params Ast
body Text
label [Text]
fvars [Text]
types =
    Text -> CompilerMonad ()
emitLabelDefinition Text
label CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
    (Text -> Text -> Int -> CompilerMonad ())
-> [Text] -> [Text] -> [Int] -> CompilerMonad ()
forall (m :: * -> *) a b c d.
Monad m =>
(a -> b -> c -> m d) -> [a] -> [b] -> [c] -> m ()
zipWith3M_ (
        \Text
n Text
t Int
i -> Text -> Text -> ScopeType -> Int -> CompilerMonad ()
registerSymbol Text
n Text
t ScopeType
ScopeCapture Int
i) [Text]
fvars [Text]
types [Int
0..] CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
    [Text] -> CompilerMonad ()
registerLambdaParams [Text]
params CompilerMonad () -> CompilerMonad () -> CompilerMonad ()
forall a b.
StateT CompilerState (Either Text) a
-> StateT CompilerState (Either Text) b
-> StateT CompilerState (Either Text) b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
    (Ast -> CompilerMonad ()) -> Ast -> CompilerMonad ()
compileTail Ast -> CompilerMonad ()
compileFn Ast
body

-- | Compiles a lambda definition.
--
-- @args
--   - compileFn: The main compilation function (recursion).
--   - params: List of parameter names.
--   - body: The body of the lambda.
--
-- @details
--   1. Identifies free variables (captures).
--   2. Emits code to push captures onto the stack.
--   3. Generates a unique label for the lambda body.
--   4. Compiles the body in an isolated scope.
--   5. Emits 'MakeClosure' to create the closure object at runtime.
--
compileDefineLambda :: (Ast -> CompilerMonad ()) -> [Text] -> Ast ->
    CompilerMonad ()
compileDefineLambda :: (Ast -> CompilerMonad ()) -> [Text] -> Ast -> CompilerMonad ()
compileDefineLambda Ast -> CompilerMonad ()
compileFn [Text]
params Ast
body = do
    let ast :: Ast
ast = [Text] -> Ast -> Ast
ADefineLambda [Text]
params Ast
body
    let fvars :: [Text]
fvars = Set Text -> [Text]
forall a. Set a -> [a]
Set.toList (Ast -> Set Text
getLambdaFreeVariables Ast
ast)
    capTypes <- (Text -> CompilerMonad Text)
-> [Text] -> StateT CompilerState (Either Text) [Text]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Text -> CompilerMonad Text
getSymbolType [Text]
fvars
    mapM_ astSymbolToAsm fvars
    ulabel <- generateUniqueLabel (pack "lambda")
    compileInIsolatedFunctionScope $
        compileLambdaScope compileFn params body ulabel fvars capTypes
    appendPseudoInstruction (MakeClosureLabel ulabel (length fvars))

-- | Registers a structure definition.
--
-- @args
--   - name: The name of the structure.
--   - fields: The list of (Field Name, Field Type) pairs.
--
-- @details
--   Updates the 'CompilerState' to store the list of field names in order.
--   No bytecode is emitted for a structure definition (it is a compile-time
--   metadata operation).
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
compileDefineStruct :: Text -> [(Text, Text)] -> CompilerMonad ()
compileDefineStruct :: Text -> [(Text, Text)] -> CompilerMonad ()
compileDefineStruct Text
name [(Text, Text)]
fields = Text -> [(Text, Text)] -> CompilerMonad ()
defineStruct Text
name [(Text, Text)]
fields

-- | Infers the type of an AST expression.
--
-- @details
--   - ASymbol: Look up the variable in the symbol table.
--   - AAccessStruct: Recursively infer the parent object's type,
--   then look up the field's type.
--
inferType :: Ast -> CompilerMonad Text
inferType :: Ast -> CompilerMonad Text
inferType (APos Int
_ Int
_ Ast
ast) = Ast -> CompilerMonad Text
inferType Ast
ast
inferType (ASymbol Text
name) = Text -> CompilerMonad Text
getSymbolType Text
name
inferType (AInteger IntValue
_) = Text -> CompilerMonad Text
forall a. a -> StateT CompilerState (Either Text) a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> Text
pack String
"int")
inferType (ABool Bool
_) = Text -> CompilerMonad Text
forall a. a -> StateT CompilerState (Either Text) a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> Text
pack String
"bool")
inferType (ACall (ASymbol Text
name) [Ast]
args) = do
    argTypes <- (Ast -> CompilerMonad Text)
-> [Ast] -> StateT CompilerState (Either Text) [Text]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Ast -> CompilerMonad Text
inferType [Ast]
args
    case getBuiltinReturnType (unpack name) argTypes of
        Right Text
builtinType -> Text -> CompilerMonad Text
forall a. a -> StateT CompilerState (Either Text) a
forall (m :: * -> *) a. Monad m => a -> m a
return Text
builtinType
        Left Text
_ -> Text -> CompilerMonad Text
getSymbolType Text
name
inferType (AAccessStruct Ast
obj Text
field) = do
    parentType <- Ast -> CompilerMonad Text
inferType Ast
obj
    (_, fieldType) <- getStructField parentType field
    return fieldType
inferType Ast
_ = Either Text Text -> CompilerMonad Text
forall (m :: * -> *) a. Monad m => m a -> StateT CompilerState m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either Text Text -> CompilerMonad Text)
-> Either Text Text -> CompilerMonad Text
forall a b. (a -> b) -> a -> b
$ Text -> Either Text Text
forall a b. a -> Either a b
Left (String -> Text
pack
    String
"Cannot infer type of expression (too complex AST).")

-- | Compiles a struct field access.
--
-- @details
--   1. Infers the type of the object being accessed (e.g., "Player").
--   2. Looks up the field index in the struct definition.
--   3. Emits the code for the object followed by GET_STRUCT_FIELD.
--
compileAccessStruct :: (Ast -> CompilerMonad ()) -> Ast -> Text -> CompilerMonad ()
compileAccessStruct :: (Ast -> CompilerMonad ()) -> Ast -> Text -> CompilerMonad ()
compileAccessStruct Ast -> CompilerMonad ()
compileFn Ast
obj Text
fieldName = do
    structType <- Ast -> CompilerMonad Text
inferType Ast
obj
    (idx, _) <- getStructField structType fieldName
    compileFn obj
    emitInstruction (GetStructField idx)

extractChar :: Ast -> CompilerMonad Char
extractChar :: Ast -> CompilerMonad Char
extractChar (AInteger (IChar Int8
c)) = Char -> CompilerMonad Char
forall a. a -> StateT CompilerState (Either Text) a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Char
chr (Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
c))
extractChar Ast
_ =
    Either Text Char -> CompilerMonad Char
forall (m :: * -> *) a. Monad m => m a -> StateT CompilerState m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either Text Char -> CompilerMonad Char)
-> Either Text Char -> CompilerMonad Char
forall a b. (a -> b) -> a -> b
$ Text -> Either Text Char
forall a b. a -> Either a b
Left (String -> Text
pack String
"attr_update: Field name must be a string of characters")

-- | Convert an AST list of IChar to Haskell Text
--
extractFieldName :: Ast -> CompilerMonad Text
extractFieldName :: Ast -> CompilerMonad Text
extractFieldName (ASymbol Text
s) = Text -> CompilerMonad Text
forall a. a -> StateT CompilerState (Either Text) a
forall (m :: * -> *) a. Monad m => a -> m a
return Text
s
extractFieldName (AList [Ast]
xs) = do
    chars <- (Ast -> CompilerMonad Char)
-> [Ast] -> StateT CompilerState (Either Text) String
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Ast -> CompilerMonad Char
extractChar [Ast]
xs
    return (pack chars)
extractFieldName Ast
_ =
    Either Text Text -> CompilerMonad Text
forall (m :: * -> *) a. Monad m => m a -> StateT CompilerState m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either Text Text -> CompilerMonad Text)
-> Either Text Text -> CompilerMonad Text
forall a b. (a -> b) -> a -> b
$ Text -> Either Text Text
forall a b. a -> Either a b
Left (String -> Text
pack String
"attr_update: Second argument must be a Sym or String")

-- | Compiles the 'attr_update' builtin special form.
--
-- @args
--   - compileFn: Recursive compiler.
--   - obj: The AST for the structure object.
--   - fieldArg: The AST for the field name (must be a Symbol).
--   - v: The AST for the new value.
--
-- @details
--   1. Infers the type of 'obj' to find its struct definition.
--   2. Resolves 'fieldArg' to an integer index.
--   3. Pushes Obj, Index, and Value onto the stack.
--   4. Emits AttrUpdate.
--
compileAttrUpdate :: (Ast -> CompilerMonad ()) -> Ast -> Ast -> Ast ->
    CompilerMonad ()
compileAttrUpdate :: (Ast -> CompilerMonad ()) -> Ast -> Ast -> Ast -> CompilerMonad ()
compileAttrUpdate Ast -> CompilerMonad ()
compileFn Ast
obj Ast
fieldArg Ast
v = do
    fieldName <- Ast -> CompilerMonad Text
extractFieldName Ast
fieldArg
    structType <- inferType obj
    (idx, _) <- getStructField structType fieldName
    compileFn obj
    emitInstruction (Push (ImmInt (I64 (fromIntegral idx))))
    compileFn v
    emitInstruction AttrUpdate

-- | Pushes a loop exit label onto the stack.
--
pushLoopExit :: Text -> CompilerMonad ()
pushLoopExit :: Text -> CompilerMonad ()
pushLoopExit Text
label = (CompilerState -> CompilerState) -> CompilerMonad ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState -> CompilerState) -> CompilerMonad ())
-> (CompilerState -> CompilerState) -> CompilerMonad ()
forall a b. (a -> b) -> a -> b
$ \CompilerState
s -> CompilerState
s { csLoopExits = label : csLoopExits s }

-- | Pops the last loop exit label.
--
popLoopExit :: CompilerMonad ()
popLoopExit :: CompilerMonad ()
popLoopExit = (CompilerState -> CompilerState) -> CompilerMonad ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CompilerState -> CompilerState) -> CompilerMonad ())
-> (CompilerState -> CompilerState) -> CompilerMonad ()
forall a b. (a -> b) -> a -> b
$ \CompilerState
s -> case CompilerState -> [Text]
csLoopExits CompilerState
s of
    (Text
_:[Text]
xs) -> CompilerState
s { csLoopExits = xs }
    []     -> CompilerState
s

-- | Main Dispatcher Function: Compiles any AST node.
--
-- @args
--   - ast: The AST node to compile.
--
-- @details
--   This is the central router of the compiler. It pattern-matches on the AST
--   constructors and delegates to the appropriate specific compilation function.
--   It passes itself ('compileAst') recursively to handle nested expressions.
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
compileAst :: Ast -> CompilerMonad ()
compileAst :: Ast -> CompilerMonad ()
compileAst (APos Int
_ Int
_ Ast
ast) = Ast -> CompilerMonad ()
compileAst Ast
ast
compileAst (ABlock [Ast]
xs) = (Ast -> CompilerMonad ()) -> [Ast] -> CompilerMonad ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Ast -> CompilerMonad ()
compileAst [Ast]
xs
compileAst (AInteger IntValue
i) = IntValue -> CompilerMonad ()
astIntToAsm IntValue
i
compileAst (ABool Bool
b) = Bool -> CompilerMonad ()
astBoolToAsm Bool
b
compileAst (ASymbol Text
s) = Text -> CompilerMonad ()
astSymbolToAsm Text
s
compileAst Ast
AVoid = () -> CompilerMonad ()
forall a. a -> StateT CompilerState (Either Text) a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
compileAst (AList [Ast]
xs) = (Ast -> CompilerMonad ()) -> [Ast] -> CompilerMonad ()
astListToAsm Ast -> CompilerMonad ()
compileAst [Ast]
xs
compileAst (ADefineFunc Text
name [(Text, Text)]
args Text
_ Ast
body) =
    (Ast -> CompilerMonad ())
-> Text -> [(Text, Text)] -> Ast -> CompilerMonad ()
compileDefineFun Ast -> CompilerMonad ()
compileAst Text
name [(Text, Text)]
args Ast
body
compileAst (ADefineLambda [Text]
params Ast
body) =
    (Ast -> CompilerMonad ()) -> [Text] -> Ast -> CompilerMonad ()
compileDefineLambda Ast -> CompilerMonad ()
compileAst [Text]
params Ast
body
compileAst (ADefineStruct Text
name [(Text, Text)]
fields) = Text -> [(Text, Text)] -> CompilerMonad ()
compileDefineStruct Text
name [(Text, Text)]
fields
compileAst (AExprStmt (ASetVar Text
name Text
typeName Ast
body)) =
    (Ast -> CompilerMonad ())
-> Text -> Text -> Ast -> Bool -> CompilerMonad ()
compileSetVar Ast -> CompilerMonad ()
compileAst Text
name Text
typeName Ast
body Bool
True
compileAst (ASetVar Text
name Text
typeName Ast
body) =
    (Ast -> CompilerMonad ())
-> Text -> Text -> Ast -> Bool -> CompilerMonad ()
compileSetVar Ast -> CompilerMonad ()
compileAst Text
name Text
typeName Ast
body Bool
False
compileAst (ASetStruct Text
name [(Text, Ast)]
fields) = (Ast -> CompilerMonad ())
-> Text -> [(Text, Ast)] -> CompilerMonad ()
compileSetStruct Ast -> CompilerMonad ()
compileAst Text
name [(Text, Ast)]
fields
compileAst (AAccessStruct Ast
obj Text
field) = (Ast -> CompilerMonad ()) -> Ast -> Text -> CompilerMonad ()
compileAccessStruct Ast -> CompilerMonad ()
compileAst Ast
obj Text
field
compileAst (AIf Ast
cond Ast
t Ast
f) = (Ast -> CompilerMonad ()) -> Ast -> Ast -> Ast -> CompilerMonad ()
compileIf Ast -> CompilerMonad ()
compileAst Ast
cond Ast
t Ast
f
compileAst (AWhile Ast
cond Ast
body) = (Ast -> CompilerMonad ()) -> Ast -> Ast -> CompilerMonad ()
compileWhile Ast -> CompilerMonad ()
compileAst Ast
cond Ast
body
compileAst (AFor Ast
i Ast
cond Ast
u Ast
body) = (Ast -> CompilerMonad ())
-> Ast -> Ast -> Ast -> Ast -> CompilerMonad ()
compileFor Ast -> CompilerMonad ()
compileAst Ast
i Ast
cond Ast
u Ast
body
compileAst (ACall (ASymbol Text
name) [Ast
obj, Ast
fieldArg, Ast
v]) 
    | Text
name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Text
pack String
"attr_update" = (Ast -> CompilerMonad ()) -> Ast -> Ast -> Ast -> CompilerMonad ()
compileAttrUpdate Ast -> CompilerMonad ()
compileAst Ast
obj Ast
fieldArg Ast
v
compileAst (ACall Ast
func [Ast]
args) = (Ast -> CompilerMonad ()) -> Ast -> [Ast] -> CompilerMonad ()
astCallToAsm Ast -> CompilerMonad ()
compileAst Ast
func [Ast]
args
compileAst Ast
ABreak = do
    state <- StateT CompilerState (Either Text) CompilerState
forall s (m :: * -> *). MonadState s m => m s
get
    case csLoopExits state of
        [] -> Either Text () -> CompilerMonad ()
forall (m :: * -> *) a. Monad m => m a -> StateT CompilerState m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Either Text () -> CompilerMonad ())
-> Either Text () -> CompilerMonad ()
forall a b. (a -> b) -> a -> b
$ Text -> Either Text ()
forall a b. a -> Either a b
Left (String -> Text
pack String
"Error: 'break' statement outside of loop")
        (Text
label:[Text]
_) -> Text -> CompilerMonad ()
emitJumpToLabel Text
label
compileAst (AReturn Ast
expr) = do
    Ast -> CompilerMonad ()
compileAst Ast
expr
    s <- StateT CompilerState (Either Text) CompilerState
forall s (m :: * -> *). MonadState s m => m s
get
    emitInstruction (Ret (csCurrentArgCount s))
compileAst (AImport Text
_) = () -> CompilerMonad ()
forall a. a -> StateT CompilerState (Either Text) a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
compileAst (AExprStmt Ast
expr) = Ast -> CompilerMonad ()
compileAst Ast
expr