{-# LANGUAGE PartialTypeSignatures, MultiParamTypeClasses, DeriveGeneric, DefaultSignatures, FlexibleContexts, StrictData #-}

{-|
Module      : Propagation
Description : A generic abstract interpretation algorithm for propagating postcondition-transformations through a control flow graph.

We assume a class where we can do predicate transformation through function @tau@,
and we can merge two predicates through function @join@.
Moreover, we assume an implementation of a function @implies@ that implements symbolic implication.
Given these functions, we provide a generic abstract interpretation algorithm.
-}



module Generic.SymbolicPropagation (
  Propagator(..),
  do_prop
 ) where

import Base
import Analysis.ControlFlow
import Analysis.Context
import qualified Data.Map as M
import qualified Data.Set as S
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Data.Foldable (find)

import Control.Monad.State.Strict hiding (join)
import Debug.Trace
import qualified X86.Instruction as X86

import Generic.SymbolicConstituents 


-- | A class that allows propagation of predicates over a CFG.
class (Show pred) => Propagator ctxt pred where
  -- | Predicate transformation for an edge in in a CFG, over a basic blocks.
  tau     :: ctxt -> [X86.Instruction] -> Maybe [X86.Instruction] -> pred -> (pred,S.Set VerificationCondition)
  -- | A lattice-join
  join    :: ctxt -> String -> pred -> pred -> pred
  -- | Symbolic implication
  implies :: ctxt -> pred -> pred -> Bool



-- The set of edges starting in $v$
out_edges :: CFG -> Int -> Set (Int, Int)
out_edges CFG
g Int
v = [(Int, Int)] -> Set (Int, Int)
forall a. Ord a => [a] -> Set a
S.fromList ([(Int, Int)] -> Set (Int, Int)) -> [(Int, Int)] -> Set (Int, Int)
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Int -> [Int]
forall a. a -> [a]
repeat Int
v) ([Int] -> [(Int, Int)]) -> [Int] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ IntSet -> [Int]
IS.toList (IntSet -> [Int]) -> IntSet -> [Int]
forall a b. (a -> b) -> a -> b
$ CFG -> Int -> IntSet
post CFG
g Int
v



propstate_m :: (a, b, c) -> a
propstate_m   (a
m,b
bag,c
vcs) = a
m
propstate_bag :: (a, b, c) -> b
propstate_bag (a
m,b
bag,c
vcs) = b
bag
propstate_vcs :: (a, b, c) -> c
propstate_vcs (a
m,b
bag,c
vcs) = c
vcs

-- pick edges from the bag, preferring edges to already visited nodes
pick_edge_from_bag :: State (IM.IntMap pred, S.Set (Int,Int), S.Set VerificationCondition) (Maybe ((Int,Int), S.Set (Int,Int)))
pick_edge_from_bag :: State
  (IntMap pred, Set (Int, Int), Set VerificationCondition)
  (Maybe ((Int, Int), Set (Int, Int)))
pick_edge_from_bag = do
  (IntMap pred
m,Set (Int, Int)
bag,Set VerificationCondition
_) <- StateT
  (IntMap pred, Set (Int, Int), Set VerificationCondition)
  Identity
  (IntMap pred, Set (Int, Int), Set VerificationCondition)
forall s (m :: * -> *). MonadState s m => m s
get
  case ((Int, Int) -> Bool) -> Set (Int, Int) -> Maybe (Int, Int)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\(Int
v0,Int
v1) -> Int -> IntMap pred -> Bool
forall a. Int -> IntMap a -> Bool
IM.member Int
v1 IntMap pred
m) Set (Int, Int)
bag of
    Just (Int, Int)
edge -> Maybe ((Int, Int), Set (Int, Int))
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition)
     (Maybe ((Int, Int), Set (Int, Int)))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ((Int, Int), Set (Int, Int))
 -> State
      (IntMap pred, Set (Int, Int), Set VerificationCondition)
      (Maybe ((Int, Int), Set (Int, Int))))
-> Maybe ((Int, Int), Set (Int, Int))
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition)
     (Maybe ((Int, Int), Set (Int, Int)))
forall a b. (a -> b) -> a -> b
$ ((Int, Int), Set (Int, Int)) -> Maybe ((Int, Int), Set (Int, Int))
forall a. a -> Maybe a
Just ((Int, Int)
edge, (Int, Int) -> Set (Int, Int) -> Set (Int, Int)
forall a. Ord a => a -> Set a -> Set a
S.delete (Int, Int)
edge Set (Int, Int)
bag)
    Maybe (Int, Int)
