+-- | For heap graphs, i.e. data structures that also represent sharing and
+-- cyclic structures, these are the entries. If the referenced value is
+-- @Nothing@, then we do not have that value in the map, most likely due to
+-- exceeding the recursion bound passed to 'buildHeapGraph'.
+data HeapGraphEntry = HeapGraphEntry WeakBox (GenClosure (Maybe HeapGraphIndex))
+ deriving (Show)
+type HeapGraphIndex = Int
+
+-- | The whole graph. The suggested interface is to only use 'lookupHeapGraph',
+-- as the internal representation may change. Nevertheless, we export it here:
+-- Sometimes the user knows better what he needs than we do.
+newtype HeapGraph = HeapGraph (M.IntMap HeapGraphEntry)
+ deriving (Show)
+
+lookupHeapGraph :: HeapGraphIndex -> HeapGraph -> Maybe HeapGraphEntry
+lookupHeapGraph i (HeapGraph m) = M.lookup i m
+
+-- | Creates a 'HeapGraph' for the value in the box, but not recursing further
+-- than the given limit.
+buildHeapGraph :: Int -> Box -> IO HeapGraph
+buildHeapGraph limit initialBox = do
+ let initialState = (M.empty, [], [0..])
+ (\(m,_,_) -> HeapGraph m) <$> execStateT (add limit initialBox) initialState
+ where
+ add 0 _ = return Nothing
+ add n b = do
+ -- If the box is in the map, return the index
+ (_,existing,_) <- get
+ case lookup b existing of
+ Just i -> return $ Just i
+ Nothing -> do
+ -- Otherwise, allocate a new index
+ i <- nextI
+ -- And register it
+ modify (\(m,a,is) -> (m,(b,i):a,is))
+ c <- lift $ getBoxedClosureData b
+ -- Find indicies for all boxes contained in the map
+ c' <- T.mapM (add (n-1)) c
+ w <- lift $ weakBox b
+ -- Add add the resulting closure to the map
+ modify (\(m,a,is) -> (M.insert i (HeapGraphEntry w c') m,a,is))
+ return $ Just i
+ nextI = do
+ (_,_,i:_) <- get
+ modify (\(m,a,is) -> (m,a,tail is))
+ return i