llvm-codegen-0.1.0.0: lib/LLVM/Codegen/IRBuilder.hs

{-# LANGUAGE RecursiveDo, PolyKinds, RoleAnnotations #-}

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

  , add
  , mul
  , sub
  , udiv
  , and
  , or
  , trunc
  , zext
  , ptrtoint
  , bitcast
  , ptrcast
  , icmp
  , alloca
  , gep
  , load
  , store
  , phi
  , call
  , ret
  , retVoid
  , br
  , condBr
  , switch
  , select

  , eq
  , ne
  , sge
  , sgt
  , sle
  , slt
  , uge
  , ugt
  , ule
  , ult
  , if'
  , loop
  , loopWhile
  , loopFor
  , pointerDiff
  , not'
  , Signedness(..)
  , minimum'
  , allocate
  , Path(..), (->>), mkPath
  , addr, deref, assign, update, increment, copy, swap

  , bit
  , int8
  , int16
  , int32
  , int64
  , intN
  , nullPtr
  ) where

import Prelude hiding (EQ, and, or)
import GHC.Stack
import Control.Monad.Fix
import qualified Data.List.NonEmpty as NE
import Data.List.NonEmpty (NonEmpty(..))
import Data.Word
import LLVM.Codegen.Name
import LLVM.Codegen.Operand
import LLVM.Codegen.Type
import LLVM.Codegen.IR
import LLVM.Codegen.IRBuilder.Monad
import LLVM.Codegen.ModuleBuilder


-- Helpers for generating instructions:

