crucible-0.7: src/Lang/Crucible/Utils/RegRewrite.hs
-----------------------------------------------------------------------
-- |
-- 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