{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE NamedFieldPuns    #-}
{-# LANGUAGE OverloadedStrings #-}

{-|
Module      :  Cardano.Prelude.GHC.Heap.Tree

This module exports functions and data types for building and printing GHC
heap object trees. Very useful for debugging issues involving the heap
representation of Haskell expressions.

-}

module Cardano.Prelude.GHC.Heap.Tree
  ( ClosureTreeOptions(..)
  , TraverseCyclicClosures(..)
  , TreeDepth(..)
  , buildClosureTree
  , buildAndRenderClosureTree
  , depth
  , isZeroOrNegativeTreeDepth
  , renderClosure
  , renderTree
  )
where

import Cardano.Prelude.Base

import Data.Text (pack, unpack)
import Data.Tree (Tree(..), drawTree, levels)
import GHC.Exts.Heap
  (Box(..), Closure, allClosures, areBoxesEqual, asBox, getClosureData)
import System.Mem (performGC)

-- | The depth of a 'Tree'.
data TreeDepth
  = TreeDepth {-# UNPACK #-} !Int
  -- ^ A specific tree depth bound.
  | AnyDepth
  -- ^ No tree depth bound.
  deriving (Int -> TreeDepth -> ShowS
[TreeDepth] -> ShowS
TreeDepth -> String
(Int -> TreeDepth -> ShowS)
-> (TreeDepth -> String)
-> ([TreeDepth] -> ShowS)
-> Show TreeDepth
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TreeDepth] -> ShowS
$cshowList :: [TreeDepth] -> ShowS
show :: TreeDepth -> String
$cshow :: TreeDepth -> String
showsPrec :: Int -> TreeDepth -> ShowS
$cshowsPrec :: Int -> TreeDepth -> ShowS
Show)

-- | Whether to traverse cyclic closures in a 'Closure' 'Tree'.
data TraverseCyclicClosures
  = TraverseCyclicClosures
  -- ^ Traverse cyclic closures.
  | NoTraverseCyclicClosures
  -- ^ Do not traverse cyclic closures.
  deriving (Int -> TraverseCyclicClosures -> ShowS
[TraverseCyclicClosures] -> ShowS
TraverseCyclicClosures -> String
(Int -> TraverseCyclicClosures -> ShowS)
-> (TraverseCyclicClosures -> String)
-> ([TraverseCyclicClosures] -> ShowS)
-> Show TraverseCyclicClosures
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TraverseCyclicClosures] -> ShowS
$cshowList :: [TraverseCyclicClosures] -> ShowS
show :: TraverseCyclicClosures -> String
$cshow :: TraverseCyclicClosures -> String
showsPrec :: Int -> TraverseCyclicClosures -> ShowS
$cshowsPrec :: Int -> TraverseCyclicClosures -> ShowS
Show)

-- | Options which detail how a 'Closure' 'Tree' should be constructed.
data ClosureTreeOptions = ClosureTreeOptions
  { ClosureTreeOptions -> TreeDepth
ctoMaxDepth       :: !TreeDepth
  -- ^ Construct a closure tree given a maximum depth.
  , ClosureTreeOptions -> TraverseCyclicClosures
ctoCyclicClosures :: !TraverseCyclicClosures
  -- ^ Whether to traverse cyclic closures while constructing a closure tree.
  } deriving (Int -> ClosureTreeOptions -> ShowS
[ClosureTreeOptions] -> ShowS
ClosureTreeOptions -> String
(Int -> ClosureTreeOptions -> ShowS)
-> (ClosureTreeOptions -> String)
-> ([ClosureTreeOptions] -> ShowS)
-> Show ClosureTreeOptions
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ClosureTreeOptions] -> ShowS
$cshowList :: [ClosureTreeOptions] -> ShowS
show :: ClosureTreeOptions -> String
$cshow :: ClosureTreeOptions -> String
showsPrec :: Int -> ClosureTreeOptions -> ShowS
$cshowsPrec :: Int -> ClosureTreeOptions -> ShowS
Show)

