Monad.hs

{-# LANGUAGE TypeFamilies, RankNTypes, MultiParamTypeClasses, UndecidableInstances, BangPatterns, TypeOperators #-}

module LLVM.Codegen.IRBuilder.Monad
  ( IRBuilderT
  , IRBuilder
  , runIRBuilderT
  , runIRBuilder
  , MonadIRBuilder(..)
  , BasicBlock(..)
  , block
  , blockNamed
  , emitBlockStart
  , emitInstr
  , emitInstrVoid
  , emitTerminator
  , renderBasicBlock
  , freshName
  ) where

-- NOTE: this module only exists to solve a cyclic import

import Prelude hiding (and)
import Control.Arrow hiding ((<+>))
import Control.Monad.State.Lazy (StateT(..), MonadState, modify)
import qualified Control.Monad.State.Strict as StrictState
import qualified Control.Monad.State.Lazy as LazyState
import qualified Control.Monad.RWS.Lazy as LazyRWS
import qualified Control.Monad.RWS.Strict as StrictRWS
import Control.Monad
import Control.Monad.Reader
import Control.Monad.Writer
import Control.Monad.Except
import Control.Monad.Morph
import Control.Monad.Fix
import Data.Functor.Identity
import qualified Data.Text as T
import qualified Data.DList as DList
import Data.DList (DList)
import qualified Data.Map as M
import Data.Map (Map)
import Data.Maybe
import Data.Monoid
import LLVM.Codegen.Operand
import LLVM.Codegen.IR
import LLVM.Codegen.Type
import LLVM.Codegen.Name
import LLVM.Pretty


data BasicBlock
  = BB
  { bbName :: !Name
  , bbInstructions :: !(DList (Maybe Operand, IR))
  , bbTerminator :: !Terminator
  } deriving Show

data PartialBlock
  = PartialBlock
  { pbName :: !Name
  , pbInstructions :: !(DList (Maybe Operand, IR))
  , pbTerminator :: !(First Terminator)
  , pbNumInstrs :: !Int
  }

data IRBuilderState
  = IRBuilderState
  { allocas :: !(DList (Maybe Operand, IR))
  , basicBlocks :: !(DList BasicBlock)
  , currentPartialBlock :: !PartialBlock
  , operandCounter :: !Int
  , nameMap :: !(Map T.Text Int)
  }

newtype IRBuilderT m a
  = IRBuilderT { unIRBuilderT :: StateT IRBuilderState m a }
  deriving ( Functor, Applicative, Monad, MonadFix, MonadIO
           , MonadError e
           )
  via StateT IRBuilderState m

instance MonadReader r m => MonadReader r (IRBuilderT m) where
  ask = lift ask
  {-# INLINEABLE ask #-}
  local = mapIRBuilderT . local
  {-# INLINEABLE local #-}

-- TODO MonadWriter

mapIRBuilderT :: (Monad m, Monad n) => (m a -> n a) -> IRBuilderT m a -> IRBuilderT n a
mapIRBuilderT f (IRBuilderT inner) =
  IRBuilderT $ do
    s <- LazyState.get
    LazyState.mapStateT (g s) inner
  where
    g s = fmap (,s) . f . fmap fst
{-# INLINEABLE mapIRBuilderT #-}

instance MonadState s m => MonadState s (IRBuilderT m) where
  state = lift . StrictState.state
  {-# INLINEABLE state #-}

instance MonadTrans IRBuilderT where
  lift = IRBuilderT . lift
  {-# INLINEABLE lift #-}

instance MFunctor IRBuilderT where
  hoist nat = IRBuilderT . hoist nat . unIRBuilderT
  {-# INLINEABLE hoist #-}

type IRBuilder = IRBuilderT Identity

runIRBuilderT :: Monad m => IRBuilderT m a -> m (a, [BasicBlock])
runIRBuilderT (IRBuilderT m) = do
  let partialBlock = PartialBlock (Name "start") mempty mempty 0
      result = runStateT m (IRBuilderState mempty mempty partialBlock 0 mempty)
  fmap (second getBlocks) result
  where
    getBlocks irState =
      case blocks of
        [] -> []
        (firstBlk:restBlks) ->
          let firstBlk' = firstBlk { bbInstructions = DList.append allocations (bbInstructions firstBlk) }
           in (firstBlk':restBlks)
      where
        previousBlocks = DList.apply (basicBlocks irState) mempty
        currentBlk = currentPartialBlock irState
        blocks = previousBlocks <> [partialBlockToBasicBlock currentBlk]
        allocations = allocas irState
{-# INLINEABLE runIRBuilderT #-}

runIRBuilder :: IRBuilder a -> (a, [BasicBlock])
runIRBuilder = runIdentity . runIRBuilderT
{-# INLINEABLE runIRBuilder #-}

partialBlockToBasicBlock :: PartialBlock -> BasicBlock
partialBlockToBasicBlock pb =
  let currentTerm = fromMaybe (Terminator $ Ret Nothing) $ getFirst $ pbTerminator pb
  in BB (pbName pb) (pbInstructions pb) currentTerm
{-# INLINEABLE partialBlockToBasicBlock #-}

block :: (MonadIRBuilder m) => m Name
block = do
  blockName <- freshName (Just "block")
  emitBlockStart blockName
  pure blockName
{-# INLINEABLE block #-}

blockNamed :: (MonadIRBuilder m) => T.Text -> m Name
blockNamed blkName = do
  blockName <- freshName (Just blkName)
  emitBlockStart blockName
  pure blockName
{-# INLINEABLE blockNamed #-}

emitBlockStart :: (MonadIRBuilder m) => Name -> m ()
emitBlockStart blockName =
  modifyIRBuilderState $ \s ->
    let currBlock = currentPartialBlock s
        hasntStartedBlock = (pbNumInstrs currBlock == 0) && isNothing (getFirst (pbTerminator currBlock))
        blocks = basicBlocks s
        -- If the current block is empty:
        --   Insert a dummy basic block that jumps directly to the next block, to avoid continuity errors.
        --   Normally, LLVM should optimize this away since it is semantically a no-op.
        -- Otherwise:
        --   Append the current block to the existing list of blocks.
        --
        -- NOTE: This is different behavior compared to the llvm-hs-pure library,
        -- but this avoids a lot of partial functions!
        newBlock =
          if hasntStartedBlock
            then BB (pbName currBlock) mempty (Terminator $ Br blockName)
            else partialBlockToBasicBlock currBlock
     in s { basicBlocks = DList.snoc blocks newBlock
          , currentPartialBlock = PartialBlock blockName mempty mempty 0
          }
{-# INLINEABLE emitBlockStart #-}

-- NOTE: Only used internally, this creates an unassigned operand
mkOperand :: (MonadIRBuilder m) => Type -> m Operand
mkOperand ty = LocalRef ty <$!> freshUnnamed
{-# INLINEABLE mkOperand #-}

freshName :: MonadIRBuilder m => Maybe T.Text -> m Name
freshName = \case
  Nothing -> freshUnnamed
  Just suggestion -> do
    nameMapping <- nameMap <$> getIRBuilderState
    let !mCount = M.lookup suggestion nameMapping
        !count = fromMaybe 0 mCount
        !newMapping = M.insert suggestion (count + 1) nameMapping
    modifyIRBuilderState $ \s -> s { nameMap = newMapping }
    pure $! Name $! suggestion <> "_" <> T.pack (show count)
{-# INLINEABLE freshName #-}

freshUnnamed :: MonadIRBuilder m => m Name
freshUnnamed = do
  !ctr <- operandCounter <$> getIRBuilderState
  let !newCount = ctr + 1
  modifyIRBuilderState $ \s -> s { operandCounter = newCount }
  pure $! Generated ctr
{-# INLINEABLE freshUnnamed #-}

emitInstr :: (MonadIRBuilder m) => Type -> IR -> m Operand
emitInstr ty = \case
  instr@(Alloca {}) -> do
    -- For performant code, all alloca instructions should be at the start of the function!
    -- https://llvm.org/docs/Frontend/PerformanceTips.html#use-of-allocas
    -- (A custom operand name is only used here to avoid having to re-number all operands.)
    operand <- LocalRef ty <$!> freshName (Just "stack.ptr")
    addAlloca operand instr
    pure operand
  instr -> do
    operand <- mkOperand ty
    addInstrToCurrentBlock (Just operand) instr
    pure operand
{-# INLINABLE emitInstr #-}

emitInstrVoid :: MonadIRBuilder m => IR -> m ()
emitInstrVoid =
  addInstrToCurrentBlock Nothing
{-# INLINABLE emitInstrVoid #-}

addInstrToCurrentBlock :: MonadIRBuilder m => Maybe Operand -> IR -> m ()
addInstrToCurrentBlock operand instr =
  modifyCurrentBlock $ \blk ->
    let instrs = DList.snoc (pbInstructions blk) (operand, instr)
        in blk { pbInstructions = instrs, pbNumInstrs = pbNumInstrs blk + 1 }
{-# INLINEABLE addInstrToCurrentBlock #-}

addAlloca :: MonadIRBuilder m => Operand -> IR -> m ()
addAlloca operand instr =
  modifyIRBuilderState $ \s ->
    s { allocas = DList.snoc (allocas s) (Just operand, instr) }
{-# INLINEABLE addAlloca #-}

emitTerminator :: MonadIRBuilder m => Terminator -> m ()
emitTerminator term =
  modifyCurrentBlock $ \blk ->
    blk { pbTerminator = pbTerminator blk <> First (Just term) }
{-# INLINABLE emitTerminator #-}

modifyCurrentBlock :: MonadIRBuilder m => (PartialBlock -> PartialBlock) -> m ()
modifyCurrentBlock f =
  modifyIRBuilderState $ \s ->
    s { currentPartialBlock = f (currentPartialBlock s) }
{-# INLINEABLE modifyCurrentBlock #-}

class Monad m => MonadIRBuilder m where
  getIRBuilderState :: m IRBuilderState

  modifyIRBuilderState :: (IRBuilderState -> IRBuilderState) -> m ()

  currentBlock :: m Name

  default getIRBuilderState
    :: (MonadTrans t, MonadIRBuilder m1, m ~ t m1)
    => m IRBuilderState
  getIRBuilderState = lift getIRBuilderState
  {-# INLINEABLE getIRBuilderState #-}

  default modifyIRBuilderState
    :: (MonadTrans t, MonadIRBuilder m1, m ~ t m1)
    => (IRBuilderState -> IRBuilderState)
    -> m ()
  modifyIRBuilderState = lift . modifyIRBuilderState
  {-# INLINEABLE modifyIRBuilderState #-}

  default currentBlock
    :: (MonadTrans t, MonadIRBuilder m1, m ~ t m1)
    => m Name
  currentBlock = lift currentBlock
  {-# INLINEABLE currentBlock #-}

instance Monad m => MonadIRBuilder (IRBuilderT m) where
  getIRBuilderState = IRBuilderT LazyState.get
  {-# INLINEABLE getIRBuilderState #-}

  modifyIRBuilderState = IRBuilderT . modify
  {-# INLINEABLE modifyIRBuilderState #-}

  currentBlock =
    IRBuilderT $ LazyState.gets (pbName . currentPartialBlock)
  {-# INLINEABLE currentBlock #-}

instance MonadIRBuilder m => MonadIRBuilder (StrictState.StateT s m)
instance MonadIRBuilder m => MonadIRBuilder (LazyState.StateT s m)
instance (MonadIRBuilder m, Monoid w) => MonadIRBuilder (StrictRWS.RWST r w s m)
instance (MonadIRBuilder m, Monoid w) => MonadIRBuilder (LazyRWS.RWST r w s m)
instance MonadIRBuilder m => MonadIRBuilder (ReaderT r m)
instance (MonadIRBuilder m, Monoid w) => MonadIRBuilder (WriterT w m)
instance MonadIRBuilder m => MonadIRBuilder (ExceptT e m)

renderBasicBlock :: Renderer BasicBlock
renderBasicBlock buf (BB name stmts (Terminator term)) =
  if null stmts
    then (renderName buf name |># ":\n  "#) `renderIR` term
    else (vsep (renderName buf name |># ":\n"#) stmts' renderStmt |># "\n  "#) `renderIR` term
  where
    stmts' = DList.apply stmts []
    renderStmt :: Buffer %1 -> (Maybe Operand, IR) -> Buffer
    renderStmt buf' (mOperand, instr) =
      withIndent buf' (\buf'' -> renderStmt' buf'' mOperand instr)
    renderStmt' :: Buffer %1 -> Maybe Operand -> IR -> Buffer
    renderStmt' buf' mOperand instr =
      renderMaybe buf' mOperand (\buf'' operand -> buf'' `renderOperand` operand |># " = "#) `renderIR` instr
{-# INLINABLE renderBasicBlock #-}