content:  ------------------------------------------------------------------------
  -- |
  -- Module           : Lang.Crucible.LLVM.Ctors
  -- Description      : Extract and manipulate the @llvm.global_ctors@ variable
  -- Copyright        : (c) Galois, Inc 2018
  -- License          : BSD3
  -- Maintainer       : Rob Dockins <rdockins@galois.com>
  -- Stability        : provisional
  ------------------------------------------------------------------------
  {-# LANGUAGE DataKinds             #-}
  {-# LANGUAGE DeriveDataTypeable    #-}
  {-# LANGUAGE DeriveGeneric         #-}
  {-# LANGUAGE FlexibleContexts      #-}
  {-# LANGUAGE GADTs                 #-}
  {-# LANGUAGE ImplicitParams        #-}
  {-# LANGUAGE OverloadedStrings     #-}
  {-# LANGUAGE RankNTypes            #-}
  {-# LANGUAGE ScopedTypeVariables   #-}
  {-# LANGUAGE TypeOperators         #-}
  
  module Lang.Crucible.LLVM.Ctors
    ( Ctor(..)
    , globalCtors
    , callCtors
    , callAllCtors
    , callCtorsCFG
    ) where
  
  import           Data.Data (Data)
  import           Data.IORef (newIORef)
  import           Data.String(fromString)
  import           Data.Typeable (Typeable)
  import qualified Data.Text as Text
  import           GHC.Generics (Generic)
  import           Data.Parameterized.Nonce
  
  import           Control.Monad (forM, forM_)
  import           Control.Monad.Except (MonadError(..))
  import           Data.List (find, sortBy)
  import           Data.Ord (comparing, Down(..))
  import           Data.Maybe (fromMaybe)
  
  import qualified Text.LLVM.AST as L
  
  import           Lang.Crucible.LLVM.Translation.Instruction (callOrdinaryFunction)
  import           Lang.Crucible.LLVM.Translation.Monad (LLVMGenerator, LLVMState(..))
  
  -- Generating CFGs
  
  import           Data.Map.Strict (empty)
  import           Data.Text (Text)
  import           GHC.TypeNats
  
  import qualified Data.Parameterized.Context.Unsafe as Ctx
  
  import           What4.FunctionName (functionNameFromText)
  import           What4.ProgramLoc (Position(InternalPos))
  
  import qualified Lang.Crucible.CFG.Core as Core
  import           Lang.Crucible.CFG.Expr (App(EmptyApp))
  import           Lang.Crucible.CFG.Generator (FunctionDef, defineFunction)
  import           Lang.Crucible.CFG.Reg (Expr(App))
  import qualified Lang.Crucible.CFG.Reg as Reg
  import           Lang.Crucible.CFG.SSAConversion (toSSA)
  import           Lang.Crucible.FunctionHandle (HandleAllocator, mkHandle')
  import           Lang.Crucible.Types (UnitType, TypeRepr(UnitRepr))
  import           Lang.Crucible.LLVM.Extension (LLVM, ArchWidth)
  import           Lang.Crucible.LLVM.Translation.Monad (LLVMContext, _llvmTypeCtx, malformedLLVMModule)
  import           Lang.Crucible.LLVM.Types (HasPtrWidth)
  
  {- Example:
  
  @llvm.global_ctors = appending global [3 x { i32, void ()*, i8* }] [{ i32, void ()*, i8* } { i32 65535, void ()* @_GLOBAL__sub_I_HkdfTest.cpp, i8* null }, { i32, void ()*, i8* } { i32 65535, void ()* @_GLOBAL__sub_I_gtest_all.cc, i8* null }, { i32, void ()*, i8* } { i32 65535, void ()* @_GLOBAL__sub_I_iostream.cpp, i8* null }]
  
  -}
  
  -- | A representation of well-typed inhabitants of the @llvm.global_ctors@ array
  --
  -- See https://llvm.org/docs/LangRef.html#the-llvm-global-ctors-global-variable
  data Ctor = Ctor
    { ctorPriority :: Integer
    , ctorFunction :: L.Symbol
    , ctorData     :: Maybe L.Symbol
    } deriving (Data, Eq, Generic, Ord, Show, Typeable)
  
  -- | Get the global variable representing @llvm.global_ctors@.
  getGlobalCtorsGlobal :: L.Module -> Maybe L.Global
  getGlobalCtorsGlobal mod_ =
    let symb = L.Symbol "llvm.global_ctors"
    in find (\x -> L.globalSym x == symb) (L.modGlobals mod_)
  
  -- | Unpack a @ctors@ value of type @{ i32, void ()*, i8* }@ from the AST
  extractCtors :: L.Value -> Maybe Ctor
  extractCtors val =
    case val of
      -- This is permissive about the integer widths... No reason to get caught up.
      L.ValStruct [ L.Typed (L.PrimType (L.Integer _w0)) (L.ValInteger priority)
                  , L.Typed (L.PtrTo (L.FunTy (L.PrimType L.Void) [] _bool)) (L.ValSymbol symb)
                  , L.Typed (L.PtrTo (L.PrimType (L.Integer _w1))) data0_
                  ] -> Just . Ctor priority symb $
                         case data0_ of
                           L.ValSymbol data_ -> Just data_
                           _                 -> Nothing
      _ -> Nothing
  
  -- | Unpack and sort the values in @llvm.global_ctors@ by priority
  globalCtors :: (MonadError String m)
              => L.Module
              -> m [Ctor]
  globalCtors mod_ =
    case getGlobalCtorsGlobal mod_ >>= L.globalValue of -- in the Maybe monad
      Just (L.ValArray _ty vs) -> do
  
        -- Assert that each value is of the expected type.
        vs' <- forM vs $ \v ->
          fromMaybe
            (throwError $ unlines $ [ "Ill-typed value in llvm.global_ctors: "
                                    , show v
                                    ])
            (pure <$> extractCtors v)
  
        -- Sort the values by priority, highest to lowest.
        pure (sortBy (comparing (Down . ctorPriority)) vs')
  
      -- @llvm.ctors value not found, assume there are no global_ctors to run
      Nothing -> return []
  
      Just v  -> throwError $ unlines $
        [ "llvm.global_ctors wasn't an array"
        , "Value: " ++ show v
        ]
  
  ----------------------------------------------------------------------
  -- ** callCtors
  
  -- | Call some or all of the functions in @llvm.global_ctors@
  callCtors :: (Ctor -> Bool) -- ^ Filter function
            -> L.Module
            -> LLVMGenerator s arch UnitType (Expr LLVM s UnitType)
  callCtors select mod_ = do
    let err msg = malformedLLVMModule "Error loading @llvm.global_ctors" [fromString msg]
    let ty = L.FunTy (L.PrimType L.Void) [] False
  
    ctors <- either err (pure . filter select) (globalCtors mod_)
    forM_ ctors $ \ctor ->
      callOrdinaryFunction Nothing False ty (L.ValSymbol (ctorFunction ctor)) [] (\_ -> pure ())
    return (App EmptyApp)
  
  -- | Call each function in @llvm.global_ctors@ in order of decreasing priority
  callAllCtors :: L.Module -> LLVMGenerator s arch UnitType (Expr LLVM s UnitType)
  callAllCtors = callCtors (const True)
  
  ----------------------------------------------------------------------
  -- ** callCtorsCFG
  
  -- | Make a 'LLVMGenerator' into a CFG by making it a function with no arguments
  -- that returns unit.
  generatorToCFG :: forall arch wptr ret. (HasPtrWidth wptr, wptr ~ ArchWidth arch, 16 <= wptr)
                 => Text
                 -> HandleAllocator
                 -> LLVMContext arch
                 -> (forall s. LLVMGenerator s arch ret (Expr LLVM s ret))
                 -> TypeRepr ret
                 -> IO (Core.SomeCFG LLVM Core.EmptyCtx ret)
  generatorToCFG name halloc llvmctx gen ret = do
    ref <- newIORef []
    let ?lc = _llvmTypeCtx llvmctx
    let def :: forall args. FunctionDef LLVM (LLVMState arch) args ret IO
        def _inputs = (state, gen)
          where state = LLVMState { _identMap     = empty
                                  , _blockInfoMap = empty
                                  , llvmContext   = llvmctx
                                  , _translationWarnings = ref
                                  , _functionSymbol = L.Symbol (Text.unpack name)
                                  }
  
    hand <- mkHandle' halloc (functionNameFromText name) Ctx.empty ret
    sng <- newIONonceGenerator
    (Reg.SomeCFG g, []) <- defineFunction InternalPos sng hand def
    return $! toSSA g
  
  -- | Create a CFG that calls some of the functions in @llvm.global_ctors@.
  callCtorsCFG :: forall arch wptr. (HasPtrWidth wptr, wptr ~ ArchWidth arch, 16 <= wptr)
               => (Ctor -> Bool) -- ^ Filter function
               -> L.Module
               -> HandleAllocator
               -> LLVMContext arch
               -> IO (Core.SomeCFG LLVM Core.EmptyCtx UnitType)
  callCtorsCFG select mod_ halloc llvmctx = do
    generatorToCFG "llvm_global_ctors" halloc llvmctx (callCtors select mod_) UnitRepr
  