{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}

-- | A monadic builder for JVM bytecode instructions.
module H2JVM.Builder.Code (
    CodeBuilder,
    emit,
    addCodeAttribute,
    runCodeBuilder,
    newLabel,
)
where

import Data.List.NonEmpty (NonEmpty)
import Effectful
import Effectful.Dispatch.Dynamic
import Effectful.State.Static.Local
import Effectful.TH (makeEffect)

import Data.List.NonEmpty qualified as NE

import H2JVM.Builder.Label
import H2JVM.ClassFile.Method (CodeAttribute)
import H2JVM.Data.TypeMergingList (TypeMergingList)
import H2JVM.Instruction

import H2JVM.Data.TypeMergingList qualified as TML

-- | The 'CodeBuilder' effect, which allows emitting instructions and adding code attributes in a monadic style.
data CodeBuilder m a where
    -- | Add a code attribute to the current code block. If an attribute of the same type already exists, it will be merged with the existing one using the 'H2JVM.Data.TypeMergingList.merge' function from the 'H2JVM.Data.TypeMergingList.DataMergeable' instance for that attribute type.
    AddCodeAttribute ::
        -- | the code attribute to add
        CodeAttribute ->
        CodeBuilder m ()
    -- | Create a new, unique 'Label'.
    NewLabel ::
        CodeBuilder m Label
    Emit' :: [Instruction] -> CodeBuilder m ()

makeEffect ''CodeBuilder

