{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Lang.Crucible.Simulator.GlobalState
  ( SymGlobalState
  , emptyGlobals
  , GlobalEntry(..)
  , insertGlobal
  , lookupGlobal
  , insertRef
  , lookupRef
  , dropRef
  , updateRef
  , globalPushBranch
  , globalAbortBranch
  , globalMuxFn
  ) where

import           Control.Applicative ((<|>))
import           Control.Monad.Trans.Class (lift)
import           Data.Functor.Identity
import           Data.Kind

import qualified Data.Parameterized.Map as MapF
import           Data.Parameterized.TraversableF

import           What4.Interface
import           What4.Partial
import           What4.ProgramLoc

import           Lang.Crucible.CFG.Core
import           Lang.Crucible.FunctionHandle
import           Lang.Crucible.Simulator.Intrinsics
import           Lang.Crucible.Simulator.RegMap
import           Lang.Crucible.Backend
import           Lang.Crucible.Panic(panic)

-- | As a map element, type @GlobalEntry sym tp@ models the contents
-- of a 'GlobalVar', which is always defined.
newtype GlobalEntry (sym :: Type) (tp :: CrucibleType) =
  GlobalEntry { globalEntryValue :: RegValue sym tp }

-- | Type @RefCellContents sym tp@ models the contents of a @RefCell@,
-- which holds a partial value. A @RefCell@ not found in the map is
-- considered to represent the 'Unassigned' value.
data RefCellContents (sym :: Type) (tp :: CrucibleType)
  = RefCellContents !(Pred sym) !(RegValue sym tp)

-- | Type @RefCellUpdate sym tp@ models an update to the contents of a
-- @RefCell@. The value @RefCellUpdate Unassigned@ represents the
-- deletion of a @RefCell@.
newtype RefCellUpdate (sym :: Type) (tp :: CrucibleType) =
  RefCellUpdate (PartExpr (Pred sym) (RegValue sym tp))

------------------------------------------------------------------------
-- GlobalTable

data GlobalTable f g =
  GlobalTable
  { globalVars :: !(MapF.MapF GlobalVar f)
  , globalRefs :: !(MapF.MapF RefCell g)
  }

updateGlobalVars ::
  (MapF.MapF GlobalVar v -> MapF.MapF GlobalVar v) ->
  GlobalTable v r -> GlobalTable v r
updateGlobalVars f (GlobalTable vs rs) = GlobalTable (f vs) rs

updateGlobalRefs ::
  (MapF.MapF RefCell r -> MapF.MapF RefCell r) ->
  GlobalTable v r -> GlobalTable v r
updateGlobalRefs f (GlobalTable vs rs) = GlobalTable vs (f rs)

-- | The empty set of global variable updates.
emptyGlobalTable :: GlobalTable v r
emptyGlobalTable = GlobalTable MapF.empty MapF.empty

-- | Take the union of two sets of updates, preferring elements from
-- the first set when duplicate keys are encountered.
mergeGlobalTable :: GlobalTable v r -> GlobalTable v r -> GlobalTable v r
mergeGlobalTable (GlobalTable vs1 rs1) (GlobalTable vs2 rs2) =
  GlobalTable (MapF.union vs1 vs2) (MapF.union rs1 rs2)

-- | Maps from global variables and global references to their values.
type GlobalContents sym = GlobalTable (GlobalEntry sym) (RefCellContents sym)

-- | A collection of updates to the global variables and global
-- references that have happened since a branch push.
type GlobalUpdates sym = GlobalTable (GlobalEntry sym) (RefCellUpdate sym)

-- | Apply a set of updates to the contents of global memory.
-- @GlobalVar@s cannot be deleted, so we just merge the @GlobalVar@
-- maps, preferring entries from the first argument. When merging the
-- @RefCell@ maps, a @RefCellUpdate Unassigned@ entry causes the
-- corresponding entry in the old map to be deleted.
applyGlobalUpdates :: forall sym . GlobalUpdates sym -> GlobalContents sym -> GlobalContents sym
applyGlobalUpdates (GlobalTable vs1 rs1) (GlobalTable vs2 rs2) =
  GlobalTable (MapF.union vs1 vs2) (runIdentity (MapF.mergeWithKeyM both left Identity rs1 rs2))
  where
    upd :: forall tp. RefCellUpdate sym tp -> Maybe (RefCellContents sym tp)
    upd (RefCellUpdate Unassigned) = Nothing
    upd (RefCellUpdate (PE p e)) = Just (RefCellContents p e)

    both :: forall tp. RefCell tp -> RefCellUpdate sym tp -> RefCellContents sym tp -> Identity (Maybe (RefCellContents sym tp))
    both _ u _ = Identity (upd u)

    left :: MapF.MapF RefCell (RefCellUpdate sym) -> Identity (MapF.MapF RefCell (RefCellContents sym))
    left m = Identity (MapF.mapMaybe upd m)

------------------------------------------------------------------------
-- GlobalFrames

-- | The state of global memory as a stack of changes separated by
-- branch pushes. The second parameter of 'BranchFrame' caches
-- the combined view of memory as of the previous branch push.
data GlobalFrames (sym :: Type) =
    InitialFrame !(GlobalContents sym)
  | BranchFrame !(GlobalUpdates sym) (GlobalContents sym) !(GlobalFrames sym)

-- | The depth of this value represents the number of symbolic
-- branches currently pending. We use this primarily as a sanity check
-- to help find bugs where we fail to match up calls to
-- 'globalPushBranch' with 'globalAbortBranch'/'globalMuxFn'.
globalPendingBranches :: GlobalFrames sym -> Int
globalPendingBranches (InitialFrame _) = 0
globalPendingBranches (BranchFrame _ _ gf) = 1 + globalPendingBranches gf

-- | The empty set of global variable bindings.
emptyGlobalFrames :: GlobalFrames sym
emptyGlobalFrames = InitialFrame emptyGlobalTable

------------------------------------------------------------------------
-- GlobalTable

-- | A map from global variables to their value.
data SymGlobalState (sym :: Type) =
  GlobalState
  { globalFrames :: !(GlobalFrames sym)
    -- ^ The stack of updates to global memory, separated by branch
    -- pushes. This field only contains values with ordinary crucible
    -- types (i.e. not intrinsic), which do not have mux operations
    -- that require branch push notifications.
  , globalIntrinsics :: !(GlobalContents sym)
    -- ^ The set of updates since initialization to vars and refs that
    -- have intrinsic or "any" types, which must be notified of
    -- branch pushes.
  }

-- | The empty set of global variable bindings.
emptyGlobals :: SymGlobalState sym
emptyGlobals = GlobalState emptyGlobalFrames emptyGlobalTable

-- | Test whether this type could be an intrinsic type, which must be
-- notified of branch pushes and aborts.
needsNotification :: TypeRepr tp -> Bool
needsNotification tr =
  case tr of
    IntrinsicRepr{} -> True
    AnyRepr -> True
    _ -> False

-- | Lookup a global variable in the state.
lookupGlobal :: GlobalVar tp -> SymGlobalState sym -> Maybe (RegValue sym tp)
lookupGlobal g gst
  | needsNotification (globalType g) =
      globalEntryValue <$> MapF.lookup g (globalVars (globalIntrinsics gst))
  | otherwise =
      globalEntryValue <$> go (globalFrames gst)
  where
    -- We never have to search more than one level deep, because the
    -- 'BranchFrame' constructor caches the combined contents of the
    -- rest of the 'GlobalFrames'.
    go (InitialFrame c) = MapF.lookup g (globalVars c)
    go (BranchFrame u c _) = MapF.lookup g (globalVars u) <|> MapF.lookup g (globalVars c)

-- | Set the value of a global in the state, or create a new global variable.
insertGlobal ::
  GlobalVar tp ->
  RegValue sym tp ->
  SymGlobalState sym ->
  SymGlobalState sym
insertGlobal g v gst
  | needsNotification (globalType g) =
      gst{ globalIntrinsics = updateGlobalVars (MapF.insert g x) (globalIntrinsics gst) }
  | otherwise =
      gst{ globalFrames = upd (globalFrames gst) }
  where
    x = GlobalEntry v
    upd (InitialFrame c) = InitialFrame (updateGlobalVars (MapF.insert g x) c)
    upd (BranchFrame u c gf) = BranchFrame (updateGlobalVars (MapF.insert g x) u) c gf
    -- NOTE: While global variables can be updated within branches, it
    -- should probably be forbidden to create a new global variable in
    -- a branch. However, we don't check for this currently. (An error
    -- will happen later when merging the branches if the set of
    -- global variables does not match.)

