content: {-# LANGUAGE UndecidableInstances, FunctionalDependencies #-}
  
  module Eclair.LLVM.Template
    ( TemplateT
    , Template
    , HasSuffix(..)
    , MonadTemplate(..)
    , Suffix
    , cmapParams
    , instantiate
    , partialInstantiate
    , function
    , typedef
    ) where
  
  
  import Control.Monad.Morph
  import LLVM.Codegen hiding (function, typedef)
  import qualified LLVM.Codegen as CG
  import qualified Control.Monad.State.Lazy as LazyState
  import qualified Control.Monad.State.Strict as StrictState
  import qualified Control.Monad.RWS.Lazy as LazyRWS
  import qualified Control.Monad.RWS.Strict as StrictRWS
  import Eclair.LLVM.Config
  
  
  type Suffix = Text
  
  -- | A MTL-like monad transformer that allows generating code in a way similar to C++ templates.
  -- Instead of complicated machinery in the compiler, this transformer just adds a suffix to all generated functions and types.
  -- It is up to the programmer to make sure all provided suffixes to one template are unique!
  -- The type variable 'p' is short for "template parameters" and can be used to tweak (specialize) the code generation.
  -- The type variable 'm' allows running this stack in a pure context, or in a stack on top of IO.
  newtype TemplateT p m a
    = TemplateT
    { unTemplateT :: ReaderT (Suffix, p) (ModuleBuilderT m) a
    } deriving (Functor, Applicative, Monad, MonadFix, MonadIO, MonadError e, MonadState s, MonadModuleBuilder)
    via ReaderT (Suffix, p) (ModuleBuilderT m)
  
  type Template p = TemplateT p Identity
  
  instance MFunctor (TemplateT p) where
    hoist nat = TemplateT . hoist (hoist nat) . unTemplateT
  
  instance MonadReader r m => MonadReader r (TemplateT p m) where
    ask = lift ask
    local f (TemplateT m) =
      TemplateT $ hoist (local f) m
  
  class HasSuffix m where
    getSuffix :: m Suffix
  
  instance Monad m => HasSuffix (TemplateT p m) where
    getSuffix = TemplateT $ asks (("_" <>) . fst)
  -- The following instance makes 'function' behave the same as in llvm-codegen
  instance Monad m => HasSuffix (ModuleBuilderT m) where
    getSuffix = pure mempty
  -- This allows getting the suffix inside a function body with llvm-codegen
  instance (Monad m, HasSuffix m) => HasSuffix (IRBuilderT m) where
    getSuffix = lift getSuffix
  -- MTL boilerplate:
  instance (Monad m, HasSuffix m) => HasSuffix (ReaderT r m) where
    getSuffix = lift getSuffix
  instance (Monad m, HasSuffix m, Monoid w) => HasSuffix (WriterT w m) where
    getSuffix = lift getSuffix
  instance (Monad m, HasSuffix m) => HasSuffix (LazyState.StateT w m) where
    getSuffix = lift getSuffix
  instance (Monad m, HasSuffix m) => HasSuffix (StrictState.StateT w m) where
    getSuffix = lift getSuffix
  instance (Monad m, HasSuffix m, Monoid w) => HasSuffix (LazyRWS.RWST r w s m) where
    getSuffix = lift getSuffix
  instance (Monad m, HasSuffix m, Monoid w) => HasSuffix (StrictRWS.RWST r w s m) where
    getSuffix = lift getSuffix
  instance (Monad m, HasSuffix m) => HasSuffix (ExceptT e m) where
    getSuffix = lift getSuffix
  instance (Monad m, HasSuffix m) => HasSuffix (ConfigT m) where
    getSuffix = lift getSuffix
  
  class MonadTemplate p m | m -> p where
    getParams :: m p
  
  instance Monad m => MonadTemplate p (TemplateT p m) where
    getParams = TemplateT $ asks snd
  
  instance (Monad m, MonadTemplate p m) => MonadTemplate p (ReaderT r m) where
    getParams = lift getParams
  instance (Monad m, MonadTemplate p m, Monoid w) => MonadTemplate p (WriterT w m) where
    getParams = lift getParams
  instance (Monad m, MonadTemplate p m, Monoid w) => MonadTemplate p (LazyState.StateT w m) where
    getParams = lift getParams
  instance (Monad m, MonadTemplate p m, Monoid w) => MonadTemplate p (StrictState.StateT w m) where
    getParams = lift getParams
  instance (Monad m, MonadTemplate p m, Monoid w) => MonadTemplate p (LazyRWS.RWST r w s m) where
    getParams = lift getParams
  instance (Monad m, MonadTemplate p m, Monoid w) => MonadTemplate p (StrictRWS.RWST r w s m) where
    getParams = lift getParams
  instance (Monad m, MonadTemplate p m) => MonadTemplate p (ExceptT e m) where
    getParams = lift getParams
  -- This allows getting the params inside a function body with llvm-codegen
  instance (Monad m, MonadTemplate p m) => MonadTemplate p (IRBuilderT m) where
    getParams = lift getParams
  instance (Monad m, MonadTemplate p m) => MonadTemplate p (ConfigT m) where
    getParams = lift getParams
  
  instance MonadTrans (TemplateT p) where
    lift m =
      TemplateT $ ReaderT $ const $ lift m
  
  -- "contramap" over the template params.
  -- Useful if you only need access to part of the template data
  -- and/or types of the params don't match.
  cmapParams :: (p2 -> p1) -> TemplateT p1 m a -> TemplateT p2 m a
  cmapParams f (TemplateT m) =
    TemplateT $ flip withReaderT m $ second f
  
  -- This instantiates a template, given a template name suffix and some template parameters.
  -- The result is the underlying ModuleBuilerT which generates specialized code based on the parameters.
  instantiate :: Suffix -> p -> TemplateT p m a -> ModuleBuilderT m a
  instantiate suffix p (TemplateT t) =
    runReaderT t (suffix, p)
  
  -- This instantiates a template and wraps it into another template.
  -- Useful for templated member functions on a templated object.
  partialInstantiate :: Monad m => p1 -> TemplateT p1 m a -> TemplateT p2 m a
  partialInstantiate p t = do
    suffix <- getSuffix
    embedIntoTemplate $ instantiate suffix p t
    where
      -- This embeds a plain ModuleBuilderT action into a template.
      -- This action has no access to the actual template params from this point onwards.
      embedIntoTemplate :: ModuleBuilderT m a -> TemplateT p m a
      embedIntoTemplate m = TemplateT $ ReaderT $ const m
  
  
  -- The next functions replace the corresponding functions defined in llvm-codegen.
  -- The functions automatically add a suffix if needed, to guarantee unique function names.
  -- In the actual codegen, you will still need to call 'getParams' to get access to the params,
  -- to do the actual specialization based on them.
  
  function :: (MonadModuleBuilder m, HasSuffix m)
           => Name -> [(Type, ParameterName)] -> Type -> ([Operand] -> IRBuilderT m a) -> m Operand
  function (unName -> name) args retTy body = do
    suffix <- getSuffix
    let nameWithSuffix = Name $ name <> suffix
    CG.function nameWithSuffix args retTy body
  
  typedef :: (MonadModuleBuilder m, HasSuffix m)
          => Name -> Flag Packed -> [Type] -> m Type
  typedef (unName -> name) packedFlag tys = do
    suffix <- getSuffix
    let nameWithSuffix = Name $ name <> suffix
    CG.typedef nameWithSuffix packedFlag tys
  