{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- | Trace-based tree trace store.
module Test.Mutagen.Tracer.Store.Tree
  ( -- * Tree trace store
    TraceStore
  , newTraceStore
  , resetTraceStore
  , saveTrace
  , printTraceStore
  )
where

import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import qualified Data.Tree as Tree
import Data.Tree.Pretty (drawVerticalForest)
import Test.Mutagen.Tracer.Store.API (TraceStoreImpl (..))
import Test.Mutagen.Tracer.Store.Types (TraceBackend (Tree))
import Test.Mutagen.Tracer.Trace (Trace (..), TraceNode)

{-------------------------------------------------------------------------------
-- * Tree trace store
-------------------------------------------------------------------------------}

instance TraceStoreImpl Tree where
  newtype TraceStore Tree = TreeStore (IORef TraceTree)

  type SaveTraceResult Tree = (Int, Int)

  newTraceStore :: TraceNode -> IO (TraceStore 'Tree)
newTraceStore TraceNode
_ = do
    IORef TraceTree -> TraceStore 'Tree
TreeStore (IORef TraceTree -> TraceStore 'Tree)
-> IO (IORef TraceTree) -> IO (TraceStore 'Tree)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TraceTree -> IO (IORef TraceTree)
forall a. a -> IO (IORef a)
newIORef TraceTree
emptyTraceTree

  resetTraceStore :: TraceStore 'Tree -> IO ()
resetTraceStore (TreeStore IORef TraceTree
ref) = do
    IORef TraceTree -> TraceTree -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef TraceTree
ref TraceTree
emptyTraceTree

  saveTrace :: Trace -> TraceStore 'Tree -> IO (SaveTraceResult 'Tree)
saveTrace Trace
trace (TreeStore IORef TraceTree
ref) = do
    tt <- IORef TraceTree -> IO TraceTree
forall a. IORef a -> IO a
readIORef IORef TraceTree
ref
    let (tt', n, d) = insertTrace trace tt
    writeIORef ref tt'
    return (n, d)

  printTraceStore :: TraceStore 'Tree -> IO ()
printTraceStore (TreeStore IORef TraceTree
ref) = do
    tt <- IORef TraceTree -> IO TraceTree
forall a. IORef a -> IO a
readIORef IORef TraceTree
ref
    putStrLn (drawTraceTree tt)

{-------------------------------------------------------------------------------
-- * Trace trees implemented using "rose maps"
-------------------------------------------------------------------------------}

-- | Recursive partial maps indexed by trace nodes.
newtype TraceTree = TraceTree (Map TraceNode TraceTree)

-- | Create an empty trace tree.
emptyTraceTree :: TraceTree
emptyTraceTree :: TraceTree
emptyTraceTree = Map TraceNode TraceTree -> TraceTree
TraceTree Map TraceNode TraceTree
forall a. Monoid a => a
mempty

-- | Insert a trace into a trace tree.
--
-- Returns the updated tree, the number of new nodes added, and the depth at
-- which the input trace was inserted into the global trace tree.
insertTrace :: Trace -> TraceTree -> (TraceTree, Int, Int)
insertTrace :: Trace -> TraceTree -> (TraceTree, TraceNode, TraceNode)
insertTrace (Trace [TraceNode]
entries) = TraceNode
-> [TraceNode] -> TraceTree -> (TraceTree, TraceNode, TraceNode)
forall {b} {t}.
(Num b, Num t) =>
t -> [TraceNode] -> TraceTree -> (TraceTree, b, t)
go TraceNode
0 [TraceNode]
entries
  where
    go :: t -> [TraceNode] -> TraceTree -> (TraceTree, b, t)
go t
d [] (TraceTree Map TraceNode TraceTree
tt) = (Map TraceNode TraceTree -> TraceTree
TraceTree Map TraceNode TraceTree
tt, b
0, t
d)
    go t
d (TraceNode
e : [TraceNode]
es) (TraceTree Map TraceNode TraceTree
tt) =
      case TraceNode -> Map TraceNode TraceTree -> Maybe TraceTree
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TraceNode
e Map TraceNode TraceTree
tt of
        Maybe TraceTree
Nothing ->
          let (TraceTree
subTree, b
new) = [TraceNode] -> (TraceTree, b)
forall {b}. Num b => [TraceNode] -> (TraceTree, b)
chain [TraceNode]
es
           in (Map TraceNode TraceTree -> TraceTree
TraceTree (TraceNode
-> TraceTree -> Map TraceNode TraceTree -> Map TraceNode TraceTree
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert TraceNode
e TraceTree
subTree Map TraceNode TraceTree
tt), b
new b -> b -> b
forall a. Num a => a -> a -> a
+ b
1, t
d)
        Just TraceTree
subTree ->
          let (TraceTree
subTree', b
new, t
d') = t -> [TraceNode] -> TraceTree -> (TraceTree, b, t)
go (t
d t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) [TraceNode]
es TraceTree
subTree
           in (Map TraceNode TraceTree -> TraceTree
TraceTree (TraceNode
-> TraceTree -> Map TraceNode TraceTree -> Map TraceNode TraceTree
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert TraceNode
e TraceTree
subTree' Map TraceNode TraceTree
tt), b
new, t
d')

    chain :: [TraceNode] -> (TraceTree, b)
chain [] = (Map TraceNode TraceTree -> TraceTree
TraceTree Map TraceNode TraceTree
forall a. Monoid a => a
mempty, b
0)
    chain (TraceNode
e : [TraceNode]
es) = (Map TraceNode TraceTree -> TraceTree
TraceTree (TraceNode -> TraceTree -> Map TraceNode TraceTree
forall k a. k -> a -> Map k a
Map.singleton TraceNode
e TraceTree
tlog'), b
n b -> b -> b
forall a. Num a => a -> a -> a
+ b
1)
      where
        (TraceTree
tlog', b
n) = [TraceNode] -> (TraceTree, b)
chain [TraceNode]
es

-- | Convert a trace tree to a forest for pretty printing.
traceTreetoForest :: TraceTree -> Tree.Forest TraceNode
traceTreetoForest :: TraceTree -> Forest TraceNode
traceTreetoForest (TraceTree Map TraceNode TraceTree
tt) =
  Map TraceNode (Tree TraceNode) -> Forest TraceNode
forall k a. Map k a -> [a]
Map.elems ((TraceNode -> TraceTree -> Tree TraceNode)
-> Map TraceNode TraceTree -> Map TraceNode (Tree TraceNode)
forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey (\TraceNode
node -> TraceNode -> Forest TraceNode -> Tree TraceNode
forall a. a -> [Tree a] -> Tree a
Tree.Node TraceNode
node (Forest TraceNode -> Tree TraceNode)
-> (TraceTree -> Forest TraceNode) -> TraceTree -> Tree TraceNode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TraceTree -> Forest TraceNode
traceTreetoForest) Map TraceNode TraceTree
tt)

-- | Pretty print a trace tree.
drawTraceTree :: TraceTree -> String
drawTraceTree :: TraceTree -> String
drawTraceTree TraceTree
tt = Forest String -> String
drawVerticalForest ((TraceNode -> String) -> Tree TraceNode -> Tree String
forall a b. (a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TraceNode -> String
forall a. Show a => a -> String
show (Tree TraceNode -> Tree String)
-> Forest TraceNode -> Forest String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TraceTree -> Forest TraceNode
toForest TraceTree
tt)
  where
    toForest :: TraceTree -> Forest TraceNode
toForest (TraceTree Map TraceNode TraceTree
tt') =
      Map TraceNode (Tree TraceNode) -> Forest TraceNode
forall k a. Map k a -> [a]
Map.elems ((TraceNode -> TraceTree -> Tree TraceNode)
-> Map TraceNode TraceTree -> Map TraceNode (Tree TraceNode)
forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey (\TraceNode
node -> TraceNode -> Forest TraceNode -> Tree TraceNode
forall a. a -> [Tree a] -> Tree a
Tree.Node TraceNode
node (Forest TraceNode -> Tree TraceNode)
-> (TraceTree -> Forest TraceNode) -> TraceTree -> Tree TraceNode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TraceTree -> Forest TraceNode
traceTreetoForest) Map TraceNode TraceTree
tt')