-----------------------------------------------------------------------
-- |
-- Module           : Lang.Crucible.Utils.RegRewrite
-- Description      : Operations for manipulating registerized CFGs
-- Copyright        : (c) Galois, Inc 2014-2018
-- License          : BSD3
-- Maintainer       : Luke Maurer <lukemaurer@galois.com>
-- Stability        : provisional
--
-- A rewrite engine for registerized CFGs.
------------------------------------------------------------------------
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module Lang.Crucible.Utils.RegRewrite
  ( -- * Main interface
    annotateCFGStmts
    -- * Annotation monad
  , Rewriter
  , addStmt
  , addInternalStmt
  , ifte
  , freshAtom
  ) where

import           Control.Monad.RWS.Strict
import           Control.Monad.State.Strict ( StateT, evalStateT )
import           Control.Monad.ST ( ST, runST )
import           Data.Foldable ( toList )
import           Data.Parameterized.Map ( MapF )
import qualified Data.Parameterized.Map as MapF
import           Data.Parameterized.Nonce ( Nonce, NonceGenerator, freshNonce
                                          , newSTNonceGenerator )
import           Data.Parameterized.Some ( Some(Some) )
import           Data.Sequence ( Seq )
import qualified Data.Sequence as Seq
import qualified Data.Set as Set

import           What4.ProgramLoc

import           Lang.Crucible.CFG.Extension
import           Lang.Crucible.CFG.Reg
import           Lang.Crucible.Types

------------------------------------------------------------------------
-- Public interface

-- | Add statements to each block in a CFG according to the given
-- instrumentation functions. See the 'Rewriter' monad for the
-- operations provided for adding code.
annotateCFGStmts :: TraverseExt ext
                 => u
                 -- ^ Initial user state
                 -> (forall s h. Posd (Stmt ext s) -> Rewriter ext h s ret u ())
                 -- ^ Action to run on each non-terminating statement;
                 -- must explicitly add the original statement back if
                 -- desired
                 -> (forall s h. Posd (TermStmt s ret) -> Rewriter ext h s ret u ())
                 -- ^ Action to run on each terminating statement
                 -> SomeCFG ext init ret
                 -- ^ Graph to rewrite
                 -> SomeCFG ext init ret