add :: (MonadIRBuilder m, HasCallStack) => Operand -> Operand -> m Operand
add lhs rhs =
  emitInstr (typeOf lhs) $ Add Off Off lhs rhs
{-# INLINEABLE add #-}

mul :: (MonadIRBuilder m, HasCallStack) => Operand -> Operand -> m Operand
mul lhs rhs =
  emitInstr (typeOf lhs) $ Mul Off Off lhs rhs
{-# INLINEABLE mul #-}

sub :: (MonadIRBuilder m, HasCallStack) => Operand -> Operand -> m Operand
sub lhs rhs =
  emitInstr (typeOf lhs) $ Sub Off Off lhs rhs
{-# INLINEABLE sub #-}

udiv :: (MonadIRBuilder m, HasCallStack) => Operand -> Operand -> m Operand
udiv lhs rhs =
  emitInstr (typeOf lhs) $ Udiv Off lhs rhs
{-# INLINEABLE udiv #-}

and :: (MonadIRBuilder m, HasCallStack) => Operand -> Operand -> m Operand
and lhs rhs =
  emitInstr (typeOf lhs) $ And lhs rhs
{-# INLINEABLE and #-}

or :: (MonadIRBuilder m, HasCallStack) => Operand -> Operand -> m Operand
or lhs rhs =
  emitInstr (typeOf lhs) $ Or lhs rhs
{-# INLINEABLE or #-}

trunc :: (MonadIRBuilder m, HasCallStack) => Operand -> Type -> m Operand
trunc val ty =
  emitInstr ty $ Trunc val ty
{-# INLINEABLE trunc #-}

zext :: (MonadIRBuilder m, HasCallStack) => Operand -> Type -> m Operand
zext val ty =
  emitInstr ty $ Zext val ty
{-# INLINEABLE zext #-}

ptrtoint :: (MonadIRBuilder m, HasCallStack) => Operand -> Type -> m Operand
ptrtoint val ty =
  emitInstr ty $ PtrToInt val ty
{-# INLINEABLE ptrtoint #-}

-- At the moment not useful because of introduction of opaque pointers.
-- Will become more useful once float or vector types are added again.
bitcast :: (MonadIRBuilder m, HasCallStack) => Operand -> Type -> m Operand
bitcast val ty =
  emitInstr ty $ Bitcast val ty
{-# INLINEABLE bitcast #-}

-- Casts a pointer to be a pointer containing type "ty".
-- This helper function is introduced to smooth the transition between
-- LLVM14 -> LLVM15+ (opaque pointer migration).
-- All bitcasts of pointers should be replaced with ptrcasts
ptrcast :: Type -> Operand -> Operand
ptrcast ty = \case
  LocalRef (PointerType _) name ->
    LocalRef (PointerType ty) name
  ConstantOperand (NullPtr _) ->
    ConstantOperand $ NullPtr ty
  _ ->
    error "'ptrcast' is only supported for pointer operands."
{-# INLINEABLE ptrcast #-}

icmp :: (MonadIRBuilder m, HasCallStack) => ComparisonType -> Operand -> Operand -> m Operand
icmp cmp a b =
  emitInstr i1 $ ICmp cmp a b
{-# INLINEABLE icmp #-}

eq, ne, sge, sgt, sle, slt, uge, ugt, ule, ult
  :: (MonadIRBuilder m, HasCallStack) => Operand -> Operand -> m Operand
eq = icmp EQ
ne = icmp NE
sge = icmp SGE
sgt = icmp SGT
sle = icmp SLE
slt = icmp SLT
uge = icmp UGE
ugt = icmp UGT
ule = icmp ULE
ult = icmp ULT
{-# INLINABLE eq #-}
{-# INLINABLE ne #-}
{-# INLINABLE sge #-}
{-# INLINABLE sgt #-}
{-# INLINABLE sle #-}
{-# INLINABLE slt #-}
{-# INLINABLE uge #-}
{-# INLINABLE ugt #-}
{-# INLINABLE ule #-}
{-# INLINABLE ult #-}

alloca :: (MonadIRBuilder m, HasCallStack) => Type -> Maybe Operand -> Int -> m Operand
alloca ty numElems alignment =
  emitInstr (ptr ty) $ Alloca ty numElems alignment
{-# INLINEABLE alloca #-}

gep :: (HasCallStack, MonadModuleBuilder m, MonadIRBuilder m)
    => Operand -> [Operand] -> m Operand
gep operand indices = do
  resultType <- computeGepType (typeOf operand) indices
  case resultType of
    Left err -> error err
    Right ty ->
      emitInstr ty $ GetElementPtr Off operand indices
{-# INLINEABLE gep #-}

computeGepType :: (MonadModuleBuilder m, HasCallStack) => Type -> [Operand] -> m (Either String Type)
computeGepType ty [] = pure $ Right $ PointerType ty
computeGepType (PointerType ty) (_ : idxs) =
  case (ty, null idxs) of
    -- If you want to load something from e.g. i8***, you need to gep + load for each pointer indirection!
    (PointerType{}, False) -> pure $ Left "Opaque pointers support only one gep offset."
    _ -> computeGepType ty idxs
computeGepType (StructureType _ elTys) (ConstantOperand (Int 32 val):is) =
  computeGepType (elTys !! fromIntegral val) is
computeGepType (StructureType _ _) (i:_) =
  pure $ Left $ "Indices into structures should be 32-bit integer constants. (Malformed AST): " <> show i
computeGepType (ArrayType _ elTy) (_:is) = computeGepType elTy is
computeGepType (NamedTypeReference n) is =
  lookupType n >>= \case
    Nothing -> pure $ Left $ "Couldn’t resolve typedef for: " <> show n
    Just ty -> computeGepType ty is
computeGepType ty _ =
  pure $ Left $ "Expecting aggregate type. (Malformed AST): " <> show ty
{-# INLINEABLE computeGepType #-}

load :: (HasCallStack, MonadIRBuilder m) => Operand -> Alignment -> m Operand
load address align =
  case typeOf address of
    PointerType ty ->
      emitInstr ty $ Load Off address Nothing align
    t ->
      error $ "Malformed AST: Expected a pointer type" <> show t
{-# INLINEABLE load #-}

store :: (MonadIRBuilder m, HasCallStack) => Operand -> Alignment -> Operand -> m ()
store address align value =
  emitInstrVoid $ Store Off address value Nothing align
{-# INLINEABLE store #-}

phi :: (HasCallStack, MonadIRBuilder m) => [(Operand, Name)] -> m Operand
phi cases
  | null cases = error "phi instruction should always have > 0 cases!"
  | otherwise =
    let neCases = NE.fromList cases
        ty = typeOf $ fst $ NE.head neCases
     in emitInstr ty $ Phi neCases
{-# INLINEABLE phi #-}

call :: (HasCallStack, MonadIRBuilder m) => Operand -> [Operand] -> m Operand
call fn args = case typeOf fn of
  FunctionType retTy _->
    emitCallInstr retTy
  PointerType (FunctionType retTy _) ->
    emitCallInstr retTy
  _ -> error "Malformed AST, expected function type in 'call' instruction"
  where
    emitCallInstr resultTy =
      if resultTy == VoidType
        then do
          emitInstrVoid $ Call Nothing C fn args
          pure $ ConstantOperand $ Undef void  -- Invalid, but isn't rendered anyway
        else emitInstr resultTy $ Call Nothing C fn args
{-# INLINEABLE call #-}

ret :: (HasCallStack, MonadIRBuilder m) => Operand -> m ()
ret val =
  emitTerminator (Terminator (Ret (Just val)))
{-# INLINEABLE ret #-}

retVoid :: (HasCallStack, MonadIRBuilder m) => m ()
retVoid =
  emitTerminator (Terminator (Ret Nothing))
{-# INLINEABLE retVoid #-}

br :: (MonadIRBuilder m, HasCallStack) => Name -> m ()
br label =
  emitTerminator (Terminator (Br label))
{-# INLINEABLE br #-}

condBr :: (HasCallStack, MonadIRBuilder m) => Operand -> Name -> Name -> m ()
condBr cond trueLabel falseLabel =
  emitTerminator (Terminator (CondBr cond trueLabel falseLabel))
{-# INLINEABLE condBr #-}

switch :: (HasCallStack, MonadIRBuilder m) => Operand -> Name -> [(Operand, Name)] -> m ()
switch value defaultDest dests =
  emitTerminator $ Terminator $ Switch value defaultDest dests
{-# INLINEABLE switch #-}

select :: (HasCallStack, MonadIRBuilder m) => Operand -> Operand -> Operand -> m Operand
select c t f =
  emitInstr (typeOf t) $ Select c t f
{-# INLINEABLE select #-}

if' :: (HasCallStack, MonadIRBuilder m, MonadFix m)
    => Operand -> m a -> m ()
if' condition asm = mdo
  condBr condition ifBlock end
  ifBlock <- blockNamed "if"
  _ <- asm
  br end
  end <- blockNamed "end_if"
  pure ()
{-# INLINEABLE if' #-}

loop :: (HasCallStack, MonadIRBuilder m, MonadFix m) => m a -> m ()
loop asm = mdo
  br begin
  begin <- blockNamed "loop"
  _ <- asm
  br begin
{-# INLINEABLE loop #-}

loopWhile :: (HasCallStack, MonadIRBuilder m, MonadFix m)
          => m Operand -> m a -> m ()
loopWhile condition asm = mdo
  br begin
  begin <- blockNamed "while_begin"
  result <- condition
  condBr result body end
  body <- blockNamed "while_body"
  _ <- asm
  br begin
  end <- blockNamed "while_end"
  pure ()
{-# INLINEABLE loopWhile #-}

loopFor :: (HasCallStack, MonadModuleBuilder m, MonadIRBuilder m, MonadFix m)
        => Operand
        -> (Operand -> m Operand)
        -> (Operand -> m Operand)
        -> (Operand -> m a)
        -> m ()
loopFor beginValue condition post asm = mdo
  start <- currentBlock
  br begin
  begin <- blockNamed "for_begin"
  loopValue <- phi [(beginValue, start), (updatedValue, bodyEnd)]
  result <- condition loopValue
  condBr result bodyStart end
  bodyStart <- blockNamed "for_body"
  _ <- asm loopValue
  updatedValue <- post loopValue
  bodyEnd <- currentBlock
  br begin
  end <- blockNamed "for_end"
  pure ()
{-# INLINEABLE loopFor #-}

-- NOTE: diff is in bytes! (Different compared to C and C++)
pointerDiff :: (HasCallStack, MonadIRBuilder m)
            => Type -> Operand -> Operand -> m Operand
pointerDiff ty a b = do
  a' <- ptrtoint a i64
  b' <- ptrtoint b i64
  result <- sub a' b'
  if ty == i64
    then pure result
    else trunc result ty
{-# INLINEABLE pointerDiff #-}

-- | Calculates the logical not of a boolean 'Operand'.
--   NOTE: This assumes the 'Operand' is of type 'i1', this is not checked!
--   Passing in an argument of another width will lead to a crash in LLVM.
not' :: (HasCallStack, MonadIRBuilder m)
     => Operand -> m Operand
not' bool =
  select bool (bit 0) (bit 1)
{-# INLINEABLE not' #-}

data Signedness = Signed | Unsigned

--   NOTE: No check is made if the 2 operands have the same 'Type'!
minimum' :: (HasCallStack, MonadIRBuilder m)
         => Signedness -> Operand -> Operand -> m Operand
minimum' sign a b = do
  let inst = case sign of
        Signed -> slt
        Unsigned -> ult
  isLessThan <- inst a b
  select isLessThan a b
{-# INLINEABLE minimum' #-}

allocate :: (HasCallStack, MonadIRBuilder m) => Type -> Operand -> m Operand
allocate ty beginValue = do
  value <- alloca ty Nothing 0
  store value 0 beginValue
  pure value
{-# INLINEABLE allocate #-}

newtype Path (a :: k) (b :: k)
  = Path (NonEmpty Operand)
  deriving (Eq, Show)
type role Path nominal nominal

mkPath :: [Operand] -> Path a b
mkPath path = Path (int32 0 :| path)

(->>) :: Path a b -> Path b c -> Path a c
Path a2b ->> Path b2c =
  let b2c' = if NE.head b2c == int32 0
               then NE.tail b2c
               else NE.toList b2c
   in Path $ NE.head a2b :| (NE.tail a2b ++ b2c')

addr :: (MonadModuleBuilder m, MonadIRBuilder m, HasCallStack)
     => Path a b -> Operand -> m Operand
addr path p = gep p (pathToIndices path)
  where
    pathToIndices :: Path a b -> [Operand]
    pathToIndices (Path indices) =
      NE.toList indices
{-# INLINEABLE addr #-}

deref :: (MonadModuleBuilder m, MonadIRBuilder m, HasCallStack)
      => Path a b -> Operand -> m Operand
deref path p = do
  address <- addr path p
  load address 0
{-# INLINEABLE deref #-}

assign :: (MonadModuleBuilder m, MonadIRBuilder m, HasCallStack)
       => Path a b -> Operand -> Operand -> m ()
assign path p value = do
  dstAddr <- addr path p
  store dstAddr 0 value
{-# INLINEABLE assign #-}

update :: (MonadModuleBuilder m, MonadIRBuilder m, HasCallStack)
       => Path a b
       -> Operand
       -> (Operand -> m Operand)
       -> m ()
update path p f = do
  dstAddr <- addr path p
  store dstAddr 0 =<< f =<< load dstAddr 0
{-# INLINEABLE update #-}

increment :: (MonadModuleBuilder m, MonadIRBuilder m, HasCallStack)
          => (Integer -> Operand) -> Path a b -> Operand -> m ()
increment ty path p =
  update path p (add (ty 1))
{-# INLINEABLE increment #-}

copy :: (MonadModuleBuilder m, MonadIRBuilder m, HasCallStack)
     => Path a b -> Operand -> Operand -> m ()
copy path src dst = do
  value <- deref path src
  assign path dst value
{-# INLINEABLE copy #-}

swap :: (MonadModuleBuilder m, MonadIRBuilder m, HasCallStack)
     => Path a b -> Operand -> Operand -> m ()
swap path lhs rhs = do
  tmp <- deref path lhs
  copy path rhs lhs
  assign path rhs tmp
{-# INLINEABLE swap #-}


bit :: Integer -> Operand
bit b =
  intN 1 $ if b == 0 then 0 else 1

int8 :: Integer -> Operand
int8 =
  intN 8

int16 :: Integer -> Operand
int16 =
  intN 16

int32 :: Integer -> Operand
int32 =
  intN 32

int64 :: Integer -> Operand
int64 =
  intN 64

intN :: Word32 -> Integer -> Operand
intN bits value =
  ConstantOperand $ Int bits value

nullPtr :: Type -> Operand
nullPtr =
  ConstantOperand . NullPtr