{-
-- EPITECH PROJECT, 2026
-- Glados
-- File description:
-- ResolveLabels
-}

module Compiler.ResolveLabels.ResolveLabels (resolveLabels) where

import Compiler.ResolveLabels.Step2CallLabel (step2CallLabel)
import Compiler.ResolveLabels.Step2GetFuncAddrLabel (step2GetFuncAddrLabel)
import Compiler.ResolveLabels.Step2JumpIfFalseLabel (step2JumpIfFalseLabel)
import Compiler.ResolveLabels.Step2JumpIfTrueLabel (step2JumpIfTrueLabel)
import Compiler.ResolveLabels.Step2JumpLabel (step2JumpLabel)
import Compiler.ResolveLabels.Step2LabelDef (step2LabelDef)
import Compiler.ResolveLabels.Step2MakeClosureLabel (step2MakeClosureLabel)
import Compiler.ResolveLabels.Step2Real (step2Real)
import Compiler.ResolveLabels.Step2TailCallLabel (step2TailCallLabel)
import Compiler.Instruction (instructionSize, Instruction(..))
import Compiler.ResolveLabels.ResolveLabelsHelpers
    ( sizeOfJumpInst
    , sizeOfJumpIfFalseInst
    , sizeOfJumpIfTrueInst
    , sizeOfCallInst
    , sizeOfTailCallInst
    )
import Compiler.PsInstruction (PsInstruction(..))
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Data.Text (Text)
import Control.Monad (foldM)
import qualified Data.Text as T

resolveLabels :: [PsInstruction] -> Either Text [Instruction]
resolveLabels :: [PsInstruction] -> Either Text [Instruction]
resolveLabels [PsInstruction]
pseudos =
    ((Int, Map Text Int, Set Int)
 -> PsInstruction -> Either Text (Int, Map Text Int, Set Int))
-> (Int, Map Text Int, Set Int)
-> [PsInstruction]
-> Either Text (Int, Map Text Int, Set Int)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Int, Map Text Int, Set Int)
-> PsInstruction -> Either Text (Int, Map Text Int, Set Int)
step1 (Int
0, Map Text Int
forall k a. Map k a
Map.empty, Set Int
forall a. Set a
Set.empty) [PsInstruction]
pseudos Either Text (Int, Map Text Int, Set Int)
-> ((Int, Map Text Int, Set Int) -> Either Text [Instruction])
-> Either Text [Instruction]
forall a b. Either Text a -> (a -> Either Text b) -> Either Text b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
        \(Int
_, Map Text Int
labelMap, Set Int
startsSet) ->
        (([Instruction], Int) -> [Instruction])
-> Either Text ([Instruction], Int) -> Either Text [Instruction]
forall a b. (a -> b) -> Either Text a -> Either Text b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Instruction], Int) -> [Instruction]
forall a b. (a, b) -> a
fst (Either Text ([Instruction], Int) -> Either Text [Instruction])
-> Either Text ([Instruction], Int) -> Either Text [Instruction]
forall a b. (a -> b) -> a -> b
$ (([Instruction], Int)
 -> PsInstruction -> Either Text ([Instruction], Int))
-> ([Instruction], Int)
-> [PsInstruction]
-> Either Text ([Instruction], Int)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Map Text Int
-> Set Int
-> ([Instruction], Int)
-> PsInstruction
-> Either Text ([Instruction], Int)
step2 Map Text Int
labelMap Set Int
startsSet) ([], Int
0) [PsInstruction]
pseudos

step1 :: (Int, Map.Map Text Int, Set.Set Int) -> PsInstruction
    -> Either Text (Int, Map.Map Text Int, Set.Set Int)
step1 :: (Int, Map Text Int, Set Int)
-> PsInstruction -> Either Text (Int, Map Text Int, Set Int)
step1 (Int
idx, Map Text Int
m, Set Int
s) (LabelDef Text
name) = case Int -> Text -> Map Text Int -> Either Text (Int, Map Text Int)
detectDuplicateLabels Int
idx Text
name Map Text Int
m of
    Left Text
