content: ---------------------------------------------------------------------------
  -- |
  -- Module          : Lang.Crucible.CFG.ExtractSubgraph
  -- Description     : Allows for construction of a function based off a subgraph
  --                   of an SSA-form function, subject to certain constraints
  -- Copyright       : (c) Galois, Inc 2015
  -- License         : BSD3
  --
  ---------------------------------------------------------------------------
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE PatternGuards #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE TypeOperators #-}
  module Lang.Crucible.CFG.ExtractSubgraph
    ( extractSubgraph
    ) where
  
  import           Control.Lens
  import qualified Data.Bimap as Bimap
  import           Data.Parameterized.Context as Ctx
  import           Data.Parameterized.Map as MapF
  import           Data.Set as S
  import qualified Data.Map as Map
  import           Debug.Trace
  
  import           What4.FunctionName
  import           What4.ProgramLoc
  
  import           Lang.Crucible.CFG.Core
  import           Lang.Crucible.FunctionHandle
  
  -- | Given a CFG @cfg@, a set of blocks @cuts@ that take the return type as their sole
  -- argument, and a block @bi@ that takes the CFG's init type as its sole argument,
  -- construct a CFG that is a maximal subgraph starting at @bi@ and not entering any
  -- block in @cuts@.  If the original graph would enter a block in @cuts@, the value
  -- passed to that block is returned.  If @bi `member` cuts@, then whenever the subgraph
  -- would transition to @bi@, it returns the value that would be passed to @bi@ instead.
  extractSubgraph :: (KnownCtx TypeRepr init, KnownRepr TypeRepr ret)
                  => CFG ext blocks init ret
                  -> Set (BlockID blocks (EmptyCtx ::> ret))
                  -> BlockID blocks init
                  -> HandleAllocator
                  -> IO (Maybe (SomeCFG ext init ret))
  extractSubgraph (CFG{cfgBlockMap = orig, cfgBreakpoints = breakpoints}) cuts bi halloc =
    extractSubgraphFirst orig cuts MapF.empty zeroSize bi $
      \(SubgraphIntermediate finalMap finalInitMap _sz entryID cb) -> do
          hn <- mkHandle halloc startFunctionName
          return $ do
            bm <- cb finalMap finalInitMap Ctx.empty
            return $ SomeCFG $ CFG
              { cfgBlockMap = bm
              , cfgEntryBlockID = entryID
              , cfgHandle = hn
              , cfgBreakpoints = Bimap.fromList $ Map.toList $
                  Map.mapMaybe (viewSome $ \bid -> Some <$> MapF.lookup bid finalMap) $
                  Bimap.toMap breakpoints
              }
  
  -- | Type for carrying intermediate results through subraph extraction
  -- the interesting field is the final one - it holds a callback for transforming
  -- the result of the previous portion of the subgraph extraction into the result
  -- of this subgraph extraction.
  data SubgraphIntermediate ext old ret init soFar new where
    SubgraphIntermediate :: MapF (BlockID old) (BlockID new)
                         -> MapF (BlockID old) (BlockID new)
                         -> Size new
                         -> BlockID new init
                         -> (forall all. (MapF (BlockID old) (BlockID all)
                                          -> MapF (BlockID old) (BlockID all)
                                          -> Assignment (Block ext all ret) soFar
                                          -> Maybe (Assignment (Block ext all ret) new)))
                         -> SubgraphIntermediate ext old ret init soFar new
  
  
  -- | The inner loop of subgraph extraction
  --   produces a callback with an existential type, in order to hide new
  extractSubgraph' :: KnownRepr TypeRepr ret
                   => BlockMap ext old ret
                   -> Set (BlockID old (EmptyCtx ::> ret))
                   -> MapF (BlockID old) (BlockID soFar)
                   -> MapF (BlockID old) (BlockID soFar)
                   -> Size soFar
                   -> BlockID old init
                   -> BlockID soFar args
                   -> forall r . (forall new. SubgraphIntermediate ext old ret args soFar new -> r)
                   -> r
  extractSubgraph' orig cuts mapF initMap sz bi ident f =
    let block = getBlock bi orig
    in  withBlockTermStmt block $ (\_ t ->
          (case t of
            Jump (JumpTarget bi' _ _) -> \sgi -> visitChildNode orig cuts bi' sgi f
            Br _ (JumpTarget bi1 _ _) (JumpTarget bi2 _ _) -> \sgi1 ->
              visitChildNode orig cuts bi1 sgi1
                $ \sgi2 -> visitChildNode orig cuts bi2 sgi2 f
            Return _ -> f
            _ -> error "extractSubgraph': unexpected case!")
                  (SubgraphIntermediate
                    (MapF.insert bi (BlockID $ nextIndex sz) (MapF.map extendBlockID mapF))
                    (MapF.map extendBlockID initMap)
                    (incSize sz)
                    (extendBlockID ident)
                    (\finalMap _finalInitMap assn ->
                      fmap (extend assn) (do
                        finalID <- MapF.lookup bi finalMap
                        cloneBlock finalMap finalID block))))
  
  -- code duplication... but the types need to be different between iterations
  -- FIXME: write a generic version that this and extractSubgraph' can be wrappers
  -- around
  extractSubgraphFirst :: KnownRepr TypeRepr ret
                       => BlockMap ext old ret
                       -> Set (BlockID old (EmptyCtx ::> ret))
                       -> MapF (BlockID old) (BlockID soFar)
                       -> Size soFar
                       -> BlockID old init
                       -> forall r . (forall new. SubgraphIntermediate ext old ret init soFar new -> r)
                       -> r
  extractSubgraphFirst orig cuts mapF sz bi f =
    let block = getBlock bi orig
    in  withBlockTermStmt block $ (\_ t ->
          (case t of
            Jump (JumpTarget bi' _ _) -> \sgi -> visitChildNode orig cuts bi' sgi f
            Br _ (JumpTarget bi1 _ _) (JumpTarget bi2 _ _) -> \sgi1 ->
              visitChildNode orig cuts bi1 sgi1
                $ \sgi2 -> visitChildNode orig cuts bi2 sgi2 f
            Return _ -> f
            _ -> error "extractSubgraphFirst: unexpected case!")
                  (SubgraphIntermediate
                    (if case S.minView cuts of
                        Just (bi', _) -> case testEquality (blockInputs block) (blockInputs $ orig Ctx.! blockIDIndex bi') of
                          Just Refl -> bi `S.member` cuts
                          Nothing -> False
                        Nothing -> False
                      then MapF.map extendBlockID mapF
                      else MapF.insert bi (BlockID $ nextIndex sz) (MapF.map extendBlockID mapF))
                    (MapF.insert bi (BlockID $ nextIndex sz) (MapF.map extendBlockID mapF))
                    (incSize sz)
                    (BlockID $ nextIndex sz)
                    (\finalMap finalInitMap assn -> fmap (extend assn) (do
                        finalID <- MapF.lookup bi finalInitMap
                        cloneBlock finalMap finalID block))))
  
  -- does the building of a new node - mutually recursive with exrtactSubgraph'
  visitChildNode :: KnownRepr TypeRepr ret
                 => BlockMap ext old ret
                 -> Set (BlockID old (EmptyCtx ::> ret))
                 -> BlockID old init
                 -> SubgraphIntermediate ext old ret args soFar prev
                 -> (forall r. (forall new . SubgraphIntermediate ext old ret args soFar new -> r)
                 -> r)
  visitChildNode orig cuts bi (SubgraphIntermediate sgMap initMap sz ident cb) f=
    case MapF.lookup bi sgMap of
      Just _bi' -> f $ SubgraphIntermediate sgMap initMap sz ident cb
      Nothing -> case S.minView cuts of
        Just (cut, _)
          | Just Refl <- testEquality (blockInputs $ orig Ctx.! blockIDIndex bi) (blockInputs $ orig Ctx.! blockIDIndex cut)
          , S.member bi cuts ->
              f $ SubgraphIntermediate
                (MapF.insert bi (BlockID $ nextIndex sz) (MapF.map extendBlockID sgMap))
                (MapF.map extendBlockID initMap)
                (incSize sz)
                (extendBlockID ident)
                (\finalMap finalCutMap assn -> do
                  assn' <- cb finalMap finalCutMap assn
                  newBlock <- mkRetBlock finalMap orig bi
                  return $ extend assn' newBlock)
        _ -> extractSubgraph' orig cuts sgMap initMap sz bi ident
              (\ (SubgraphIntermediate sgMap' initMap' sz' ident' ccb) ->
                f $ SubgraphIntermediate sgMap' initMap' sz' ident'
                  (\finalMap finalCutMap assn ->
                    ccb finalMap finalCutMap  =<< cb finalMap finalCutMap assn))
  
  
  mkRetBlock :: MapF (BlockID old) (BlockID new)
             -> BlockMap ext old ret
             -> BlockID old (EmptyCtx ::> ret)
             -> Maybe (Block ext new ret (EmptyCtx ::> ret))
  mkRetBlock mapF bm ident =
    case MapF.lookup ident mapF of
      Just id' ->
        let block = bm Ctx.! blockIDIndex ident
        in Just $
             let name = plFunction (blockLoc block)
                 term = Return lastReg
                in Block{ blockID       = id'
                        , blockInputs   = blockInputs block
                        , _blockStmts   = TermStmt (mkProgramLoc name InternalPos) term
                        }
      Nothing -> trace ("could not lookup return block id " ++ show (blockIDIndex ident)) Nothing
  
  
  cloneBlock :: MapF (BlockID old) (BlockID new)
             -> BlockID new ctx -> Block ext old ret ctx -> Maybe (Block ext new ret ctx)
  cloneBlock mapF newID b = do
    stmts' <- cloneStmtSeq mapF (b^.blockStmts)
    return Block{ blockID       = newID
                , blockInputs   = blockInputs b
                , _blockStmts   = stmts'
                }
  
  cloneStmtSeq :: MapF (BlockID old) (BlockID new) -> StmtSeq ext old ret ctx -> Maybe (StmtSeq ext new ret ctx)
  cloneStmtSeq mapF (ConsStmt loc stmt rest) = do
    rest' <- cloneStmtSeq mapF rest
    return $ ConsStmt loc stmt rest'
  cloneStmtSeq mapF (TermStmt loc term) = do
    term' <- cloneTerm mapF term
    return $ TermStmt loc term'
  
  cloneTerm :: MapF (BlockID old) (BlockID new) -> TermStmt old ret ctx -> Maybe (TermStmt new ret ctx)
  cloneTerm mapF (Jump jt) = fmap Jump $ cloneJumpTarget mapF jt
  cloneTerm mapF (Br reg jt1 jt2) = do
    jt1' <- cloneJumpTarget mapF jt1
    jt2' <- cloneJumpTarget mapF jt2
    return $ Br reg jt1' jt2'
  cloneTerm _mapF (Return reg) = Just $ Return reg
  cloneTerm _ _ = error "cloneTerm: unexpected case!"
  
  cloneJumpTarget :: MapF (BlockID blocks1) (BlockID blocks2)
                  -> JumpTarget blocks1 t
                  -> Maybe (JumpTarget blocks2 t)
  cloneJumpTarget mapF (JumpTarget ident args assn) = do
    case MapF.lookup ident mapF of
      Just id' -> Just $ JumpTarget id' args assn
      Nothing -> trace ("could not lookup jump target id " ++ show (blockIDIndex ident)) Nothing
  