Codegen.hs

{-# LANGUAGE TupleSections #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RecursiveDo #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Microc.Codegen
  ( codegenProgram
  )
where

import qualified LLVM.AST.IntegerPredicate     as IP
import qualified LLVM.AST.FloatingPointPredicate
                                               as FP
import           LLVM.AST                       ( Operand )
import qualified LLVM.AST                      as AST
import qualified LLVM.AST.Float                as AST
import qualified LLVM.AST.Type                 as AST
import qualified LLVM.AST.Constant             as C
import           LLVM.AST.Name
import           LLVM.AST.Typed                 ( typeOf )

import qualified LLVM.IRBuilder.Module         as L
import qualified LLVM.IRBuilder.Monad          as L
import qualified LLVM.IRBuilder.Instruction    as L
import qualified LLVM.IRBuilder.Constant       as L
import           LLVM.Prelude                   ( ShortByteString )

import qualified Data.Map                      as M
import           Control.Monad.State
import           Data.String                    ( fromString )

import           Microc.Utils
import           Microc.Sast
import           Microc.Ast                     ( Type(..)
                                                , Op(..)
                                                , Uop(..)
                                                , Bind(..)
                                                , Struct(..)
                                                )

import           Data.String.Conversions
import qualified Data.Text                     as T
import           Data.Text                      ( Text )
import           Data.Word                      ( Word32 )
import           Data.List                      ( find )

-- When using the IRBuilder, both functions and variables have the type Operand
data Env = Env { operands :: M.Map Text Operand
               , structs :: [ Struct ]
               , strings :: M.Map Text Operand
               }
  deriving (Eq, Show)

-- LLVM and Codegen type synonyms allow us to emit module definitions and basic
-- block instructions at the top level without being forced to pass explicit
-- module and builder parameters to every function
type LLVM = L.ModuleBuilderT (State Env)
type Codegen = L.IRBuilderT LLVM

registerOperand :: MonadState Env m => Text -> Operand -> m ()
registerOperand name op =
  modify $ \env -> env { operands = M.insert name op (operands env) }

getFields :: MonadState Env m => Text -> m [Bind]
getFields name = do
  ss <- gets structs
  case find (\s -> structName s == name) ss of
    Nothing               -> error "Internal error - struct not found"
    Just (Struct _ binds) -> pure binds

instance ConvertibleStrings Text ShortByteString where
  convertString = fromString . T.unpack

ltypeOfTyp :: MonadState Env m => Type -> m AST.Type
ltypeOfTyp = \case
  TyVoid         -> pure AST.void
  TyInt          -> pure AST.i32
  TyChar         -> pure AST.i8
  TyFloat        -> pure AST.double
  TyBool         -> pure AST.i1
  -- (void *) is invalid LLVM
  Pointer TyVoid -> pure $ charStar
  -- special case to handle recursively defined structures
  -- TODO: add real cycle checking so that improperly defined
  -- recursive types case the compiler to hang forever
  Pointer (TyStruct n) ->
    pure $ AST.ptr (AST.NamedTypeReference (mkName $ cs ("struct." <> n)))
  Pointer  t -> fmap AST.ptr (ltypeOfTyp t)
  TyStruct n -> do
    fields <- getFields n
    typs   <- mapM (ltypeOfTyp . bindType) fields
    -- Packed structs aren't great for performance but very easy to code for now
    pure $ AST.StructureType { AST.isPacked = True, AST.elementTypes = typs }

charStar :: AST.Type
charStar = AST.ptr AST.i8

sizeof :: MonadState Env m => Type -> m Word32
sizeof = \case
  TyBool     -> pure 1
  TyChar     -> pure 1
  TyInt      -> pure 4
  TyFloat    -> pure 8
  TyVoid     -> pure 0
  Pointer  _ -> pure 8
  TyStruct n -> do
    fields <- getFields n
    sizes  <- mapM (sizeof . bindType) fields
    pure (sum sizes)

codegenLVal :: LValue -> Codegen Operand
codegenLVal (SId    name) = gets ((M.! name) . operands)
codegenLVal (SDeref e   ) = codegenSexpr e

codegenLVal (SAccess e i) = do
  e' <- codegenLVal e
  L.gep e' [L.int32 0, L.int32 (fromIntegral i)]

codegenSexpr :: SExpr -> Codegen Operand
codegenSexpr (TyInt         , SLiteral i ) = pure $ L.int32 (fromIntegral i)
codegenSexpr (TyFloat       , SFliteral f) = pure $ L.double f
codegenSexpr (TyBool        , SBoolLit b ) = pure $ L.bit (if b then 1 else 0)
codegenSexpr (TyChar        , SCharLit c ) = pure $ L.int8 (fromIntegral c)
codegenSexpr (Pointer TyChar, SStrLit s  ) = do
  -- Generate a new unique global variable for every string literal we see
  strs <- gets strings
  case M.lookup s strs of
    Nothing -> do
      let nm = mkName (show (M.size strs) <> ".str")
      op <- L.globalStringPtr (cs s) nm
      modify $ \env -> env { strings = M.insert s (AST.ConstantOperand op) strs }
      pure (AST.ConstantOperand op)
    Just op -> pure op

codegenSexpr (t    , SNull          ) = L.inttoptr (L.int64 0) =<< ltypeOfTyp t
codegenSexpr (TyInt, SSizeof t      ) = L.int32 . fromIntegral <$> sizeof t

-- All LVals are already memory addresses.
codegenSexpr (_    , SAddr e        ) = codegenLVal e
codegenSexpr (_    , LVal e         ) = flip L.load 0 =<< codegenLVal e
codegenSexpr (_    , SAssign lhs rhs) = do
  rhs' <- codegenSexpr rhs
  lhs' <- codegenLVal lhs
  L.store lhs' 0 rhs'
  return rhs'
codegenSexpr (t, SBinop op lhs rhs) = do
  lhs' <- codegenSexpr lhs
  rhs' <- codegenSexpr rhs
  case op of
    Add -> case (fst lhs, fst rhs) of
      (Pointer _, TyInt    ) -> L.gep lhs' [rhs']
      (TyInt    , Pointer _) -> L.gep rhs' [lhs']
      (TyInt    , TyInt    ) -> L.add lhs' rhs'
      (TyFloat  , TyFloat  ) -> L.fadd lhs' rhs'
      _                      -> error "Internal error - semant failed"
    Sub -> case (fst lhs, fst rhs) of
      (Pointer typ, Pointer typ') -> if typ' /= typ
        then error "Internal error - semant failed"
        else do
          lhs''  <- L.ptrtoint lhs' AST.i64
          rhs''  <- L.ptrtoint rhs' AST.i64
          diff   <- L.sub lhs'' rhs''
          width  <- L.int64 . fromIntegral <$> sizeof typ
          result <- L.sdiv diff width
          L.trunc result AST.i32
      (Pointer _, TyInt) -> do
        rhs'' <- L.sub (L.int32 0) rhs'
        L.gep lhs' [rhs'']
      (TyInt  , TyInt  ) -> L.sub lhs' rhs'
      (TyFloat, TyFloat) -> L.fsub lhs' rhs'
      _                  -> error "Internal error - semant failed"
    Mult -> case t of
      TyInt   -> L.mul lhs' rhs'
      TyFloat -> L.fmul lhs' rhs'
      _       -> error "Internal error - semant failed"
    Div -> case t of
      TyInt   -> L.sdiv lhs' rhs'
      TyFloat -> L.fdiv lhs' rhs'
      _       -> error "Internal error - semant failed"
    -- We implement int ** int directly in llvm
    Power -> mdo
      enclosing <- L.currentBlock
      L.br loop
      loop <- L.block `L.named` "loop_pow"
      acc <- L.phi [(L.int32 1, enclosing), (nextAcc, continueBlock)] `L.named` "acc"
      expt <- L.phi [(rhs', enclosing), (nextExpt, continueBlock)] `L.named` "expt"
      done <- L.icmp IP.EQ expt (L.int32 0)
      L.condBr done doneBlock continueBlock
      continueBlock <- L.block `L.named` "continue"
      nextAcc       <- L.mul acc lhs' `L.named` "next_acc"
      nextExpt      <- L.sub expt (L.int32 1) `L.named` "next_expt"
      L.br loop
      doneBlock <- L.block `L.named` "done"
      pure acc

    -- Only relational operators defined on chars, not arithmetic
    Equal -> case fst lhs of
      TyInt     -> L.icmp IP.EQ lhs' rhs'
      TyBool    -> L.icmp IP.EQ lhs' rhs'
      TyChar    -> L.icmp IP.EQ lhs' rhs'
      Pointer _ -> L.icmp IP.EQ lhs' rhs'
      TyFloat   -> L.fcmp FP.OEQ lhs' rhs'
      _         -> error "Internal error - semant failed"
    Neq -> case fst lhs of
      TyInt     -> L.icmp IP.NE lhs' rhs'
      TyBool    -> L.icmp IP.NE lhs' rhs'
      TyChar    -> L.icmp IP.NE lhs' rhs'
      Pointer _ -> L.icmp IP.NE lhs' rhs'
      TyFloat   -> L.fcmp FP.ONE lhs' rhs'
      _         -> error "Internal error - semant failed"
    Less -> case fst lhs of
      TyInt   -> L.icmp IP.SLT lhs' rhs'
      TyBool  -> L.icmp IP.SLT lhs' rhs'
      TyChar  -> L.icmp IP.ULT lhs' rhs'
      TyFloat -> L.fcmp FP.OLT lhs' rhs'
      _       -> error "Internal error - semant failed"
    Leq -> case fst lhs of
      TyInt   -> L.icmp IP.SLE lhs' rhs'
      TyBool  -> L.icmp IP.SLE lhs' rhs'
      TyChar  -> L.icmp IP.ULE lhs' rhs'
      TyFloat -> L.fcmp FP.OLE lhs' rhs'
      _       -> error "Internal error - semant failed"
    Greater -> case fst lhs of
      TyInt   -> L.icmp IP.SGT lhs' rhs'
      TyBool  -> L.icmp IP.SGT lhs' rhs'
      TyChar  -> L.icmp IP.UGT lhs' rhs'
      TyFloat -> L.fcmp FP.OGT lhs' rhs'
      _       -> error "Internal error - semant failed"
    Geq -> case fst lhs of
      TyInt   -> L.icmp IP.SGE lhs' rhs'
      TyBool  -> L.icmp IP.SGE lhs' rhs'
      TyChar  -> L.icmp IP.UGE lhs' rhs'
      TyFloat -> L.fcmp FP.OGE lhs' rhs'
      _       -> error "Internal error - semant failed"
    -- Relational operators all emit the same instructions
    -- Calling And between floats or pointers doesn't make sense,
    -- but semant catches that. Calling BitAnd between them IS allowed,
    -- and makes even less sense, but this is C, and we "trust" programmers
    -- to know why they want to do such things.
    And    -> L.and lhs' rhs'
    Or     -> L.or lhs' rhs'
    BitAnd -> L.and lhs' rhs'
    BitOr  -> L.or lhs' rhs'

codegenSexpr (t, SUnop op e) = do
  e' <- codegenSexpr e
  case op of
    Neg -> case t of
      TyInt   -> L.sub (L.int32 0) e'
      TyFloat -> L.fsub (L.double 0) e'
      _       -> error "Internal error - semant failed"
    Not -> case t of
      TyBool -> L.xor e' (L.bit 1)
      _      -> error "Internal error - semant failed"

codegenSexpr (_, SCall fun es) = do
  es' <- mapM (fmap (, []) . codegenSexpr) es
  f   <- gets ((M.! fun) . operands)
  L.call f es'

codegenSexpr (_, SCast t' e@(t, _)) = do
  e'       <- codegenSexpr e
  llvmType <- ltypeOfTyp t'
  case (t', t) of
    (Pointer _, Pointer _) -> L.bitcast e' llvmType
    (Pointer _, TyInt    ) -> L.inttoptr e' llvmType
    (TyInt    , Pointer _) -> L.ptrtoint e' llvmType
    (TyFloat  , TyInt    ) -> L.sitofp e' llvmType
    _ -> error $ "Internal error - semant failed. Invalid sexpr " <> show
      (t', SCast t e)

codegenSexpr (_, SNoexpr) = pure $ L.int64 0

-- Final catchall
codegenSexpr sx =
  error $ "Internal error - semant failed. Invalid sexpr " <> show sx


codegenStatement :: SStatement -> Codegen ()
codegenStatement (SExpr   e) = void $ codegenSexpr e

codegenStatement (SReturn e) = case e of
  (TyVoid, SNoexpr) -> L.retVoid
  _                 -> L.ret =<< codegenSexpr e

codegenStatement (SBlock ss     ) = mapM_ codegenStatement ss

codegenStatement (SIf p cons alt) = mdo
  bool <- codegenSexpr p
  L.condBr bool thenBlock elseBlock
  thenBlock <- L.block `L.named` "then"
  do
    codegenStatement cons
    mkTerminator $ L.br mergeBlock
  elseBlock <- L.block `L.named` "else"
  do
    codegenStatement alt
    mkTerminator $ L.br mergeBlock
  mergeBlock <- L.block `L.named` "merge"
  return ()

codegenStatement (SDoWhile p body) = mdo
  L.br whileBlock
  whileBlock <- L.block `L.named` "while_body"
  do
    codegenStatement body
    continue <- codegenSexpr p
    mkTerminator $ L.condBr continue whileBlock mergeBlock
  mergeBlock <- L.block `L.named` "merge"
  return ()

mkTerminator :: Codegen () -> Codegen ()
mkTerminator instr = do
  check <- L.hasTerminator
  unless check instr

-- | Generate a function and add both the function name and variable names to
-- the map.
codegenFunc :: SFunction -> LLVM ()
codegenFunc f = mdo
  -- We need to forward reference the generated function and insert it into the
  -- environment _before_ generating its body in order to handle the
  -- possibility of the function calling itself recursively
  registerOperand (sname f) function
  -- We wrap generating the function inside of the `locally` combinator in
  -- order to prevent local variables from escaping the scope of the function
  (function, strs) <- locally $ do
    retty    <- ltypeOfTyp (styp f)
    args     <- mapM mkParam (sformals f)
    fun      <- L.function name args retty genBody
    strings' <- gets strings
    pure (fun, strings')
  modify $ \e -> e { strings = strs }
 where
  name = mkName (cs $ sname f)
  mkParam (Bind t n) = (,) <$> ltypeOfTyp t <*> pure (L.ParameterName (cs n))

  -- Generate the body of the function:
  genBody :: [Operand] -> Codegen ()
  genBody ops = do
    _entry <- L.block `L.named` "entry"
    -- Add the formal parameters to the map, allocate them on the stack,
    -- and then emit the necessary store instructions
    forM_ (zip ops (sformals f)) $ \(op, Bind _ n) -> do
      addr <- L.alloca (typeOf op) Nothing 0
      L.store addr 0 op
      registerOperand n addr
    -- Same for the locals, except we do not emit the store instruction for
    -- them
    forM_ (slocals f) $ \(Bind t n) -> do
      ltype <- ltypeOfTyp t
      addr  <- L.alloca ltype Nothing 0
      registerOperand n addr
    -- Evaluate the actual body of the function after making the necessary
    -- allocations
    codegenStatement (sbody f)

emitBuiltIn :: (String, [AST.Type], AST.Type) -> LLVM ()
emitBuiltIn (name, argtys, retty) = do
  func <- L.extern (mkName name) argtys retty
  registerOperand (cs name) func

builtIns :: [(String, [AST.Type], AST.Type)]
builtIns =
  [ ("printbig"     , [AST.i32]               , AST.void)
  , ("llvm.pow.f64" , [AST.double, AST.double], AST.double)
  , ("llvm.powi.f64", [AST.double, AST.i32]   , AST.double)
  , ("malloc"       , [AST.i32]               , AST.ptr AST.i8)
  , ("free"         , [AST.ptr AST.i8]        , AST.void)
  ]

codegenGlobal :: Bind -> LLVM ()
codegenGlobal (Bind t n) = do
  typ <- ltypeOfTyp t
  let name    = mkName $ cs n
      initVal = case t of
        Pointer  _ -> C.Int 64 0
        TyStruct _ -> C.AggregateZero typ
        TyInt      -> C.Int 32 0
        TyBool     -> C.Int 1 0
        TyFloat    -> C.Float (AST.Double 0)
        TyChar     -> C.Int 8 0
        TyVoid     -> error "Global void variables illegal"
  var <- L.global name typ initVal
  registerOperand n var

emitTypeDef :: Struct -> LLVM AST.Type
emitTypeDef (Struct name _) = do
  typ <- ltypeOfTyp (TyStruct name)
  L.typedef (mkName (cs ("struct." <> name))) (Just typ)


codegenProgram :: SProgram -> AST.Module
codegenProgram (structs, globals, funcs) =
  flip evalState (Env { operands = M.empty, structs, strings = M.empty })
    $ L.buildModuleT "microc"
    $ do
        printf <- L.externVarArgs (mkName "printf") [charStar] AST.i32
        registerOperand "printf" printf
        mapM_ emitBuiltIn   builtIns
        mapM_ emitTypeDef   structs
        mapM_ codegenGlobal globals
        mapM_ codegenFunc   funcs