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

{-|
Module : Compiler.ASM.AstToAsm
Description : Atomic converters from AST nodes to Assembly instructions.
Stability : experimental
-}
module Compiler.ASM.AstToAsm
    ( astIntToAsm
    , astBoolToAsm
    , astSymbolToAsm
    , astListToAsm
    , astCallToAsm
    , builtinMap
    ) where

import Control.Monad.State (get, lift)
import Data.Text (Text, pack)
import qualified Data.Map.Strict as Map

import Compiler.ASM.Builtins (builtinMap)
import Compiler.ASM.CompilerMonad
    ( CompilerMonad
    , emitInstruction
    , emitCallToLabel
    , lookupSymbol
    )
import Compiler.CompilerState (CompilerState(..), ScopeType(..))
import Compiler.Instruction (Instruction(..), Immediate(..))
import Common.Type.Integer (IntValue(..))
import AST.Ast (Ast(..))

-- | Pushes an Integer AST value directly to the stack.
--
-- @args
--   - val: The IntValue from the AST.
--
-- @details
--   Since the AST construction (Parser) already uses 'fitInteger' to determine
--   the optimal storage size (I8, I16, etc.), we simply wrap the value in
--   an 'ImmInt' and emit the PUSH instruction.
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
astIntToAsm :: IntValue -> CompilerMonad ()
astIntToAsm :: IntValue -> CompilerMonad ()
astIntToAsm IntValue
v = Instruction -> CompilerMonad ()
emitInstruction (Immediate -> Instruction
Push (IntValue -> Immediate
ImmInt IntValue
v))

-- | Converts a native Haskell Bool to a PUSH instruction.
--
-- @args
--   - b: The boolean to convert.
--
-- @details
--   Emits a Push instruction with an immediate boolean value.
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
astBoolToAsm :: Bool -> CompilerMonad ()
astBoolToAsm :: Bool -> CompilerMonad ()
astBoolToAsm Bool
b = Instruction -> CompilerMonad ()
emitInstruction (Immediate -> Instruction
Push (Bool -> Immediate
ImmBool Bool
b))

-- | Converts a symbol (variable) to a LOAD instruction.
--
-- @args
--   - name: The name of the symbol.
--
-- @details
--   Looks up the symbol in the 'csSymbols' table. Based on the associated
--   'ScopeType', it emits either 'LoadGlobal', 'LoadLocal', or 'LoadCapture'
--   with the correct index. Returns an error if the symbol is undefined.
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
astSymbolToAsm :: Text -> CompilerMonad ()
astSymbolToAsm :: Text -> CompilerMonad ()
astSymbolToAsm Text
name = do
    currentState <- StateT CompilerState (Either Text) CompilerState
forall s (m :: * -> *). MonadState s m => m s
get
    case Map.lookup name (csSymbols currentState) of
        Just (ScopeType
ScopeGlobal, Int
idx, Text
_)  -> Instruction -> CompilerMonad ()
emitInstruction (Int -> Instruction
LoadGlobal Int
idx)
        Just (ScopeType
ScopeLocal, Int
idx, Text
_)   -> Instruction -> CompilerMonad ()
emitInstruction (Int -> Instruction
LoadLocal Int
idx)
        Just (ScopeType
ScopeCapture, Int
idx, Text
_) -> Instruction -> CompilerMonad ()
emitInstruction (Int -> Instruction
LoadCapture Int
idx)
        Maybe (ScopeType, Int, Text)
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
"Undefined symbol: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name)

-- | Converts a literal AST list to assembly instructions.
--
-- @args
--   - compileFn: The recursive compilation function.
--   - elements: The list of AST nodes to compile.
--
-- @details
--   Compiles each element pushing them onto the stack, then calls the
--   builtin "list" function to construct the list object.
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
astListToAsm :: (Ast -> CompilerMonad ()) -> [Ast] -> CompilerMonad ()
astListToAsm :: (Ast -> CompilerMonad ()) -> [Ast] -> CompilerMonad ()
astListToAsm Ast -> CompilerMonad ()
compileFn [Ast]
elements = (Ast -> CompilerMonad ()) -> [Ast] -> CompilerMonad ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Ast -> CompilerMonad ()
compileFn [Ast]
elements 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
BuildList ([Ast] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Ast]
elements))

-- | Compiles a function call (Builtin or User-defined).
--
-- @args
--   - compileFn: The recursive compilation function.
--   - callee: The AST representing the function being called.
--   - args: The list of argument ASTs.
--
-- @details
--   If the callee is a known builtin, emits the specific instruction.
--   Otherwise, emits a 'CallLabel' to the named function.
--   (Higher-order calls are explicitly not supported in this version).
--
-- @return
--   Unit value wrapped in 'CompilerMonad'.
--
astCallToAsm :: (Ast -> CompilerMonad ()) -> Ast -> [Ast] -> CompilerMonad ()
astCallToAsm :: (Ast -> CompilerMonad ()) -> Ast -> [Ast] -> CompilerMonad ()
astCallToAsm Ast -> CompilerMonad ()
compileFn (ASymbol Text
name) [Ast
a, Ast
b] | Text
name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Text
pack String
">" =
    Ast -> CompilerMonad ()
compileFn Ast
b 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
a 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 Instruction
Lt
astCallToAsm Ast -> CompilerMonad ()
compileFn (ASymbol Text
name) [Ast
a, Ast
b] | Text
name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Text
pack String
">=" =
    Ast -> CompilerMonad ()
compileFn Ast
b 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
a 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 Instruction
Le
astCallToAsm Ast -> CompilerMonad ()
compileFn (ASymbol Text
name) [Ast
a, Ast
b] | Text
name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Text
pack String
"neq?" =
    Ast -> CompilerMonad ()
compileFn Ast
a 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
b 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 Instruction
Eq 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 Instruction
Not
astCallToAsm Ast -> CompilerMonad ()
compileFn (ASymbol Text
name) [Ast
a, Ast
b] | Text
name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Text
pack String
"tneq?" =
    Ast -> CompilerMonad ()
compileFn Ast
a 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
b 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 Instruction
TEq 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 Instruction
Not
astCallToAsm Ast -> CompilerMonad ()
compileFn (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 -> (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
>> Instruction -> CompilerMonad ()
emitInstruction Instruction
instr
    Maybe Instruction
Nothing -> do
        sym <- Text -> CompilerMonad (Maybe (ScopeType, Int, Text))
lookupSymbol Text
name
        case sym of
            Just (ScopeType, Int, Text)
_ -> (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
>> Ast -> CompilerMonad ()
compileFn (Text -> Ast
ASymbol 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
>>
                Instruction -> CompilerMonad ()
emitInstruction Instruction
CallIndirect
            Maybe (ScopeType, Int, Text)
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
>>
                Text -> CompilerMonad ()
emitCallToLabel (String -> Text
pack String
"fun_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name)
astCallToAsm Ast -> CompilerMonad ()
_ (AInteger IntValue
_) [Ast]
_ =
    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: Cannot call an Integer")
astCallToAsm Ast -> CompilerMonad ()
_ (ABool Bool
_) [Ast]
_ =
    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: Cannot call a Boolean")
astCallToAsm Ast -> CompilerMonad ()
_ Ast
AVoid [Ast]
_ =
    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: Cannot call Void")
astCallToAsm Ast -> CompilerMonad ()
compileFn Ast
expr [Ast]
args =
    (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
>> Ast -> CompilerMonad ()
compileFn Ast
expr 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 Instruction
CallIndirect