annotateCFGStmts u fS fT (SomeCFG cfg) =
  runRewriter u $
    do cfg1 <- renameAll cfg
       blocks' <- mapM (annotateBlockStmts fS fT) (cfgBlocks cfg1)
       SomeCFG <$> newCFG cfg1 (concat blocks')

-- | Monad providing operations for modifying a basic block by adding
-- statements and/or splicing in conditional braches. Also provides a
-- 'MonadState' instance for storing user state.
newtype Rewriter ext h s (ret :: CrucibleType) u a =
  Rewriter (RWST (NonceGenerator (ST h) s)
                 (Seq (ComplexStmt ext s))
                 u (ST h) a)
  deriving ( Functor, Applicative, Monad, MonadState u
           , MonadWriter (Seq (ComplexStmt ext s))
           )

-- | Add a new statement at the current position.
addStmt :: Posd (Stmt ext s) -> Rewriter ext h s ret u ()
addStmt stmt = tell (Seq.singleton (Stmt stmt))

-- | Add a new statement at the current position, marking it as
-- internally generated.
addInternalStmt :: Stmt ext s -> Rewriter ext h s ret u ()
addInternalStmt = addStmt . Posd InternalPos

-- | Add a conditional at the current position. This will cause the
-- current block to end and new blocks to be generated for the two
-- branches and the remaining statements in the original block.
ifte :: Atom s BoolType
       -> Rewriter ext h s ret u ()
       -> Rewriter ext h s ret u ()
       -> Rewriter ext h s ret u ()
ifte atom thn els =
  do (~(), thnSeq) <- gather thn
     (~(), elsSeq) <- gather els
     tell $ Seq.singleton (IfThenElse atom thnSeq elsSeq)

-- | Create a new atom with a freshly allocated id. The id will not
-- have been used anywhere in the original CFG.
freshAtom :: TypeRepr tp -> Rewriter ext h s ret u (Atom s tp)
freshAtom tp =
  do ng <- Rewriter $ ask
     n <- Rewriter $ lift $ freshNonce ng
     return $ Atom { atomPosition = InternalPos
                   , atomId = n
                   , atomSource = Assigned
                   , typeOfAtom = tp }

------------------------------------------------------------------------
-- Monad
--
-- For each block, rewriting occurs in two stages:
--
-- 1. Generate a sequence of "complex statements", each of which may
--    be an internal if-then-else.
-- 2. Rebuild the block from the complex statements, creating
--    additional blocks for internal control flow.
--
-- Step 1 occurs through a simple writer monad, leaving the nasty details
-- of block mangling to step 2.

data ComplexStmt ext s
  = Stmt (Posd (Stmt ext s))
  | IfThenElse (Atom s BoolType)
               (Seq (ComplexStmt ext s))
               (Seq (ComplexStmt ext s))

runRewriter :: forall u ext ret a
             . u -> (forall h s. Rewriter ext h s ret u a) -> a
runRewriter u m = runST $ do
  Some ng <- newSTNonceGenerator
  case m of
    -- Have to do this pattern match *after* unpacking the Some from
    -- newSTNonceGenerator for obscure reasons involving Skolem
    -- functions
    Rewriter f -> do
      (a, _, _) <- runRWST f ng u
      return a

freshLabel :: forall ext h s ret u. Rewriter ext h s ret u (Label s)
freshLabel =
  do ng <- Rewriter $ ask
     n <- Rewriter $ lift $ freshNonce ng
     return $ Label { labelId = n }

-- | Return the output of a writer action without passing it onward.
gather :: MonadWriter w m => m a -> m (a, w)
gather m = censor (const mempty) $ listen m

------------------------------------------------------------------------
-- Implementation

-- Give fresh names to everything.  The only point of this is that the
-- new names come from a known nonce generator, so we can now generate
-- more names.  We do this in a separate pass up front so that we
-- don't have to juggle two namespaces afterward.
renameAll :: forall s0 s ext init ret h u
           . ( TraverseExt ext )
          => CFG ext s0 init ret
          -> Rewriter ext h s ret u (CFG ext s init ret)
renameAll cfg = do
  ng <- Rewriter $ ask
  Rewriter $ lift $ evalStateT (substCFG (rename ng) cfg) MapF.empty
  where
    rename :: NonceGenerator (ST h) s
           -> Nonce s0 (tp :: CrucibleType)
           -> StateT (MapF @CrucibleType (Nonce s0) (Nonce s)) (ST h) (Nonce s tp)
    rename ng n = do
      mapping <- get
      case MapF.lookup n mapping of
        Just n' ->
          return n'
        Nothing -> do
          n' <- lift $ freshNonce ng
          modify (MapF.insert n n')
          return n'

newCFG :: CFG ext s init ret
       -> [Block ext s ret]
       -> Rewriter ext h s ret u (CFG ext s init ret)
newCFG cfg blocks = do
  return $ cfg { cfgBlocks = blocks }

annotateBlockStmts :: TraverseExt ext
                   => (Posd (Stmt ext s) -> Rewriter ext h s ret u ())
                   -> (Posd (TermStmt s ret) -> Rewriter ext h s ret u ())
                   -> Block ext s ret
                   -> Rewriter ext h s ret u [Block ext s ret]
annotateBlockStmts fS fT block =
  do -- Step 1
     stmts <- annotateAsComplexStmts fS fT block
     -- Step 2
     rebuildBlock stmts block

annotateAsComplexStmts :: (Posd (Stmt ext s) -> Rewriter ext h s ret u ())
                       -> (Posd (TermStmt s ret) -> Rewriter ext h s ret u ())
                       -> Block ext s ret
                       -> Rewriter ext h s ret u (Seq (ComplexStmt ext s))
annotateAsComplexStmts fS fT block =
  do (~(), stmts) <- gather $
       do mapM_ fS (blockStmts block)
          fT (blockTerm block)
     return stmts

rebuildBlock :: TraverseExt ext
             => Seq (ComplexStmt ext s)
             -> Block ext s ret
             -> Rewriter ext h s ret u [Block ext s ret]
rebuildBlock stmts block =
  toList <$> go stmts Seq.empty Seq.empty
     (blockID block) (blockExtraInputs block) (blockTerm block)
  where
    go :: TraverseExt ext
       => Seq (ComplexStmt ext s) -- Statements to process
       -> Seq (Posd (Stmt ext s)) -- Statements added to current block
       -> Seq (Block ext s ret)   -- Blocks created so far
       -> BlockID s               -- Id of current block
       -> ValueSet s              -- Extra inputs to current block
       -> Posd (TermStmt s ret)   -- Terminal statement of current block
       -> Rewriter ext h s ret u (Seq (Block ext s ret))
    go s accStmts accBlocks bid ext term = case s of
      Seq.Empty ->
        return $ accBlocks Seq.|> mkBlock bid ext accStmts term
      (Stmt stmt Seq.:<| s') ->
        go s' (accStmts Seq.|> stmt) accBlocks bid ext term
      (IfThenElse a thn els Seq.:<| s') ->
        do thnLab <- freshLabel
           elsLab <- freshLabel
           newLab <- freshLabel
           -- End the block, terminating with a branch statement
           let branch = Posd InternalPos (Br a thnLab elsLab)
               thisBlock = mkBlock bid ext accStmts branch
           -- Make the branches into (sets of) blocks
           let jump = Posd InternalPos (Jump newLab)
           thnBlocks <-
             go thn Seq.empty Seq.empty (LabelID thnLab) Set.empty jump
           elsBlocks <-
             go els Seq.empty Seq.empty (LabelID elsLab) Set.empty jump
           -- Keep going with a new, currently empty block
           let accBlocks' = (accBlocks Seq.|> thisBlock) Seq.><
                            thnBlocks Seq.>< elsBlocks
           go s' Seq.empty accBlocks' (LabelID newLab) Set.empty term