-- | Look up the value of a reference cell in the state.
lookupRef :: RefCell tp -> SymGlobalState sym -> PartExpr (Pred sym) (RegValue sym tp)
lookupRef r gst
  | needsNotification (refType r) = post $ MapF.lookup r (globalRefs (globalIntrinsics gst))
  | otherwise = go (globalFrames gst)
  where
    -- We never have to search more than one level deep, because the
    -- 'BranchFrame' constructor caches the combined contents of the
    -- rest of the 'GlobalFrames'.
    post = maybe Unassigned (\(RefCellContents p e) -> PE p e)
    go (InitialFrame c) = post $ MapF.lookup r (globalRefs c)
    go (BranchFrame u c _) =
      case MapF.lookup r (globalRefs u) of
        Just (RefCellUpdate pe) -> pe
        Nothing -> post $ MapF.lookup r (globalRefs c)

-- | Set the value of a reference cell in the state.
insertRef ::
  IsExprBuilder sym =>
  sym ->
  RefCell tp ->
  RegValue sym tp ->
  SymGlobalState sym ->
  SymGlobalState sym
insertRef sym r v = updateRef r (PE (truePred sym) v)

-- | Write a partial value to a reference cell in the state.
updateRef ::
  -- IsExprBuilder sym =>
  RefCell tp ->
  PartExpr (Pred sym) (RegValue sym tp) ->
  SymGlobalState sym ->
  SymGlobalState sym
