content: -----------------------------------------------------------------------
  -- |
  -- Module           : Lang.Crucible.Analysis.Postdom
  -- Description      : Populates postdominator entries in CFG blocks.
  -- Copyright        : (c) Galois, Inc 2014
  -- License          : BSD3
  -- Maintainer       : Joe Hendrix <jhendrix@galois.com>
  -- Stability        : provisional
  --
  -- This module provides a method for populating the postdominator fields
  -- in blocks of a Core SSA-form CFG.
  ------------------------------------------------------------------------
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TupleSections #-}
  module Lang.Crucible.Analysis.Postdom
    ( postdomInfo
    , breakpointPostdomInfo
    , validatePostdom
    ) where
  
  import           Control.Monad.State
  import qualified Data.Bimap as Bimap
  import           Data.Functor.Const
  import qualified Data.Graph.Inductive as G
  import           Data.Map (Map)
  import qualified Data.Map as Map
  import           Data.Maybe
  import qualified Data.Parameterized.Context as Ctx
  import           Data.Parameterized.TraversableFC
  import qualified Data.Set as Set
  
  import           Lang.Crucible.CFG.Core
  
  -- | Convert a block ID to a node
  toNode :: BlockID blocks ctx -> G.Node
  toNode (BlockID b) = 1 + Ctx.indexVal b
  
  -- | Create
  reverseEdge :: Int -> Block ext blocks ret ctx -> G.LEdge ()
  reverseEdge d b = (d, toNode (blockID b), ())
  
  -- | For a given block with out edges l, return edges from
  -- each block in @l@ to @b@.
  inEdges :: Block ext blocks ret ctx -> [G.LEdge ()]
  inEdges b =
    case withBlockTermStmt b (\_ -> termStmtNextBlocks) of
      Nothing -> [reverseEdge 0 b]
      Just l -> (\(Some n) -> toNode n `reverseEdge` b) <$> l
  
  inEdgeGraph :: BlockMap ext blocks ret -> [Some (BlockID blocks)] -> G.UGr
  inEdgeGraph m breakpointIds = G.mkGraph ((,()) <$> nodes) edges
    where nodes = 0 : toListFC (toNode . blockID) m
          cfgEdges = foldMapFC inEdges m
          breakpointEdges = map (\(Some bid) -> reverseEdge 0 (getBlock bid m))
                                breakpointIds
          edges = cfgEdges ++ breakpointEdges
  
  -- | Return subgraph of nodes reachable from given node.
  reachableSubgraph :: G.Node -> G.UGr -> G.UGr
  reachableSubgraph initNode g = G.mkGraph nl el
    where reachableNodes = Set.fromList $ G.bfs initNode g
          keepNode = (`Set.member` reachableNodes)
          nl = filter (\(n,_) -> keepNode n) (G.labNodes g)
  
          keepEdge (s,e,_) = keepNode s && keepNode e
          el = filter keepEdge (G.labEdges g)
  
  nodeToBlockIDMap :: BlockMap ext blocks ret
                   -> Map G.Node (Some (BlockID blocks))
  nodeToBlockIDMap =
    foldrFC (\b -> Map.insert (toNode (blockID b)) (Some (blockID b)))
            Map.empty
  
  postdomMap :: forall ext blocks ret
              . BlockMap ext blocks ret
             -> [Some (BlockID blocks)]
             -> Map (Some (BlockID blocks)) [Some (BlockID blocks)]
  postdomMap m breakpointIds = r
    where g0 = inEdgeGraph m breakpointIds
          g = reachableSubgraph 0 g0
  
          idMap = nodeToBlockIDMap m
          f :: Int -> Maybe (Some (BlockID blocks))
          f 0 = Nothing
          f i = Map.lookup i idMap
          -- Map each block to the postdominator for the block.
          r = Map.fromList
            [ (pd_id, mapMaybe f l)
            | (pd,_:l)  <- G.dom g 0
            , pd > 0
            -- Get the post dominator blkID, using a total version of:
            --    let Just pd_id = Map.lookup pd idMap
            , pd_id <- catMaybes [ Map.lookup pd idMap ]
            ]
  
  postdomAssignment :: forall ext blocks ret
                     . BlockMap ext blocks ret
                    -> [Some (BlockID blocks)]
                    -> CFGPostdom blocks
  postdomAssignment m breakpointIds = fmapFC go m
    where pd = postdomMap m breakpointIds
          go :: Block ext blocks ret c -> Const [Some (BlockID blocks)] c
          go b = Const $ fromMaybe [] (Map.lookup (Some (blockID b)) pd)
  
  -- | Compute posstdom information for CFG.
  postdomInfo :: CFG ext b i r -> CFGPostdom b
  postdomInfo g = postdomAssignment (cfgBlockMap g) []
  
  breakpointPostdomInfo :: CFG ext b i r -> [BreakpointName] -> CFGPostdom b
  breakpointPostdomInfo g breakpointNames = postdomAssignment (cfgBlockMap g) $
    mapMaybe (\nm -> Bimap.lookup nm (cfgBreakpoints g)) breakpointNames
  
  blockEndsWithError :: Block ext blocks ret args -> Bool
  blockEndsWithError b =
    withBlockTermStmt b $ \_ ts ->
      case ts of
        ErrorStmt{} -> True
        _ -> False
  
  addErrorIf :: Bool -> String -> State [String] ()
  addErrorIf True msg = modify $ (msg:)
  addErrorIf False _ = return ()
  
  validateTarget :: CFG ext blocks init ret
                 -> CFGPostdom blocks
                 -> String
                 -- ^ Identifier for error.
                 -> [Some (BlockID blocks)]
                 -- ^ Postdoms for source block.
                 -> Some (BlockID blocks)
                 -- ^ Target
                 -> State [String] ()
  validateTarget _ pdInfo src (Some pd:src_postdoms) (Some tgt)
    | isJust (testEquality pd tgt) =
        addErrorIf (src_postdoms /= tgt_postdoms) $
          "Unexpected postdominators from " ++ src ++ " to " ++ show tgt ++ "."
    where Const tgt_postdoms = pdInfo Ctx.! blockIDIndex tgt
  validateTarget g pdInfo src src_postdoms (Some tgt)
    | blockEndsWithError tgt_block =
      return ()
    | otherwise = do
        let tgt_len = length tgt_postdoms
        let src_len = length src_postdoms
        addErrorIf (tgt_len < src_len) $
          "Unexpected postdominators from " ++ src ++ " to " ++ show tgt ++ "."
        let tgt_prefix = drop (tgt_len - src_len) tgt_postdoms
        addErrorIf (src_postdoms /= tgt_prefix) $
          "Unexpected postdominators from " ++ src ++ " to " ++ show tgt ++ "."
    where tgt_block = getBlock tgt (cfgBlockMap g)
          Const tgt_postdoms = pdInfo Ctx.! blockIDIndex tgt
  
  validatePostdom :: CFG ext blocks init ret
                  -> CFGPostdom blocks
                  -> [String]
  validatePostdom g pdInfo = flip execState [] $ do
    forFC_ (cfgBlockMap g) $ \b -> do
      let Const b_pd = pdInfo Ctx.! blockIDIndex (blockID b)
      let loc = show (cfgHandle g) ++ show (blockID b)
      mapM_ (validateTarget g pdInfo loc b_pd) (nextBlocks b)
  
      withBlockTermStmt b $ \_ ts -> do
        case ts of
          Jump tgt -> do
            validateTarget g pdInfo loc b_pd (jumpTargetID tgt)
          Br _ tgt1 tgt2  -> do
            validateTarget g pdInfo loc b_pd (jumpTargetID tgt1)
            validateTarget g pdInfo loc b_pd (jumpTargetID tgt2)
          MaybeBranch _ _ x y -> do
            validateTarget g pdInfo loc b_pd (switchTargetID x)
            validateTarget g pdInfo loc b_pd (jumpTargetID   y)
          VariantElim _ _ s -> do
            traverseFC_ (validateTarget g pdInfo loc b_pd . switchTargetID) s
          Return{} -> do
            addErrorIf (not (null b_pd)) $
              "Expected empty postdom in " ++ loc ++ "."
          TailCall{} -> do
            addErrorIf (not (null b_pd)) $
              "Expected empty postdom in " ++ loc ++ "."
          ErrorStmt{} -> do
            addErrorIf (not (null b_pd)) $
              "Expected empty postdom in " ++ loc ++ "."
  