data CodeState = CodeState
    { CodeState -> Word
currentLabel :: {-# UNPACK #-} !Word
    -- ^ the current label index.
    , CodeState -> TypeMergingList CodeAttribute
attributes :: TypeMergingList CodeAttribute
    -- ^ the code attributes that have been added to this code block.
    , CodeState -> [Instruction]
code :: [Instruction]
    -- ^ the code that has been emitted so far. Note that this is stored in reverse order for efficient appending, so it should be reversed before being returned.
    }

-- | An empty code state, with an infinite supply of labels.
initialCodeState :: CodeState
initialCodeState :: CodeState
initialCodeState = CodeState{currentLabel :: Word
currentLabel = Word
0, attributes :: TypeMergingList CodeAttribute
attributes = TypeMergingList CodeAttribute
forall a. Monoid a => a
mempty, code :: [Instruction]
code = []}

-- | Emit a single instruction.
emit :: CodeBuilder :> r => Instruction -> Eff r ()
emit :: forall (r :: [Effect]).
(CodeBuilder :> r) =>
Instruction -> Eff r ()
emit = [Instruction] -> Eff r ()
forall {k} (es :: [Effect]).
(HasCallStack, CodeBuilder :> es) =>
[Instruction] -> Eff es ()
emit' ([Instruction] -> Eff r ())
-> (Instruction -> [Instruction]) -> Instruction -> Eff r ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure @[]

-- | Re-interpret a 'CodeBuilder' effect as a 'State' effect, accumulating the emitted code and attributes in the state.
codeBuilderToState :: Eff (CodeBuilder : es) a -> Eff es (a, CodeState)
codeBuilderToState :: forall (es :: [Effect]) a.
Eff (CodeBuilder : es) a -> Eff es (a, CodeState)
codeBuilderToState =
    (Eff (State CodeState : es) a -> Eff es (a, CodeState))
-> EffectHandler CodeBuilder (State CodeState : es)
-> Eff (CodeBuilder : es) a
-> Eff es (a, CodeState)
forall (e :: Effect) (handlerEs :: [Effect]) a (es :: [Effect]) b.
(HasCallStack, DispatchOf e ~ 'Dynamic) =>
(Eff handlerEs a -> Eff es b)
-> EffectHandler e handlerEs -> Eff (e : es) a -> Eff es b
reinterpret
        (CodeState -> Eff (State CodeState : es) a -> Eff es (a, CodeState)
forall s (es :: [Effect]) a.
HasCallStack =>
s -> Eff (State s : es) a -> Eff es (a, s)
runState CodeState
initialCodeState)
        ( \LocalEnv localEs (State CodeState : es)
_ -> \case
            AddCodeAttribute CodeAttribute
ca -> (CodeState -> CodeState) -> Eff (State CodeState : es) ()
forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
(s -> s) -> Eff es ()
modify ((CodeState -> CodeState) -> Eff (State CodeState : es) ())
-> (CodeState -> CodeState) -> Eff (State CodeState : es) ()
forall a b. (a -> b) -> a -> b
$ \CodeState
s -> CodeState
s{attributes = s.attributes `TML.snoc` ca}
            CodeBuilder (Eff localEs) a
NewLabel -> do
                c <- (CodeState -> Word) -> Eff (State CodeState : es) Word
forall s (es :: [Effect]) a.
(HasCallStack, State s :> es) =>
(s -> a) -> Eff es a
gets (.currentLabel)
                modify $ \CodeState
s -> CodeState
s{currentLabel = c + 1}
                pure (unsafeMkLabel c)
            Emit' [Instruction]
is -> (CodeState -> CodeState) -> Eff (State CodeState : es) ()
forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
(s -> s) -> Eff es ()
modify ((CodeState -> CodeState) -> Eff (State CodeState : es) ())
-> (CodeState -> CodeState) -> Eff (State CodeState : es) ()
forall a b. (a -> b) -> a -> b
$ \CodeState
s -> CodeState
s{code = reverse is <> s.code}
        )

{- | Run a 'CodeBuilder' effect, returning the emitted code and attributes.
Throws an imprecise exception if no code was emitted, since a code block with no instructions is invalid.
-}
runCodeBuilder ::
    forall r a.
    HasCallStack =>
    Eff (CodeBuilder ': r) a ->
    -- | the result of the computation, the list of emitted code attributes, and the list of emitted instructions, in the order they were emitted.
    Eff r (a, [CodeAttribute], NonEmpty Instruction)
runCodeBuilder :: forall (r :: [Effect]) a.
HasCallStack =>
Eff (CodeBuilder : r) a
-> Eff r (a, [CodeAttribute], NonEmpty Instruction)
runCodeBuilder =
    ((a, CodeState) -> (a, [CodeAttribute], NonEmpty Instruction))
-> Eff r (a, CodeState)
-> Eff r (a, [CodeAttribute], NonEmpty Instruction)
forall a b. (a -> b) -> Eff r a -> Eff r b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, CodeState) -> (a, [CodeAttribute], NonEmpty Instruction)
forall {r} {a} {a} {a}.
(HasField "code" r [a],
 HasField "attributes" r (TypeMergingList a)) =>
(a, r) -> (a, [a], NonEmpty a)
formatResult (Eff r (a, CodeState)
 -> Eff r (a, [CodeAttribute], NonEmpty Instruction))
-> (Eff (CodeBuilder : r) a -> Eff r (a, CodeState))
-> Eff (CodeBuilder : r) a
-> Eff r (a, [CodeAttribute], NonEmpty Instruction)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Eff (CodeBuilder : r) a -> Eff r (a, CodeState)
forall (es :: [Effect]) a.
Eff (CodeBuilder : es) a -> Eff es (a, CodeState)
codeBuilderToState
  where
    formatResult :: (a, r) -> (a, [a], NonEmpty a)
formatResult (a
res, r
s) =
        ( a
res
        , TypeMergingList a -> [a]
forall a. TypeMergingList a -> [a]
TML.toList r
s.attributes
        , case [a] -> [a]
forall a. [a] -> [a]
reverse r
s.code of
            [] -> [Char] -> NonEmpty a
forall a. HasCallStack => [Char] -> a
error [Char]
"runCodeBuilder: No code emitted"
            [a]
is -> [a] -> NonEmpty a
forall a. HasCallStack => [a] -> NonEmpty a
NE.fromList [a]
is
        )