{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE UndecidableInstances #-}

module Codegen.State where

import qualified Codegen.Errors as CC
import qualified Codegen.Utils as CU
import qualified Control.Monad.Except as E
import qualified Control.Monad.Fix as F
import qualified Control.Monad.State as S
import qualified LLVM.AST as AST
import qualified LLVM.IRBuilder.Module as M
import qualified LLVM.IRBuilder.Monad as IRM

type MonadCodegen m =
  ( IRM.MonadIRBuilder m,
    M.MonadModuleBuilder m,
    F.MonadFix m,
    S.MonadState CodegenState m,
    E.MonadError CC.CodegenError m
  )

-- | Type alias for the local code generation state.
type LocalState = [(String, AST.Operand)]

-- | Type alias for the global code generation state.
type GlobalState = [(String, AST.Operand)]

-- | Type alias for the loop code generation state.
type LoopState = Maybe (AST.Name, AST.Name)

-- | Type alias for the variables name .
type UniqueNameState = Integer

-- | Combined state for code generation.
data CodegenState = CodegenState
  { CodegenState -> LocalState
localState :: LocalState,
    CodegenState -> LocalState
globalState :: GlobalState,
    CodegenState -> LoopState
loopState :: LoopState,
    CodegenState -> LocalState
allocatedVars :: LocalState,
    CodegenState -> UniqueNameState
uniqueNameState :: UniqueNameState
  }
  deriving (Int -> CodegenState -> ShowS
[CodegenState] -> ShowS
CodegenState -> String
(Int -> CodegenState -> ShowS)
-> (CodegenState -> String)
-> ([CodegenState] -> ShowS)
-> Show CodegenState
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CodegenState -> ShowS
showsPrec :: Int -> CodegenState -> ShowS
$cshow :: CodegenState -> String
show :: CodegenState -> String
$cshowList :: [CodegenState] -> ShowS
showList :: [CodegenState] -> ShowS
Show)

-- | Variable binding typeclass.
class (Monad m) => VarBinding m where
  getVar :: String -> m (Maybe AST.Operand)
  addVar :: String -> AST.Operand -> m ()
  getGlobalVar :: String -> m (Maybe AST.Operand)
  addGlobalVar :: String -> AST.Operand -> m ()

instance (S.MonadState CodegenState m, Monad m) => VarBinding m where
  getVar :: (S.MonadState CodegenState m, Monad m) => String -> m (Maybe AST.Operand)
  getVar :: (MonadState CodegenState m, Monad m) => String -> m (Maybe Operand)
getVar String
name = do
    CodegenState
state <- m CodegenState
forall s (m :: * -> *). MonadState s m => m s
S.get
    Maybe Operand -> m (Maybe Operand)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Operand -> m (Maybe Operand))
-> Maybe Operand -> m (Maybe Operand)
forall a b. (a -> b) -> a -> b
$
      String -> LocalState -> Maybe Operand
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
name (CodegenState -> LocalState
allocatedVars CodegenState
state)
        Maybe Operand -> Maybe Operand -> Maybe Operand
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`S.mplus` String -> LocalState -> Maybe Operand
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
name (CodegenState -> LocalState
localState CodegenState
state)
        Maybe Operand -> Maybe Operand -> Maybe Operand
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`S.mplus` String -> LocalState -> Maybe Operand
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
name (CodegenState -> LocalState
globalState CodegenState
state)

  addVar :: (S.MonadState CodegenState m, Monad m) => String -> AST.Operand -> m ()
  addVar :: (MonadState CodegenState m, Monad m) => String -> Operand -> m ()
addVar String
name Operand
operand = (CodegenState -> CodegenState) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
S.modify (\CodegenState
s -> CodegenState
s {localState :: LocalState
localState = (String
name, Operand
operand) (String, Operand) -> LocalState -> LocalState
forall a. a -> [a] -> [a]
: CodegenState -> LocalState
localState CodegenState
s})

  getGlobalVar :: (S.MonadState CodegenState m, Monad m) => String -> m (Maybe AST.Operand)
  getGlobalVar :: (MonadState CodegenState m, Monad m) => String -> m (Maybe Operand)
getGlobalVar String
name = (CodegenState -> Maybe Operand) -> m (Maybe Operand)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
S.gets (String -> LocalState -> Maybe Operand
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
name (LocalState -> Maybe Operand)
-> (CodegenState -> LocalState) -> CodegenState -> Maybe Operand
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CodegenState -> LocalState
globalState)

  addGlobalVar :: (S.MonadState CodegenState m, Monad m) => String -> AST.Operand -> m ()
  addGlobalVar :: (MonadState CodegenState m, Monad m) => String -> Operand -> m ()
addGlobalVar String
name Operand
operand = (CodegenState -> CodegenState) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
S.modify (\CodegenState
s -> CodegenState
s {globalState :: LocalState
globalState = (String
name, Operand
operand) (String, Operand) -> LocalState -> LocalState
forall a. a -> [a] -> [a]
: CodegenState -> LocalState
globalState CodegenState
s})

-- Generates a fresh unique name.
fresh :: (S.MonadState CodegenState m) => m AST.Name
fresh :: forall (m :: * -> *). MonadState CodegenState m => m Name
fresh = do
  CodegenState
state <- m CodegenState
forall s (m :: * -> *). MonadState s m => m s
S.get
  let uniqueName :: UniqueNameState
uniqueName = CodegenState -> UniqueNameState
uniqueNameState CodegenState
state
  CodegenState -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
S.put (CodegenState -> m ()) -> CodegenState -> m ()
forall a b. (a -> b) -> a -> b
$ CodegenState
state {uniqueNameState :: UniqueNameState
uniqueNameState = UniqueNameState
uniqueName UniqueNameState -> UniqueNameState -> UniqueNameState
forall a. Num a => a -> a -> a
+ UniqueNameState
1}
  let fullName :: String
fullName = String
"_" String -> ShowS
forall a. [a] -> [a] -> [a]
++ UniqueNameState -> String
forall a. Show a => a -> String
show UniqueNameState
uniqueName
  Name -> m Name
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> m Name) -> Name -> m Name
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Name
AST.Name (String -> ShortByteString
CU.stringToByteString String
fullName)

-- Generates a fresh unique name with the given prefix.
freshName :: (S.MonadState CodegenState m) => String -> m AST.Name
freshName :: forall (m :: * -> *). MonadState CodegenState m => String -> m Name
freshName String
prefix = do
  CodegenState
state <- m CodegenState
forall s (m :: * -> *). MonadState s m => m s
S.get
  let uniqueName :: UniqueNameState
uniqueName = CodegenState -> UniqueNameState
uniqueNameState CodegenState
state
  CodegenState -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
S.put (CodegenState -> m ()) -> CodegenState -> m ()
forall a b. (a -> b) -> a -> b
$ CodegenState
state {uniqueNameState :: UniqueNameState
uniqueNameState = UniqueNameState
uniqueName UniqueNameState -> UniqueNameState -> UniqueNameState
forall a. Num a => a -> a -> a
+ UniqueNameState
1}
  let fullName :: String
fullName = String
prefix String -> ShowS
forall a. [a] -> [a] -> [a]
++ UniqueNameState -> String
forall a. Show a => a -> String
show UniqueNameState
uniqueName
  Name -> m Name
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> m Name) -> Name -> m Name
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Name
AST.Name (String -> ShortByteString
CU.stringToByteString String
fullName)