depth :: Tree a -> Int
depth :: Tree a -> Int
depth = [[a]] -> Int
forall a. HasLength a => a -> Int
length ([[a]] -> Int) -> (Tree a -> [[a]]) -> Tree a -> Int
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Tree a -> [[a]]
forall a. Tree a -> [[a]]
levels

isZeroOrNegativeTreeDepth :: TreeDepth -> Bool
isZeroOrNegativeTreeDepth :: TreeDepth -> Bool
isZeroOrNegativeTreeDepth TreeDepth
AnyDepth = Bool
False
isZeroOrNegativeTreeDepth (TreeDepth Int
d)
  | Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0    = Bool
True
  | Bool
otherwise = Bool
False

renderClosure :: Closure -> Text
renderClosure :: Closure -> Text
renderClosure = Closure -> Text
forall a b. (Show a, ConvertText String b) => a -> b
show

renderTree :: Tree a -> (a -> Text) -> Text
renderTree :: Tree a -> (a -> Text) -> Text
renderTree Tree a
tree a -> Text
renderA = String -> Text
pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ Tree String -> String
drawTree ((a -> String) -> Tree a -> Tree String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Text -> String
unpack (Text -> String) -> (a -> Text) -> a -> String
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. a -> Text
renderA) Tree a
tree)

-- | Given a Haskell expression, build a 'Tree' which reflects its heap object
-- representation.
buildClosureTree :: ClosureTreeOptions -> a -> IO (Maybe (Tree Closure))
buildClosureTree :: ClosureTreeOptions -> a -> IO (Maybe (Tree Closure))
buildClosureTree ClosureTreeOptions
opts a
x = do
  IO ()
performGC
  ClosureTreeOptions -> [Box] -> Box -> IO (Maybe (Tree Closure))
go ClosureTreeOptions
opts [] (Box -> IO (Maybe (Tree Closure)))
-> Box -> IO (Maybe (Tree Closure))
forall a b. (a -> b) -> a -> b
$ a -> Box
forall a. a -> Box
asBox a
x
 where
  go :: ClosureTreeOptions -> [Box] -> Box -> IO (Maybe (Tree Closure))
  go :: ClosureTreeOptions -> [Box] -> Box -> IO (Maybe (Tree Closure))
go (ClosureTreeOptions { TreeDepth
ctoMaxDepth :: TreeDepth
ctoMaxDepth :: ClosureTreeOptions -> TreeDepth
ctoMaxDepth, TraverseCyclicClosures
ctoCyclicClosures :: TraverseCyclicClosures
ctoCyclicClosures :: ClosureTreeOptions -> TraverseCyclicClosures
ctoCyclicClosures }) ![Box]
vs b :: Box
b@(Box Any
y)
    | TreeDepth -> Bool
isZeroOrNegativeTreeDepth TreeDepth
ctoMaxDepth = Maybe (Tree Closure) -> IO (Maybe (Tree Closure))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Tree Closure)
forall a. Maybe a
Nothing
    | Bool
otherwise = do
      let
        nextMaxDepth :: TreeDepth
nextMaxDepth = case TreeDepth
ctoMaxDepth of
          TreeDepth
AnyDepth    -> TreeDepth
AnyDepth
          TreeDepth Int
d -> Int -> TreeDepth
TreeDepth (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        nextOpts :: ClosureTreeOptions
nextOpts = TreeDepth -> TraverseCyclicClosures -> ClosureTreeOptions
ClosureTreeOptions TreeDepth
nextMaxDepth TraverseCyclicClosures
ctoCyclicClosures
      case TraverseCyclicClosures
ctoCyclicClosures of
        TraverseCyclicClosures
NoTraverseCyclicClosures -> do
          Bool
isElem <- ([Bool] -> Bool) -> IO [Bool] -> IO Bool
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or (IO [Bool] -> IO Bool) -> IO [Bool] -> IO Bool
forall a b. (a -> b) -> a -> b
$ (Box -> IO Bool) -> [Box] -> IO [Bool]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Box -> Box -> IO Bool
areBoxesEqual Box
b) [Box]
vs
          if Bool
