------------------------------------------------------------------------
-- |
-- Module           : Lang.Crucible.LLVM.ArraySizeProfile
-- Description      : Execution feature to observe argument buffer sizes
-- Copyright        : (c) Galois, Inc 2020
-- License          : BSD3
-- Maintainer       : Samuel Breese <sbreese@galois.com>
-- Stability        : provisional
------------------------------------------------------------------------

{-# Options -Wall #-}
{-# Language TemplateHaskell #-}
{-# Language OverloadedStrings #-}
{-# Language LambdaCase #-}
{-# Language MultiWayIf #-}
{-# Language ImplicitParams #-}
{-# Language ViewPatterns #-}
{-# Language PatternSynonyms #-}
{-# Language BangPatterns #-}
{-# Language FlexibleContexts #-}
{-# Language ScopedTypeVariables #-}
{-# Language DataKinds #-}
{-# Language KindSignatures #-}
{-# Language TypeFamilies #-}
{-# Language TypeApplications #-}
{-# Language GADTs #-}

module Lang.Crucible.LLVM.ArraySizeProfile
 ( FunctionProfile(..), funProfileName, funProfileArgs
 , ArgProfile(..), argProfileSize, argProfileInitialized
 , arraySizeProfile
 ) where

import Control.Lens.TH

import Control.Lens

import Data.Type.Equality (testEquality)
import Data.IORef
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Vector as Vector
import Data.Map (Map)
import qualified Data.Map as Map

import qualified Data.BitVector.Sized as BV
import Data.Parameterized.SymbolRepr
import qualified Data.Parameterized.Context as Ctx
import Data.Parameterized.TraversableFC

import qualified Lang.Crucible.Backend as C
import qualified Lang.Crucible.CFG.Core as C
import qualified Lang.Crucible.Simulator.CallFrame as C
import qualified Lang.Crucible.Simulator.EvalStmt as C
import qualified Lang.Crucible.Simulator.ExecutionTree as C
import qualified Lang.Crucible.Simulator.GlobalState as C
import qualified Lang.Crucible.Simulator.Intrinsics as C
import qualified Lang.Crucible.Simulator.RegMap as C

import qualified Lang.Crucible.LLVM.DataLayout as C
import qualified Lang.Crucible.LLVM.Extension as C
import qualified Lang.Crucible.LLVM.MemModel as C
import qualified Lang.Crucible.LLVM.MemModel.Generic as G
import qualified Lang.Crucible.LLVM.Translation.Monad as C

import qualified What4.Interface as W4

------------------------------------------------------------------------
-- Profiles

data ArgProfile = ArgProfile
  { _argProfileSize :: Maybe Int
  , _argProfileInitialized :: Bool
  } deriving (Show, Eq, Ord)
makeLenses ''ArgProfile

data FunctionProfile = FunctionProfile
  { _funProfileName :: Text
  , _funProfileArgs :: [ArgProfile]
  } deriving (Show, Eq, Ord)
makeLenses ''FunctionProfile

------------------------------------------------------------------------
-- Learning a profile from an ExecState

ptrStartsAlloc ::
  W4.IsExpr (W4.SymExpr sym) =>
  C.LLVMPtr sym w ->
  Bool
ptrStartsAlloc (C.llvmPointerView -> (_, W4.asBV -> Just (BV.BV 0))) = True
ptrStartsAlloc _ = False

ptrAllocSize ::
  forall sym w.
  C.IsSymInterface sym =>
  G.Mem sym ->
  C.LLVMPtr sym w ->
  Maybe Int
ptrAllocSize mem (C.llvmPointerView -> (blk, _)) =
  do a <- W4.asNat blk
     G.AllocInfo _ msz _ _ _ <- G.possibleAllocInfo a (G.memAllocs mem)
     sz <- msz
     fromIntegral <$> BV.asUnsigned <$> W4.asBV sz

ptrArraySize ::
  C.IsSymInterface sym =>
  G.Mem sym ->
  C.LLVMPtr sym w ->
  Maybe Int
ptrArraySize mem ptr
  | ptrStartsAlloc ptr = ptrAllocSize mem ptr
  | otherwise = Nothing

ptrIsInitialized ::
  ( C.IsSymInterface sym, C.HasLLVMAnn sym, C.HasPtrWidth w
  , ?memOpts :: C.MemOptions ) =>
  sym ->
  G.Mem sym ->
  C.LLVMPtr sym w ->
  IO Bool
ptrIsInitialized sym mem ptr =
  G.readMem sym C.PtrWidth Nothing ptr (C.bitvectorType 1) C.noAlignment mem >>= \case
  C.NoErr{} -> pure True
  _ -> pure False

intrinsicArgProfile ::
  ( C.IsSymInterface sym, C.HasLLVMAnn sym, C.HasPtrWidth w
  , ?memOpts :: C.MemOptions ) =>
  sym ->
  G.Mem sym ->
  SymbolRepr nm ->
  C.CtxRepr ctx ->
  C.Intrinsic sym nm ctx ->
  IO ArgProfile
intrinsicArgProfile sym mem
  (testEquality (knownSymbol :: SymbolRepr "LLVM_pointer") -> Just Refl)
  (Ctx.Empty Ctx.:> C.BVRepr (testEquality ?ptrWidth -> Just Refl)) i =
   ArgProfile (ptrArraySize mem i) <$> ptrIsInitialized sym mem i
intrinsicArgProfile _ _ _ _ _ = pure $ ArgProfile Nothing False

regValueArgProfile ::
  ( C.IsSymInterface sym, C.HasLLVMAnn sym, C.HasPtrWidth w
  , ?memOpts :: C.MemOptions ) =>
  sym ->
  G.Mem sym ->
  C.TypeRepr tp ->
  C.RegValue sym tp ->
  IO ArgProfile
regValueArgProfile sym mem (C.IntrinsicRepr nm ctx) i = intrinsicArgProfile sym mem nm ctx i
regValueArgProfile _ _ _ _ = pure $ ArgProfile Nothing False

regEntryArgProfile ::
  ( C.IsSymInterface sym, C.HasLLVMAnn sym, C.HasPtrWidth w
  , ?memOpts :: C.MemOptions ) =>
  sym ->
  G.Mem sym ->
  C.RegEntry sym tp ->
  IO ArgProfile
regEntryArgProfile sym mem (C.RegEntry t v) = regValueArgProfile sym mem t v

newtype Wrap a (b :: C.CrucibleType) = Wrap { unwrap :: a }
argProfiles ::
  ( C.IsSymInterface sym, C.HasLLVMAnn sym, C.HasPtrWidth w
  , ?memOpts :: C.MemOptions ) =>
  sym ->
  G.Mem sym ->
  Ctx.Assignment (C.RegEntry sym) ctx ->
  IO [ArgProfile]
argProfiles sym mem as =
  sequence (Vector.toList $ Ctx.toVector (fmapFC (Wrap . regEntryArgProfile sym mem) as) unwrap)

------------------------------------------------------------------------
-- Execution feature for learning profiles

updateProfiles ::
  ( C.IsSymInterface sym, C.HasLLVMAnn sym, C.HasPtrWidth (C.ArchWidth arch)
  , ?memOpts :: C.MemOptions ) =>
  C.LLVMContext arch ->
  IORef (Map Text [FunctionProfile]) ->
  C.ExecState p sym ext rtp ->
  IO ()
updateProfiles llvm cell state
  | C.CallState _ (C.CrucibleCall _ frame) sim <- state
  , C.CallFrame { C._frameCFG = cfg, C._frameRegs = regs } <- frame
  , Just mem <- C.memImplHeap <$> C.lookupGlobal (C.llvmMemVar llvm) (sim ^. C.stateGlobals)
  = do
      argProfs <- argProfiles (sim ^. C.stateSymInterface) mem $ C.regMap regs
      modifyIORef' cell $ \profs ->
        let name = Text.pack . show $ C.cfgHandle cfg
            funProf = FunctionProfile name argProfs
        in case Map.lookup name profs of
             Nothing -> Map.insert name [funProf] profs
             Just variants
               | funProf `elem` variants -> profs
               | otherwise -> Map.insert name (funProf:variants) profs
  | otherwise = pure ()

arraySizeProfile ::
  forall sym ext arch p rtp.
  ( C.IsSymInterface sym, C.HasLLVMAnn sym, C.HasPtrWidth (C.ArchWidth arch)
  , ?memOpts :: C.MemOptions ) =>
  C.LLVMContext arch ->
  IORef (Map Text [FunctionProfile]) ->
  IO (C.ExecutionFeature p sym ext rtp)
arraySizeProfile llvm profiles = do
  pure . C.ExecutionFeature $ \s -> do
    updateProfiles llvm profiles s
    pure C.ExecutionFeatureNoChange