content: ------------------------------------------------------------------------
  -- |
  -- Module           : Lang.Crucible.Utils.CoreRewrite
  -- Description      : Operations for manipulating Core CFGs
  -- Copyright        : (c) Galois, Inc 2016
  -- License          : BSD3
  -- Maintainer       : Simon Winwood <sjw@galois.com>
  -- Stability        : provisional
  --
  ------------------------------------------------------------------------
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE InstanceSigs #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE RankNTypes #-}
  
  module Lang.Crucible.Utils.CoreRewrite
  ( annotateCFGStmts
  ) where
  
  import           Control.Lens
  
  import qualified Data.Parameterized.Context as Ctx
  import           Data.Parameterized.Map (Pair(..))
  import           Data.Parameterized.TraversableFC
  
  import           Lang.Crucible.CFG.Core
  import           Lang.Crucible.CFG.Extension
  
  ------------------------------------------------------------------------
  -- CFG annotation
  
  
  -- | This function walks through all the blocks in the CFG calling
  -- @fS@ on each @Stmt@ and @fT@ on each @TermStmt@.  These functions
  -- return a possible annotaition statement (which has access to the
  -- result of the statement, if any) along with a context diff which
  -- describes any new variables.
  annotateCFGStmts ::
     TraverseExt ext =>
     (forall cin cout. Some (BlockID blocks) -> Ctx.Size cout -> Stmt ext cin cout -> Maybe (StmtSeq ext blocks UnitType cout))
    -- ^ This is the annotation function.  The resulting @StmtSeq@ gets
    -- spliced in after the statement so that they can inspect the
    -- result if desired.  The terminal statement is ignored.
    -> (forall ctx'. Some (BlockID blocks)  -> Ctx.Size ctx' -> TermStmt blocks ret ctx' -> Maybe (StmtSeq ext blocks UnitType ctx'))
    -- ^ As above but for the final term stmt, where the annotation will
    -- be _before_ the term stmt.
    -> CFG ext blocks ctx ret -> CFG ext blocks ctx ret
  annotateCFGStmts fS fT = mapCFGBlocks (annotateBlockStmts fS fT)
  
  mapCFGBlocks :: (forall x. Block ext blocks ret x -> Block ext blocks ret x)
               -> CFG ext blocks ctx ret -> CFG ext blocks ctx ret
  mapCFGBlocks f cfg = cfg { cfgBlockMap = fmapFC f (cfgBlockMap cfg) }
  
  annotateBlockStmts ::
    forall ext blocks ret ctx.
    TraverseExt ext =>
    (forall cin cout. Some (BlockID blocks) -> Ctx.Size cout -> Stmt ext cin cout -> Maybe (StmtSeq ext blocks UnitType cout))
    -- ^ This is the annotation function.  Annotation statements go
    -- after the statement so that they can inspect the result if
    -- desired.  We use Diff here over CtxEmbedding as the remainder of
    -- the statements can't use the result of the annotation function
    -> (forall ctx'. Some (BlockID blocks) -> Ctx.Size ctx' -> TermStmt blocks ret ctx' -> Maybe (StmtSeq ext blocks UnitType ctx'))
    -- ^ As above but for the final term stmt, where the annotation will
    -- be _before_ the term stmt.
    -> Block ext blocks ret ctx
    -> Block ext blocks ret ctx
  annotateBlockStmts fS fT b = b & blockStmts %~ goStmts initialCtxe
    where
      initialCtxe = Ctx.identityEmbedding (Ctx.size (blockInputs b))
      goStmts :: forall ctx' ctx''. Ctx.CtxEmbedding ctx' ctx''
              -> StmtSeq ext blocks ret ctx' -> StmtSeq ext blocks ret ctx''
      goStmts ctxe (ConsStmt loc stmt rest) =
        case applyEmbeddingStmt ctxe stmt of
          Pair stmt' ctxe' ->
            case fS (Some $ blockID b) (ctxe' ^. Ctx.ctxeSize) stmt' of
              Nothing  -> ConsStmt loc stmt' (goStmts ctxe' rest)
              Just annotSeq ->
                ConsStmt loc stmt' (appendStmtSeq ctxe' annotSeq (flip goStmts rest))
      goStmts ctxe (TermStmt loc term) =
        let term' = Ctx.applyEmbedding ctxe term in
        case fT (Some $ blockID b) (ctxe ^. Ctx.ctxeSize) term' of
          Nothing -> TermStmt loc term'
          Just annotSeq ->
            -- FIXME: we could use extendContext here instead
            let restf :: forall fctx. Ctx.CtxEmbedding ctx' fctx -> StmtSeq ext blocks ret fctx
                restf ctxe'' = TermStmt loc (Ctx.applyEmbedding ctxe'' term)
            in appendStmtSeq ctxe annotSeq restf
  
  stmtDiff :: Stmt ext ctx ctx' -> Ctx.Diff ctx ctx'
  stmtDiff stmt =
    case stmt of
      SetReg {}        -> Ctx.knownDiff
      ExtendAssign{}   -> Ctx.knownDiff
      CallHandle {}    -> Ctx.knownDiff
      Print {}         -> Ctx.knownDiff
      ReadGlobal {}    -> Ctx.knownDiff
      WriteGlobal {}   -> Ctx.knownDiff
      FreshConstant{}  -> Ctx.knownDiff
      FreshFloat{}     -> Ctx.knownDiff
      FreshNat{}       -> Ctx.knownDiff
      NewRefCell {}    -> Ctx.knownDiff
      NewEmptyRefCell{}-> Ctx.knownDiff
      ReadRefCell {}   -> Ctx.knownDiff
      WriteRefCell {}  -> Ctx.knownDiff
      DropRefCell {}   -> Ctx.knownDiff
      Assert {}        -> Ctx.knownDiff
      Assume {}        -> Ctx.knownDiff
  
  -- | This appends two @StmtSeq@, throwing away the @TermStmt@ from the first @StmtSeq@
  -- It could probably be generalized to @Ctx.Diff@ instead of an embedding.
  appendStmtSeq :: forall ext blocks ret ret' ctx ctx'.
                   Ctx.CtxEmbedding ctx ctx'
                -> StmtSeq ext blocks ret  ctx'
                -> (forall ctx''. Ctx.CtxEmbedding ctx ctx'' -> StmtSeq ext blocks ret' ctx'')
                -> StmtSeq ext blocks ret' ctx'
  appendStmtSeq ctxe seq1 seq2f = go ctxe seq1
    where
      go :: forall ctx''.
            Ctx.CtxEmbedding ctx ctx''
            -> StmtSeq ext blocks ret ctx''
            -> StmtSeq ext blocks ret' ctx''
      go ctxe' (ConsStmt loc stmt rest) =
        -- This just throws away the new variables, which is OK as seq2
        -- can't reference them.
        let ctxe'' = Ctx.extendEmbeddingRightDiff (stmtDiff stmt) ctxe'
        in ConsStmt loc stmt (go ctxe'' rest)
      go ctxe' (TermStmt _loc _term)    = seq2f ctxe'
  