err -> Text -> Either Text (Int, Map Text Int, Set Int)
forall a b. a -> Either a b
Left Text
err
    Right (Int
_, Map Text Int
newM) -> (Int, Map Text Int, Set Int)
-> Either Text (Int, Map Text Int, Set Int)
forall a b. b -> Either a b
Right (Int
idx, Map Text Int
newM, Set Int
s)
step1 (Int
idx, Map Text Int
m, Set Int
s) PsInstruction
pseudo = (Int, Map Text Int, Set Int)
-> Either Text (Int, Map Text Int, Set Int)
forall a b. b -> Either a b
Right (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ PsInstruction -> Int
pseudoSize PsInstruction
pseudo, Map Text Int
m, Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
Set.insert Int
idx Set Int
s)

pseudoSize :: PsInstruction -> Int
pseudoSize :: PsInstruction -> Int
pseudoSize (Real Instruction
instr) = Instruction -> Int
instructionSize Instruction
instr
pseudoSize (JumpLabel Text
_) = Int
sizeOfJumpInst
pseudoSize (JumpIfFalseLabel Text
_) = Int
sizeOfJumpIfFalseInst
pseudoSize (JumpIfTrueLabel Text
_) = Int
sizeOfJumpIfTrueInst
pseudoSize (CallLabel Text
_) = Int
sizeOfCallInst
pseudoSize (TailCallLabel Text
_) = Int
sizeOfTailCallInst
pseudoSize (MakeClosureLabel Text
_ Int
_) = Instruction -> Int
instructionSize (Int -> Int -> Instruction
MakeClosure Int
0 Int
0)
pseudoSize (GetFuncAddrLabel Text
_) = Instruction -> Int
instructionSize (Int -> Instruction
GetFuncAddr Int
0)
pseudoSize (LabelDef Text
_) = Int
0

detectDuplicateLabels :: Int -> Text -> Map.Map Text Int
    -> Either Text (Int, Map.Map Text Int)
detectDuplicateLabels :: Int -> Text -> Map Text Int -> Either Text (Int, Map Text Int)
detectDuplicateLabels Int
idx Text
name Map Text Int
m =
    case Text -> Map Text Int -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member Text
name Map Text Int
m of
        Bool
True  -> Text -> Either Text (Int, Map Text Int)
forall a b. a -> Either a b
Left (String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ String
"Duplicate label: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Text -> String
T.unpack Text
name)
        Bool
False -> (Int, Map Text Int) -> Either Text (Int, Map Text Int)
forall a b. b -> Either a b
Right (Int
idx, Text -> Int -> Map Text Int -> Map Text Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
name Int
idx Map Text Int
m)

step2 :: Map.Map Text Int -> Set.Set Int -> ([Instruction], Int)
    -> PsInstruction -> Either Text ([Instruction], Int)
step2 :: Map Text Int
-> Set Int
-> ([Instruction], Int)
-> PsInstruction
-> Either Text ([Instruction], Int)
step2 Map Text Int
_ Set Int
_ ([Instruction]
out, Int
idx) (Real Instruction
instr) = [Instruction]
-> Int -> Instruction -> Either Text ([Instruction], Int)
step2Real [Instruction]
out Int
idx Instruction
instr
step2 Map Text Int
_ Set Int
_ ([Instruction]
out, Int
idx) (LabelDef Text
_) = [Instruction] -> Int -> Either Text ([Instruction], Int)
step2LabelDef [Instruction]
out Int
idx
step2 Map Text Int
lm Set Int
ss ([Instruction]
out, Int
idx) PsInstruction
pseudo = Map Text Int
-> Set Int
-> [Instruction]
-> Int
-> PsInstruction
-> Either Text ([Instruction], Int)
step2Pseudo Map Text Int
lm Set Int
ss [Instruction]
out Int
idx PsInstruction
pseudo