isElem
            then Maybe (Tree Closure) -> IO (Maybe (Tree Closure))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Tree Closure)
forall a. Maybe a
Nothing
            else do
              Closure
closure  <- Any -> IO Closure
forall a. HasHeapRep a => a -> IO Closure
getClosureData Any
y
              [Maybe (Tree Closure)]
subtrees <- (Box -> IO (Maybe (Tree Closure)))
-> [Box] -> IO [Maybe (Tree Closure)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (ClosureTreeOptions -> [Box] -> Box -> IO (Maybe (Tree Closure))
go ClosureTreeOptions
nextOpts (Box
b Box -> [Box] -> [Box]
forall a. a -> [a] -> [a]
: [Box]
vs)) (Closure -> [Box]
forall b. GenClosure b -> [b]
allClosures Closure
closure)
              Maybe (Tree Closure) -> IO (Maybe (Tree Closure))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Tree Closure -> Maybe (Tree Closure)
forall a. a -> Maybe a
Just (Closure -> [Tree Closure] -> Tree Closure
forall a. a -> [Tree a] -> Tree a
Node Closure
closure ([Maybe (Tree Closure)] -> [Tree Closure]
forall a. [Maybe a] -> [a]
catMaybes [Maybe (Tree Closure)]
subtrees)))
        TraverseCyclicClosures
TraverseCyclicClosures -> do
          Closure
closure  <- Any -> IO Closure
forall a. HasHeapRep a => a -> IO Closure
getClosureData Any
y
          [Maybe (Tree Closure)]
subtrees <- (Box -> IO (Maybe (Tree Closure)))
-> [Box] -> IO [Maybe (Tree Closure)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (ClosureTreeOptions -> [Box] -> Box -> IO (Maybe (Tree Closure))
go ClosureTreeOptions
nextOpts (Box
b Box -> [Box] -> [Box]
forall a. a -> [a] -> [a]
: [Box]
vs)) (Closure -> [Box]
forall b. GenClosure b -> [b]
allClosures Closure
closure)
          Maybe (Tree Closure) -> IO (Maybe (Tree Closure))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Tree Closure -> Maybe (Tree Closure)
forall a. a -> Maybe a
Just (Closure -> [Tree Closure] -> Tree Closure
forall a. a -> [Tree a] -> Tree a
Node Closure
closure ([Maybe (Tree Closure)] -> [Tree Closure]
forall a. [Maybe a] -> [a]
catMaybes [Maybe (Tree Closure)]
subtrees)))

-- | Given a Haskell expression, build a 'Tree' which reflects its heap object
-- representation and render it as 'Text'.
buildAndRenderClosureTree :: ClosureTreeOptions -> a -> IO Text
buildAndRenderClosureTree :: ClosureTreeOptions -> a -> IO Text
buildAndRenderClosureTree ClosureTreeOptions
opts a
x = do
  Maybe (Tree Closure)
mbTr <- ClosureTreeOptions -> a -> IO (Maybe (Tree Closure))
forall a. ClosureTreeOptions -> a -> IO (Maybe (Tree Closure))
buildClosureTree ClosureTreeOptions
opts a
x
  case Maybe (Tree Closure)
mbTr of
    Maybe (Tree Closure)
Nothing -> Text -> IO Text
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
forall a. Monoid a => a
mempty
    Just Tree Closure
tr -> Text -> IO Text
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Tree Closure -> (Closure -> Text) -> Text
forall a. Tree a -> (a -> Text) -> Text
renderTree Tree Closure
tr Closure -> Text
renderClosure)