Nothing   -> Maybe ((Int, Int), Set (Int, Int))
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition)
     (Maybe ((Int, Int), Set (Int, Int)))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ((Int, Int), Set (Int, Int))
 -> State
      (IntMap pred, Set (Int, Int), Set VerificationCondition)
      (Maybe ((Int, Int), Set (Int, Int))))
-> Maybe ((Int, Int), Set (Int, Int))
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition)
     (Maybe ((Int, Int), Set (Int, Int)))
forall a b. (a -> b) -> a -> b
$ Set (Int, Int) -> Maybe ((Int, Int), Set (Int, Int))
forall a. Set a -> Maybe (a, Set a)
S.minView Set (Int, Int)
bag
  

-- propagation
-- The state consists of 
-- 		1.) the current mapping of addresses to predicates 
-- 		2.) the current bag (a set of edges to be explored)
prop :: Propagator ctxt pred => ctxt -> CFG -> State (IM.IntMap pred, S.Set (Int,Int), S.Set VerificationCondition) ()
prop :: ctxt
-> CFG
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition) ()
prop ctxt
ctxt CFG
g = do
  Maybe ((Int, Int), Set (Int, Int))
pick <- State
  (IntMap pred, Set (Int, Int), Set VerificationCondition)
  (Maybe ((Int, Int), Set (Int, Int)))
forall pred.
State
  (IntMap pred, Set (Int, Int), Set VerificationCondition)
  (Maybe ((Int, Int), Set (Int, Int)))
pick_edge_from_bag
  case Maybe ((Int, Int), Set (Int, Int))
pick of
    Maybe ((Int, Int), Set (Int, Int))
Nothing ->
      ()
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition) ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Just ((Int
v0,Int
v1),Set (Int, Int)
bag') -> do
      -- take an edge (v0,v1) out of the bag
      IntMap pred
m <- ((IntMap pred, Set (Int, Int), Set VerificationCondition)
 -> IntMap pred)
-> StateT
     (IntMap pred, Set (Int, Int), Set VerificationCondition)
     Identity
     (IntMap pred)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (IntMap pred, Set (Int, Int), Set VerificationCondition)
-> IntMap pred
forall a b c. (a, b, c) -> a
propstate_m
      -- do predicate transformation on the currently available precondition of v0
      let p :: pred
p = [Char] -> IntMap pred -> Int -> pred
forall p. [Char] -> IntMap p -> Int -> p
im_lookup [Char]
"v0 must have predicate" IntMap pred
m Int
v0
      -- this produces q: the precondition for v1
      let (pred
q,Set VerificationCondition
vcs') = ctxt
-> [Instruction]
-> Maybe [Instruction]
-> pred
-> (pred, Set VerificationCondition)
forall ctxt pred.
Propagator ctxt pred =>
ctxt
-> [Instruction]
-> Maybe [Instruction]
-> pred
-> (pred, Set VerificationCondition)
tau ctxt
ctxt (CFG -> Int -> [Instruction]
fetch_block CFG
g Int
v0) ([Instruction] -> Maybe [Instruction]
forall a. a -> Maybe a
Just ([Instruction] -> Maybe [Instruction])
-> [Instruction] -> Maybe [Instruction]
forall a b. (a -> b) -> a -> b
$ CFG -> Int -> [Instruction]
fetch_block CFG
g Int
v1) pred
p
      -- add verification conditions
      ((IntMap pred, Set (Int, Int), Set VerificationCondition)
 -> (IntMap pred, Set (Int, Int), Set VerificationCondition))
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\(IntMap pred
m,Set (Int, Int)
_,Set VerificationCondition
vcs) -> (IntMap pred
m,Set (Int, Int)
bag', Set VerificationCondition
-> Set VerificationCondition -> Set VerificationCondition
forall a. Ord a => Set a -> Set a -> Set a
S.union Set VerificationCondition
vcs Set VerificationCondition
vcs'))
      -- store q
      pred
-> Int
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition) ()
forall (m :: * -> *) a c.
(MonadState (IntMap a, Set (Int, Int), c) m, Propagator ctxt a) =>
a -> Int -> m ()
add_predicate pred
q Int
v1
      -- continue
      ctxt
-> CFG
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition) ()
forall ctxt pred.
Propagator ctxt pred =>
ctxt
-> CFG
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition) ()
prop ctxt
ctxt CFG
g
 where
  add_predicate :: a -> Int -> m ()