step2Pseudo :: Map.Map Text Int -> Set.Set Int -> [Instruction]
    -> Int -> PsInstruction -> Either Text ([Instruction], Int)
step2Pseudo :: Map Text Int
-> Set Int
-> [Instruction]
-> Int
-> PsInstruction
-> Either Text ([Instruction], Int)
step2Pseudo Map Text Int
lm Set Int
ss [Instruction]
out Int
idx (JumpLabel Text
name) = Map Text Int
-> Set Int
-> [Instruction]
-> Int
-> Text
-> Either Text ([Instruction], Int)
step2JumpLabel Map Text Int
lm Set Int
ss [Instruction]
out Int
idx Text
name
step2Pseudo Map Text Int
lm Set Int
ss [Instruction]
out Int
idx (JumpIfFalseLabel Text
name) =
    Map Text Int
-> Set Int
-> [Instruction]
-> Int
-> Text
-> Either Text ([Instruction], Int)
step2JumpIfFalseLabel Map Text Int
lm Set Int
ss [Instruction]
out Int
idx Text
name
step2Pseudo Map Text Int
lm Set Int
ss [Instruction]
out Int
idx (JumpIfTrueLabel Text
name) =
    Map Text Int
-> Set Int
-> [Instruction]
-> Int
-> Text
-> Either Text ([Instruction], Int)
step2JumpIfTrueLabel Map Text Int
lm Set Int
ss [Instruction]
out Int
idx Text
name
step2Pseudo Map Text Int
lm Set Int
ss [Instruction]
out Int
idx PsInstruction
other = Map Text Int
-> Set Int
-> [Instruction]
-> Int
-> PsInstruction
-> Either Text ([Instruction], Int)
step2PseudoRest Map Text Int
lm Set Int
ss [Instruction]
out Int
idx PsInstruction
other

step2PseudoRest :: Map.Map Text Int -> Set.Set Int -> [Instruction]
    -> Int -> PsInstruction -> Either Text ([Instruction], Int)
step2PseudoRest :: Map Text Int
-> Set Int
-> [Instruction]
-> Int
-> PsInstruction
-> Either Text ([Instruction], Int)
step2PseudoRest Map Text Int
lm Set Int
ss [Instruction]
out Int
idx (CallLabel Text
name) =
    Map Text Int
-> Set Int
-> [Instruction]
-> Int
-> Text
-> Either Text ([Instruction], Int)
step2CallLabel Map Text Int
lm Set Int
ss [Instruction]
out Int
idx Text
name
step2PseudoRest Map Text Int
lm Set Int
ss [Instruction]
out Int
idx (TailCallLabel Text
name) =
    Map Text Int
-> Set Int
-> [Instruction]
-> Int
-> Text
-> Either Text ([Instruction], Int)
step2TailCallLabel Map Text Int
lm Set Int
ss [Instruction]
out Int
idx Text
name
step2PseudoRest Map Text Int
lm Set Int
ss [Instruction]
out Int
idx (MakeClosureLabel Text
name Int
n) =
    Map Text Int
-> Set Int
-> [Instruction]
-> Int
-> Text
-> Int
-> Either Text ([Instruction], Int)
step2MakeClosureLabel Map Text Int
lm Set Int
ss [Instruction]
out Int
idx Text
name Int
n
step2PseudoRest Map Text Int
lm Set Int
ss [Instruction]
out Int
idx (GetFuncAddrLabel Text
name) =
    Map Text Int
-> Set Int
-> [Instruction]
-> Int
-> Text
-> Either Text ([Instruction], Int)
step2GetFuncAddrLabel Map Text Int
lm Set Int
ss [Instruction]
out Int
idx Text
name
step2PseudoRest Map Text Int
_ Set Int
_ [Instruction]
out Int
idx PsInstruction
_ = ([Instruction], Int) -> Either Text ([Instruction], Int)
forall a b. b -> Either a b
Right ([Instruction]
out, Int
idx)