content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  module GHC.Debug.Dominators (computeDominators
                                     , retainerSize
                                     , convertToHeapGraph
                                     , annotateWithRetainerSize ) where
  
  import Data.Maybe       ( catMaybes, fromJust )
  import qualified Data.IntMap as IM
  import qualified Data.IntSet as IS
  import GHC.Debug.Types.Ptr
  import GHC.Debug.Types.Closures
  import qualified Data.List.NonEmpty as NE
  import qualified Data.Foldable as F
  import qualified Data.Graph.Dom as DO
  import qualified Data.Tree as Tree
  import GHC.Debug.Types.Graph
  
  
  
  -- Dominators
  closurePtrToInt :: ClosurePtr -> Int
  closurePtrToInt (ClosurePtr p) = fromIntegral p
  
  intToClosurePtr :: Int -> ClosurePtr
  intToClosurePtr i = mkClosurePtr (fromIntegral i)
  
  convertToDom :: HeapGraph a -> DO.Rooted
  convertToDom  (HeapGraph groots is) = (0, new_graph)
    where
      rootNodes = IS.fromList (map closurePtrToInt (NE.toList groots))
      new_graph = IM.insert 0 rootNodes (IM.foldlWithKey' collectNodes IM.empty is)
      collectNodes newMap k h =  IM.insert k (IS.fromList (map closurePtrToInt (catMaybes (allClosures (hgeClosure h))))) newMap
  
  computeDominators :: HeapGraph a -> [Tree.Tree (HeapGraphEntry a)]
  computeDominators hg = map (fmap (fromJust . flip lookupHeapGraph hg . intToClosurePtr)) gentries
    where
      gentries = case DO.domTree (convertToDom hg) of
                  Tree.Node 0 es -> es
                  _ -> error "Dominator tree must contain 0"
  
  retainerSize :: HeapGraph Size -> [Tree.Tree (HeapGraphEntry (Size, RetainerSize))]
  retainerSize hg = map bottomUpSize doms
    where
      doms = computeDominators hg
  
  annotateWithRetainerSize :: HeapGraph Size -> HeapGraph (Size, RetainerSize)
  annotateWithRetainerSize h@(HeapGraph rs _) =
    HeapGraph rs (foldMap convertToHeapGraph (retainerSize h))
  
  bottomUpSize :: Tree.Tree (HeapGraphEntry Size) -> Tree.Tree (HeapGraphEntry (Size, RetainerSize))
  bottomUpSize (Tree.Node rl sf) =
    let ts = map bottomUpSize sf
        s'@(Size s) =  hgeData rl
        RetainerSize children_size = foldMap (snd . hgeData . Tree.rootLabel) ts
        inclusive_size :: RetainerSize
        !inclusive_size = RetainerSize  (s + children_size)
        rl' = rl { hgeData = (s', inclusive_size) }
    in Tree.Node rl' ts
  
  convertToHeapGraph ::  Tree.Tree (HeapGraphEntry a) -> IM.IntMap (HeapGraphEntry a)
  convertToHeapGraph t = IM.fromList ([(fromIntegral cp, c) | c <- F.toList t, let ClosurePtr cp = hgeClosurePtr c ])
  
  
  
  