add_predicate a
q Int
v1 = do
   (IntMap a
m,Set (Int, Int)
bag,c
vcs) <- m (IntMap a, Set (Int, Int), c)
forall s (m :: * -> *). MonadState s m => m s
get
   case Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
v1 IntMap a
m of
     Maybe a
Nothing -> do
       -- first time visit, store q and explore all outgoing edges
       (IntMap a, Set (Int, Int), c) -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
v1 a
q IntMap a
m, Set (Int, Int) -> Set (Int, Int) -> Set (Int, Int)
forall a. Ord a => Set a -> Set a -> Set a
S.union Set (Int, Int)
bag (Set (Int, Int) -> Set (Int, Int))
-> Set (Int, Int) -> Set (Int, Int)
forall a b. (a -> b) -> a -> b
$ CFG -> Int -> Set (Int, Int)
out_edges CFG
g Int
v1, c
vcs)
     Just a
p -> do
       if ctxt -> a -> a -> Bool
forall ctxt pred.
Propagator ctxt pred =>
ctxt -> pred -> pred -> Bool
implies ctxt
ctxt a
p a
q then
         -- previously visited, no need for further exploration
         () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return () -- put (IM.insert v1 j m, bag, vcs)
       else do
         let j :: a
j = ctxt -> [Char] -> a -> a -> a
forall ctxt pred.
Propagator ctxt pred =>
ctxt -> [Char] -> pred -> pred -> pred
join ctxt
ctxt [Char]
"prop" a
p a
q
         -- previously visited, need to weaken invariant by joining
         (IntMap a, Set (Int, Int), c) -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
v1 a
j IntMap a
m,Set (Int, Int) -> Set (Int, Int) -> Set (Int, Int)
forall a. Ord a => Set a -> Set a -> Set a
S.union Set (Int, Int)
bag (Set (Int, Int) -> Set (Int, Int))
-> Set (Int, Int) -> Set (Int, Int)
forall a b. (a -> b) -> a -> b
$ CFG -> Int -> Set (Int, Int)
out_edges CFG
g Int
v1, c
vcs)

-- | Start propagation at the given entry address with the given initial predicate.
-- Returns a set of invariants, i.e., a mapping of instruction addresses to predicates.
do_prop :: Propagator ctxt pred => 
  ctxt      -- ^ The context
  -> CFG    -- ^ The CFG
  -> Int    -- ^ The entry address
  -> pred   -- ^ The initial predicate
  -> (IM.IntMap pred, S.Set VerificationCondition)
do_prop :: ctxt
-> CFG -> Int -> pred -> (IntMap pred, Set VerificationCondition)
do_prop ctxt
ctxt CFG
g Int
entry pred
p = 
  let (IntMap pred
m,Set (Int, Int)
_,Set VerificationCondition
vcs) = State (IntMap pred, Set (Int, Int), Set VerificationCondition) ()
-> (IntMap pred, Set (Int, Int), Set VerificationCondition)
-> (IntMap pred, Set (Int, Int), Set VerificationCondition)
forall s a. State s a -> s -> s
execState (ctxt
-> CFG
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition) ()
forall ctxt pred.
Propagator ctxt pred =>
ctxt
-> CFG
-> State
     (IntMap pred, Set (Int, Int), Set VerificationCondition) ()
prop ctxt
ctxt CFG
g) ((IntMap pred, Set (Int, Int), Set VerificationCondition)
 -> (IntMap pred, Set (Int, Int), Set VerificationCondition))
-> (IntMap pred, Set (Int, Int), Set VerificationCondition)
-> (IntMap pred, Set (Int, Int), Set VerificationCondition)
forall a b. (a -> b) -> a -> b
$ (Int -> pred -> IntMap pred
forall a. Int -> a -> IntMap a
IM.singleton Int
entry pred
p, CFG -> Int -> Set (Int, Int)
out_edges CFG
g Int
entry, Set VerificationCondition
forall a. Set a
S.empty) in
    (IntMap pred
m,Set VerificationCondition
vcs)