updateRef r pe gst
  | needsNotification (refType r) =
      gst{ globalIntrinsics = updateGlobalRefs ins (globalIntrinsics gst) }
  | otherwise =
      gst{ globalFrames = upd (globalFrames gst) }
  where
    ins =
      case pe of
        Unassigned -> MapF.delete r
        PE p e -> MapF.insert r (RefCellContents p e)
    upd (InitialFrame c) = InitialFrame (updateGlobalRefs ins c)
    upd (BranchFrame u c gf) =
      BranchFrame (updateGlobalRefs (MapF.insert r (RefCellUpdate pe)) u) c gf

-- | Reset a reference cell to the uninitialized state. @'dropRef' r@ is
-- equivalent to @'updateRef' r 'Unassigned'@.
dropRef :: RefCell tp -> SymGlobalState sym -> SymGlobalState sym
dropRef r = updateRef r Unassigned

-- | Mark a branch point in the global state. Later calls to
-- 'globalMuxFn' will assume that the input states are identical up
-- until the most recent branch point.
globalPushBranch ::
  forall sym .
  IsSymInterface sym =>
  sym ->
  IntrinsicTypes sym ->
  SymGlobalState sym ->
  IO (SymGlobalState sym)
globalPushBranch sym iTypes (GlobalState gf (GlobalTable vs rs)) =
  do -- Notify intrinsic-typed vars and refs of the branch push.
     vs' <- MapF.traverseWithKey
            (\v (GlobalEntry e) ->
              GlobalEntry <$> pushBranchForType sym iTypes (globalType v) e)
            vs
     rs' <- MapF.traverseWithKey
            (\r (RefCellContents p e) ->
              RefCellContents p <$> pushBranchForType sym iTypes (refType r) e)
            rs
     --loc <- getCurrentProgramLoc sym
     --putStrLn $ unwords ["PUSH BRANCH:", show d, show $ plSourceLoc loc]
     let gf' = BranchFrame emptyGlobalTable cache gf
     return (GlobalState gf' (GlobalTable vs' rs'))
  where
    cache =
      case gf of
        InitialFrame c -> c
        BranchFrame u c _ -> applyGlobalUpdates u c

-- | Merge a set of updates into the outermost frame of a stack of global frames.
abortBranchFrame :: GlobalUpdates sym -> GlobalFrames sym -> GlobalFrames sym
abortBranchFrame u (InitialFrame c) = InitialFrame (applyGlobalUpdates u c)
abortBranchFrame u (BranchFrame u' c gf) = BranchFrame (mergeGlobalTable u u') c gf

-- | Remove the most recent branch point marker, and thus cancel the
-- effect of the most recent 'globalPushBranch'.
globalAbortBranch ::
  forall sym .
  IsSymInterface sym =>
  sym ->
  IntrinsicTypes sym ->
  SymGlobalState sym ->
  IO (SymGlobalState sym)

globalAbortBranch sym iTypes (GlobalState (BranchFrame u _ gf) (GlobalTable vs rs)) =
  do -- Notify intrinsic-typed vars and refs of the branch abort.
     vs' <- MapF.traverseWithKey
            (\v (GlobalEntry e) ->
              GlobalEntry <$> abortBranchForType sym iTypes (globalType v) e)
            vs
     rs' <- MapF.traverseWithKey
            (\r (RefCellContents p e) ->
              RefCellContents p <$> abortBranchForType sym iTypes (refType r) e)
            rs
     --loc <- getCurrentProgramLoc sym
     --putStrLn $ unwords ["ABORT BRANCH:", show (d-1), show $ plSourceLoc loc]
     let gf' = abortBranchFrame u gf
     return (GlobalState gf' (GlobalTable vs' rs'))

globalAbortBranch sym _ (GlobalState (InitialFrame _) _) =
  do loc <- getCurrentProgramLoc sym
     panic "GlobalState.globalAbortBranch"
       [ "Attempting to commit global changes at branch depth 0"
       , "*** Location: " ++ show (plSourceLoc loc)
       ]

muxPartialRegForType ::
  IsSymInterface sym =>
  sym ->
  IntrinsicTypes sym ->
  TypeRepr tp ->
  MuxFn (Pred sym) (PartExpr (Pred sym) (RegValue sym tp))
muxPartialRegForType sym iteFns tp =
  mergePartial sym (\c u v -> lift $ muxRegForType sym iteFns tp c u v)

-- | A symbolic mux function for @GlobalContents@.
muxGlobalContents ::
  forall sym .
  IsSymInterface sym =>
  sym ->
  IntrinsicTypes sym ->
  MuxFn (Pred sym) (GlobalContents sym)
muxGlobalContents sym iteFns c (GlobalTable vs1 rs1) (GlobalTable vs2 rs2) =
  do vs' <- MapF.mergeWithKeyM muxEntry checkNullMap checkNullMap vs1 vs2
     rs' <- MapF.mergeWithKeyM muxRef refLeft refRight rs1 rs2
     return (GlobalTable vs' rs')
  where
    muxEntry :: GlobalVar tp
             -> GlobalEntry sym tp
             -> GlobalEntry sym tp
             -> IO (Maybe (GlobalEntry sym tp))
    muxEntry g (GlobalEntry u) (GlobalEntry v) =
      Just . GlobalEntry <$> muxRegForType sym iteFns (globalType g) c u v

    muxRef :: RefCell tp
           -> RefCellContents sym tp
           -> RefCellContents sym tp
           -> IO (Maybe (RefCellContents sym tp))
    muxRef r (RefCellContents pu u) (RefCellContents pv v) =
      do uv <- muxRegForType sym iteFns (refType r) c u v
         p <- itePred sym c pu pv
         return . Just $ RefCellContents p uv

    -- Make a partial value undefined unless the given predicate holds.
    restrictRefCellContents :: Pred sym -> RefCellContents sym tp -> IO (RefCellContents sym tp)
    restrictRefCellContents p1 (RefCellContents p2 x) =
      do p' <- andPred sym p1 p2
         return (RefCellContents p' x)

    refLeft :: MapF.MapF RefCell (RefCellContents sym) -> IO (MapF.MapF RefCell (RefCellContents sym))
    refLeft m = traverseF (restrictRefCellContents c) m

    refRight :: MapF.MapF RefCell (RefCellContents sym) -> IO (MapF.MapF RefCell (RefCellContents sym))
    refRight m =
      do cnot <- notPred sym c
         traverseF (restrictRefCellContents cnot) m

    -- Sets of global variables are required to be the same in both branches.
    checkNullMap :: MapF.MapF GlobalVar (GlobalEntry sym)
                 -> IO (MapF.MapF GlobalVar (GlobalEntry sym))
    checkNullMap m
      | MapF.null m = return m
      | otherwise =
        panic "GlobalState.globalMuxFn"
                [ "Different global variables in each branch." ]

data EitherOrBoth f g (tp :: k) =
  JustLeft (f tp) | Both (f tp) (g tp) | JustRight (g tp)

-- | A symbolic mux function for @GlobalUpdates@. For cases where a
-- pre-existing value is updated only on one side, we require a
-- @GlobalContents@ to look up the previous value to use in place of
-- the missing update.
muxGlobalUpdates ::
  forall sym .
  IsSymInterface sym =>
  sym ->
  IntrinsicTypes sym ->
  GlobalContents sym ->
  MuxFn (Pred sym) (GlobalUpdates sym)
muxGlobalUpdates sym iteFns (GlobalTable vs0 rs0) c (GlobalTable vs1 rs1) (GlobalTable vs2 rs2) =
  do -- Zip together the two maps of globals.
     vs3 <- MapF.mergeWithKeyM
            (\_ x y -> return (Just (Both x y)))
            (traverseF (return . JustLeft))
            (traverseF (return . JustRight))
            vs1 vs2
     vs' <- MapF.mergeWithKeyM
            (\k x0 e ->
              Just <$>
              case e of
                JustLeft x1 -> muxEntry k x1 x0 -- use old value x0 for right side
                JustRight x2 -> muxEntry k x0 x2 -- use old value x0 for left side
                Both x1 x2 -> muxEntry k x1 x2)
            (\_ -> return MapF.empty) -- old values updated on neither side are excluded from merged updates
            (MapF.traverseWithKey $ \k e ->
              case e of
                JustLeft _ -> panicNull -- panic if there is no old value
                JustRight _ -> panicNull -- panic if there is no old value
                Both x1 x2 -> muxEntry k x1 x2)
            vs0
            vs3

     -- Zip together the two maps of references.
     rs3 <- MapF.mergeWithKeyM
            (\_ x y -> return (Just (Both x y)))
            (traverseF (return . JustLeft))
            (traverseF (return . JustRight))
            rs1 rs2
     rs' <- MapF.mergeWithKeyM
            (\k (RefCellContents p e) eb ->
              let x0 = RefCellUpdate (PE p e) in
              Just <$>
              case eb of
                JustLeft x1 -> muxRef k x1 x0 -- use old value x0 for right side
                JustRight x2 -> muxRef k x0 x2 -- use old value x0 for left side
                Both x1 x2 -> muxRef k x1 x2)
            (\_ -> return MapF.empty) -- old values updated on neither side are excluded from merged updates
            (MapF.traverseWithKey $ \k e ->
              case e of
                JustLeft x1 -> muxRef k x1 undef -- x1 was newly-created on left side, undefined on right
                JustRight x2 -> muxRef k undef x2 -- x2 was newly-created on right side, undefined on left
                Both x1 x2 -> muxRef k x1 x2)
            rs0
            rs3
     return (GlobalTable vs' rs')
  where
    undef :: forall tp. RefCellUpdate sym tp
    undef = RefCellUpdate Unassigned

    muxEntry :: GlobalVar tp
             -> GlobalEntry sym tp
             -> GlobalEntry sym tp
             -> IO (GlobalEntry sym tp)
    muxEntry g (GlobalEntry u) (GlobalEntry v) =
      GlobalEntry <$> muxRegForType sym iteFns (globalType g) c u v

    muxRef :: RefCell tp
           -> RefCellUpdate sym tp
           -> RefCellUpdate sym tp
           -> IO (RefCellUpdate sym tp)
    muxRef r (RefCellUpdate pe1) (RefCellUpdate pe2) =
      RefCellUpdate <$> muxPartialRegForType sym iteFns (refType r) c pe1 pe2

    panicNull =
      panic "GlobalState.globalMuxFn"
            [ "Different global variables in each branch." ]

-- | Compute a symbolic if-then-else on two global states. The
-- function assumes that the two states were identical up until the
-- most recent branch point marked by 'globalPushBranch'. This most
-- recent branch point marker is also popped from the stack.
globalMuxFn ::
  forall sym .
  IsSymInterface sym =>
  sym ->
  IntrinsicTypes sym ->
  MuxFn (Pred sym) (SymGlobalState sym)

globalMuxFn sym iteFns cond
  (GlobalState (BranchFrame u1 cache1 gf1) s1)
  (GlobalState (BranchFrame u2 _cache2 gf2) s2)
  | globalPendingBranches gf1 == globalPendingBranches gf2 =
    -- We assume gf1 is in fact equal to gf2, which should be the case
    -- if we've followed the appropriate branching discipline.
    do u3 <- muxGlobalUpdates sym iteFns cache1 cond u1 u2
       s3 <- muxGlobalContents sym iteFns cond s1 s2
       --let gf' = updateFrame (mergeGlobalUpdates x') gf1
       let gf3 =
             case gf1 of
               InitialFrame c' -> InitialFrame (applyGlobalUpdates u3 c')
               BranchFrame u' c' gf' -> BranchFrame (mergeGlobalTable u3 u') c' gf'
       return (GlobalState gf3 s3)

globalMuxFn sym _ _ (GlobalState gf1 _) (GlobalState gf2 _) =
  do loc <- getCurrentProgramLoc sym
     panic "GlobalState.globalMuxFn"
           [ "Attempting to merge global states of incorrect branch depths:"
           , " *** Depth 1:  " ++ show (globalPendingBranches gf1)
           , " *** Depth 2:  " ++ show (globalPendingBranches gf2)
           , " *** Location: " ++ show (plSourceLoc loc)
           ]