content: {-# LANGUAGE GADTs                      #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE OverloadedStrings          #-}
  {-# LANGUAGE RecordWildCards            #-}
  {-# LANGUAGE ScopedTypeVariables        #-}
  {-# LANGUAGE TemplateHaskell            #-}
  {-# LANGUAGE TypeApplications           #-}
  {-# OPTIONS_HADDOCK hide #-}
  -- |
  -- Module      : Data.Array.Accelerate.LLVM.CodeGen.Monad
  -- Copyright   : [2015..2020] The Accelerate Team
  -- License     : BSD3
  --
  -- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
  -- Stability   : experimental
  -- Portability : non-portable (GHC extensions)
  --
  
  module Data.Array.Accelerate.LLVM.CodeGen.Monad (
  
    CodeGen,
    evalCodeGen,
    liftCodeGen,
  
    -- codegen context
    getLLVMversion,
  
    -- declarations
    fresh, freshLocalName, freshGlobalName,
    declareGlobalVar,
    declareExternFunc,
    typedef,
    intrinsic,
  
    -- basic blocks
    Block,
    newBlock, setBlock, beginBlock, createBlocks,
  
    -- instructions
    instr, instr', do_, return_, retval_, br, cbr, switch, phi, phi', phi1,
    instr_,
  
    -- metadata
    addMetadata,
  
  ) where
  
  import Data.Array.Accelerate.Error
  import Data.Array.Accelerate.LLVM.CodeGen.IR
  import Data.Array.Accelerate.LLVM.CodeGen.Intrinsic
  import Data.Array.Accelerate.LLVM.CodeGen.Module
  import Data.Array.Accelerate.LLVM.CodeGen.Sugar                     ( IROpenAcc(..) )
  import Data.Array.Accelerate.LLVM.State                             ( LLVM )
  import Data.Array.Accelerate.LLVM.Target
  import Data.Array.Accelerate.Representation.Tag
  import Data.Array.Accelerate.Representation.Type
  import qualified Data.Array.Accelerate.Debug.Internal               as Debug
  
  import LLVM.AST.Type.Constant
  import LLVM.AST.Type.Downcast
  import LLVM.AST.Type.Instruction
  import LLVM.AST.Type.Metadata
  import LLVM.AST.Type.Name
  import LLVM.AST.Type.Operand
  import LLVM.AST.Type.Representation
  import LLVM.AST.Type.Terminator
  import qualified Data.Array.Accelerate.LLVM.Internal.LLVMPretty     as LP
  import qualified Data.Array.Accelerate.LLVM.Internal.LLVMPretty.PP  as LP
  import qualified Data.Array.Accelerate.LLVM.Internal.LLVMPretty.Triple.Parse as LP
  
  import Control.Applicative
  import Control.Monad
  import Control.Monad.Reader                                         ( ReaderT, MonadReader, runReaderT, ask, asks )
  import Control.Monad.State
  import Data.ByteString.Short                                        ( ShortByteString )
  import qualified Data.ByteString.Short.Char8                        as SBS8
  import Data.Function
  import Data.HashMap.Strict                                          ( HashMap )
  import Data.Sequence                                                ( Seq )
  import Data.String
  import Data.Text.Lazy.Builder                                       ( Builder )
  import Formatting
  import Prelude
  import qualified Data.Foldable                                      as F
  import qualified Data.HashMap.Strict                                as HashMap
  import qualified Data.Map.Strict                                    as Map
  import qualified Data.Sequence                                      as Seq
  import qualified Data.ByteString.Short                              as B
  
  
  -- Code generation
  -- ===============
  
  -- | The code generation state for scalar functions and expressions.
  --
  -- We use two records: one to hold all the code generation state as it walks the
  -- AST, and one for each of the basic blocks that are generated during the walk.
  --
  data CodeGenState = CodeGenState
    { blockChain          :: Seq Block                                      -- blocks for this function
    , globalvarTable      :: HashMap Label LP.Global                        -- global variable definitions
    , fundeclTable        :: HashMap Label LP.Declare                       -- global (external) function declarations
    , typedefTable        :: HashMap Label LP.Type                          -- global type definitions
    , metadataTable       :: HashMap ShortByteString (Seq [Maybe Metadata]) -- module metadata to be collected
    , intrinsicTable      :: HashMap ShortByteString Label                  -- standard math intrinsic functions
    , local               :: {-# UNPACK #-} !Word                           -- a name supply
    , global              :: {-# UNPACK #-} !Word                           -- a name supply for global variables
    }
  
  data Block = Block
    { blockLabel          :: {-# UNPACK #-} !Label -- block label
    , instructions        :: Seq LP.Stmt           -- stack of instructions
    , terminator          :: LP.Stmt               -- block terminator
    }
  
  -- | Code generation context: read-only additional context of the compilation.
  --
  -- This contains the LLVM version.
  data CodeGenContext = CodeGenContext
    { codegenLLVMversion :: LP.LLVMVer }
  
  newtype CodeGen target a = CodeGen
    { runCodeGen :: ReaderT CodeGenContext (StateT CodeGenState (LLVM target)) a }
    deriving (Functor, Applicative, Monad, MonadReader CodeGenContext, MonadState CodeGenState)
  
  liftCodeGen :: LLVM arch a -> CodeGen arch a
  liftCodeGen = CodeGen . lift . lift
  
  getLLVMversion :: CodeGen arch LP.LLVMVer
  getLLVMversion = asks codegenLLVMversion
  
  
  {-# INLINEABLE evalCodeGen #-}
  evalCodeGen
      :: forall arch aenv a. (HasCallStack, Target arch, Intrinsic arch)
      => CodeGen arch (IROpenAcc arch aenv a)
      -> LLVM    arch (Module    arch aenv a)
  evalCodeGen ll = do
    llvmver <- ask
    let context = CodeGenContext
          { codegenLLVMversion = llvmver }
    (IROpenAcc ks, st)   <- runStateT (runReaderT (runCodeGen ll) context)
                          $ CodeGenState
                              { blockChain        = initBlockChain
                              , globalvarTable    = HashMap.empty
                              , fundeclTable      = HashMap.empty
                              , typedefTable      = HashMap.empty
                              , metadataTable     = HashMap.empty
                              , intrinsicTable    = intrinsicForTarget @arch
                              , local             = 0
                              , global            = 0
                              }
  
    let (kernels, md)     = let (fs, as) = unzip [ (f , (LP.defName f, a)) | Kernel f a <- ks ]
                            in  (fs, Map.fromList as)
  
        createTypedefs    = map (\(n,t) -> LP.TypeDecl (labelToPrettyI n) t) . HashMap.toList
  
        (namedMd, unnamedMd) = createMetadata (metadataTable st)
  
    return $
      Module { moduleMetadata = md
             , unModule       = LP.Module
                              { LP.modSourceName = Nothing
                              , LP.modDataLayout = []  -- TODO: targetDataLayout @arch; this will be important for PTX
                              , LP.modTriple     =
                                  case targetTriple @arch of
                                    Just s -> LP.parseTriple (SBS8.unpack s)
                                    Nothing -> error "TODO: module target triple"
                              , LP.modTypes      = createTypedefs (typedefTable st)
                              , LP.modNamedMd    = namedMd
                              , LP.modUnnamedMd  = unnamedMd
                              , LP.modComdat     = mempty
                              , LP.modGlobals    = HashMap.elems (globalvarTable st)
                              , LP.modDeclares   = HashMap.elems (fundeclTable st)
                              , LP.modDefines    = kernels
                              , LP.modInlineAsm  = []
                              , LP.modAliases    = []
                              }
             }
  
  
  -- Basic Blocks
  -- ============
  
  -- | An initial block chain
  --
  initBlockChain :: HasCallStack => Seq Block
  initBlockChain
    = Seq.singleton
    $ Block "entry" Seq.empty (internalError "block has no terminator")
  
  
  -- | Create a new basic block, but don't yet add it to the block chain. You need
  -- to call 'setBlock' to append it to the chain, so that subsequent instructions
  -- are added to this block.
  --
  -- Note: [Basic blocks]
  --
  -- The names of basic blocks are generated based on the base name provided to
  -- the 'newBlock' function, as well as the current state (length) of the block
  -- stream. By not immediately adding new blocks to the stream, we have the
  -- advantage that:
  --
  --   1. Instructions are generated "in order", and are always appended to the
  --      stream. There is no need to search the stream for a block of the right
  --      name.
  --
  --   2. Blocks are named in groups, which helps readability. For example, the
  --      blocks for the then and else branches of a conditional, created at the
  --      same time, will be named similarly: 'if4.then' and 'if4.else', etc.
  --
  -- However, this leads to a slight awkwardness when walking the AST. Since a new
  -- naming group scheme is only applied *after* a call to 'setBlock',
  -- encountering (say) nested conditionals in the walk will generate logically
  -- distinct blocks that happen to have the same name. This means that
  -- instructions might be added to the wrong blocks, or the first set of blocks
  -- will be emitted empty and/or without a terminator.
  --
  newBlock :: HasCallStack => String -> CodeGen arch Block
  newBlock nm =
    state $ \s ->
      let idx     = Seq.length (blockChain s)
          label   = let (h,t) = break (== '.') nm in (h ++ shows idx t)
          next    = Block (fromString label) Seq.empty err
          err     = internalError ("block `" % string % "' has no terminator") label
      in
      ( next, s )
  
  
  -- | Add this block to the block stream. Any instructions pushed onto the stream
  -- by 'instr' and friends will now apply to this block.
  --
  setBlock :: Block -> CodeGen arch ()
  setBlock next =
    modify $ \s -> s { blockChain = blockChain s Seq.|> next }
  
  
  -- | Generate a new block and branch unconditionally to it.
  --
  beginBlock :: HasCallStack => String -> CodeGen arch Block
  beginBlock nm = do
    next <- newBlock nm
    _    <- br next
    setBlock next
    return next
  
  
  -- | Extract the block state and construct the basic blocks that form a function
  -- body. The block stream is re-initialised, but module-level state such as the
  -- global symbol table is left intact.
  --
  createBlocks :: HasCallStack => CodeGen arch [LP.BasicBlock]
  createBlocks
    = state
    $ \s -> let s'     = s { blockChain = initBlockChain, local = 0 }
                blocks = makeBlock `fmap` blockChain s
                m      = Seq.length (blockChain s)
                n      = F.foldl' (\i b -> i + Seq.length (instructions b)) 0 (blockChain s)
            in
            trace (bformat ("generated " % int % " instructions in " % int % " blocks") (n+m) m) ( F.toList blocks , s' )
    where
      makeBlock Block{..} =
        LP.BasicBlock (Just (labelToPrettyBL blockLabel)) (F.toList instructions ++ [terminator])
  
  
  -- Instructions
  -- ------------
  
  -- | Generate a fresh local reference
  --
  fresh :: TypeR a -> CodeGen arch (Operands a)
  fresh TupRunit         = return OP_Unit
  fresh (TupRpair t2 t1) = OP_Pair <$> fresh t2 <*> fresh t1
  fresh (TupRsingle t)   = ir t . LocalReference (PrimType (ScalarPrimType t)) <$> freshLocalName
  
  -- | Generate a fresh local (un)name
  --
  freshLocalName :: CodeGen arch (Name a)
  freshLocalName = state $ \s@CodeGenState{..} -> ( UnName local, s { local = local + 1 } )
  
  -- | Generate a fresh global (un)name
  --
  freshGlobalName :: CodeGen arch (Name a)
  freshGlobalName = state $ \s@CodeGenState{..} -> ( UnName global, s { global = global + 1 } )
  
  
  -- | Add an instruction to the state of the currently active block so that it is
  -- computed, and return the operand (LocalReference) that can be used to later
  -- refer to it.
  --
  instr :: HasCallStack => Instruction a -> CodeGen arch (Operands a)
  instr ins = ir (typeOf ins) <$> instr' ins
  
  instr' :: HasCallStack => Instruction a -> CodeGen arch (Operand a)
  instr' ins =
    -- LLVM-5 does not allow instructions of type void to have a name.
    case typeOf ins of
      VoidType -> do
        do_ ins
        return $ LocalReference VoidType (Name B.empty)
      --
      ty -> do
        name <- freshLocalName
        instr_ $ LP.Result (nameToPrettyI name) (downcast ins) []
        return $ LocalReference ty name
  
  -- | Execute an unnamed instruction
  --
  do_ :: HasCallStack => Instruction () -> CodeGen arch ()
  do_ ins = instr_ $ LP.Effect (downcast ins) []
  
  -- | Add raw assembly instructions to the execution stream
  --
  instr_ :: HasCallStack => LP.Stmt -> CodeGen arch ()
  instr_ ins =
    modify $ \s ->
      case Seq.viewr (blockChain s) of
        Seq.EmptyR  -> internalError "empty block chain"
        bs Seq.:> b -> s { blockChain = bs Seq.|> b { instructions = instructions b Seq.|> ins } }
  
  
  -- | Return void from a basic block
  --
  return_ :: HasCallStack => CodeGen arch ()
  return_ = void $ terminate Ret
  
  -- | Return a value from a basic block
  --
  retval_ :: HasCallStack => Operand a -> CodeGen arch ()
  retval_ x = void $ terminate (RetVal x)
  
  
  -- | Unconditional branch. Return the name of the block that was branched from.
  --
  br :: HasCallStack => Block -> CodeGen arch Block
  br target = terminate $ Br (blockLabel target)
  
  
  -- | Conditional branch. Return the name of the block that was branched from.
  --
  cbr :: HasCallStack => Operands Bool -> Block -> Block -> CodeGen arch Block
  cbr (OP_Bool cond) t f = terminate $ CondBr cond (blockLabel t) (blockLabel f)
  
  -- | Switch statement. Return the name of the block that was branched from.
  --
  switch :: HasCallStack => Operands TAG -> Block -> [(TAG, Block)] -> CodeGen arch Block
  switch tag def eqs = terminate $ Switch (op scalarType tag) (blockLabel def) [(ScalarConstant scalarType t, blockLabel b) | (t,b) <- eqs]
  
  -- | Add a phi node to the top of the current block
  --
  phi :: forall arch a. HasCallStack => TypeR a -> [(Operands a, Block)] -> CodeGen arch (Operands a)
  phi tp incoming = do
    crit  <- fresh tp
    block <- state $ \s -> case Seq.viewr (blockChain s) of
                             Seq.EmptyR -> internalError "empty block chain"
                             _ Seq.:> b -> ( b, s )
    phi' tp block crit incoming
  
  phi' :: HasCallStack => TypeR a -> Block -> Operands a -> [(Operands a, Block)] -> CodeGen arch (Operands a)
  phi' tp target = go tp
    where
      go :: TypeR t -> Operands t -> [(Operands t, Block)] -> CodeGen arch (Operands t)
      go TupRunit OP_Unit _
        = return OP_Unit
      go (TupRpair t2 t1) (OP_Pair n2 n1) inc
        = OP_Pair <$> go t2 n2 [ (x, b) | (OP_Pair x _, b) <- inc ]
                  <*> go t1 n1 [ (y, b) | (OP_Pair _ y, b) <- inc ]
      go (TupRsingle t) tup inc
        | LocalReference _ v <- op t tup = ir t <$> phi1 target v [ (op t x, b) | (x, b) <- inc ]
        | otherwise                      = internalError "expected critical variable to be local reference"
  
  
  phi1 :: HasCallStack => Block -> Name a -> [(Operand a, Block)] -> CodeGen arch (Operand a)
  phi1 target crit incoming =
    let cmp       = (==) `on` blockLabel
        update b  = b { instructions = LP.Result (nameToPrettyI crit) (downcast $ Phi t [ (p,blockLabel) | (p,Block{..}) <- incoming ]) [] Seq.<| instructions b }
        t         = case incoming of
                      []        -> internalError "no incoming values specified"
                      (o,_):_   -> case typeOf o of
                                     VoidType    -> internalError "operand has void type"
                                     PrimType x  -> x
    in
    state $ \s ->
      case Seq.findIndexR (cmp target) (blockChain s) of
        Nothing -> internalError "unknown basic block"
        Just i  -> ( LocalReference (PrimType t) crit
                   , s { blockChain = Seq.adjust update i (blockChain s) } )
  
  
  -- | Add a termination condition to the current instruction stream. Also return
  -- the block that was just terminated.
  --
  terminate :: HasCallStack => Terminator a -> CodeGen arch Block
  terminate term =
    state $ \s ->
      case Seq.viewr (blockChain s) of
        Seq.EmptyR  -> internalError "empty block chain"
        bs Seq.:> b -> ( b, s { blockChain = bs Seq.|> b { terminator = LP.Effect (downcast term) [] } } )
  
  
  -- | Add a global variable declaration to the module
  --
  declareGlobalVar :: HasCallStack => LP.Global -> CodeGen arch ()
  declareGlobalVar g =
    let unique (Just q) | g /= q    = internalError "duplicate symbol"
                        | otherwise = Just g
        unique _                    = Just g
  
        name = case LP.globalSym g of
                 LP.Symbol n -> Label (fromString n)
    in
    modify (\s -> s { globalvarTable = HashMap.alter unique name (globalvarTable s) })
  
  
  -- | Add an external function declaration to the module
  --
  declareExternFunc :: HasCallStack => LP.Declare -> CodeGen arch ()
  declareExternFunc g =
    let unique (Just q) | g /= q    = internalError "duplicate symbol"
                        | otherwise = Just g
        unique _                    = Just g
  
        name = case LP.decName g of
                 LP.Symbol n -> Label (fromString n)
    in
    modify (\s -> s { fundeclTable = HashMap.alter unique name (fundeclTable s) })
  
  
  -- | Add a global type definition
  --
  typedef :: HasCallStack => Label -> LP.Type -> CodeGen arch ()
  typedef name t =
    let unique (Just s) | t /= s    = internalError "duplicate typedef"
                        | otherwise = Just t
        unique _                    = Just t
    in
    modify (\s -> s { typedefTable = HashMap.alter unique name (typedefTable s) })
  
  
  -- | Get name of the corresponding intrinsic function implementing a given C
  -- function. If there is no mapping, the C function name is used.
  --
  intrinsic :: ShortByteString -> CodeGen arch Label
  intrinsic key =
    state $ \s ->
      let name = HashMap.lookupDefault (Label key) key (intrinsicTable s)
      in  (name, s)
  
  
  
  -- Metadata
  -- ========
  
  -- | Insert a metadata key/value pair into the current module.
  --
  addMetadata :: ShortByteString -> [Maybe Metadata] -> CodeGen arch ()
  addMetadata key val =
    modify $ \s ->
      s { metadataTable = HashMap.insertWith (flip (Seq.><)) key (Seq.singleton val) (metadataTable s) }
  
  
  -- | Generate the metadata definitions for the file. Every key in the map
  -- represents a named metadata definition. The values associated with that key
  -- represent the metadata node definitions that will be attached to that
  -- definition.
  --
  createMetadata :: HashMap ShortByteString (Seq [Maybe Metadata]) -> ([LP.NamedMd], [LP.UnnamedMd])
  createMetadata md = go (HashMap.toList md) (Seq.empty, Seq.empty)
    where
      go :: [(ShortByteString, Seq [Maybe Metadata])]
         -> (Seq LP.NamedMd, Seq LP.UnnamedMd) -- accumulator of (names, metadata)
         -> ([LP.NamedMd], [LP.UnnamedMd])
      go []     (k,d) = (F.toList k, F.toList d)
      go (x:xs) (k,d) =
        let (k',d') = meta (Seq.length d) x
        in  go xs (k Seq.|> k', d Seq.>< d')
  
      meta :: Int                                         -- number of metadata node definitions so far
           -> (ShortByteString, Seq [Maybe Metadata])     -- current assoc of the metadata map
           -> (LP.NamedMd, Seq LP.UnnamedMd)
      meta n (key, vals)
        = let node i      = i + n
              nodes       = Seq.mapWithIndex (\i x -> LP.UnnamedMd (node i) (LP.ValMdNode (downcast (F.toList x))) False) vals
              name        = LP.NamedMd (SBS8.unpack key) [ node i | i <- [0 .. Seq.length vals - 1] ]
          in
          (name, nodes)
  
  
  -- Debug
  -- =====
  
  {-# INLINE trace #-}
  trace :: Builder -> a -> a
  trace msg = Debug.trace Debug.dump_cc ("llvm: " <> msg)
  
  