{-# LANGUAGE RecordWildCards #-}

module Test.Tasty.AutoCollect.Utils.TreeMap (
  TreeMap (..),
  fromList,
  foldTreeMap,
) where

import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe)

data TreeMap k v = TreeMap
  { TreeMap k v -> Maybe v
value :: Maybe v
  , TreeMap k v -> Map k (TreeMap k v)
children :: Map k (TreeMap k v)
  }
  deriving (Int -> TreeMap k v -> ShowS
[TreeMap k v] -> ShowS
TreeMap k v -> String
(Int -> TreeMap k v -> ShowS)
-> (TreeMap k v -> String)
-> ([TreeMap k v] -> ShowS)
-> Show (TreeMap k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show v, Show k) => Int -> TreeMap k v -> ShowS
forall k v. (Show v, Show k) => [TreeMap k v] -> ShowS
forall k v. (Show v, Show k) => TreeMap k v -> String
showList :: [TreeMap k v] -> ShowS
$cshowList :: forall k v. (Show v, Show k) => [TreeMap k v] -> ShowS
show :: TreeMap k v -> String
$cshow :: forall k v. (Show v, Show k) => TreeMap k v -> String
showsPrec :: Int -> TreeMap k v -> ShowS
$cshowsPrec :: forall k v. (Show v, Show k) => Int -> TreeMap k v -> ShowS
Show, TreeMap k v -> TreeMap k v -> Bool
(TreeMap k v -> TreeMap k v -> Bool)
-> (TreeMap k v -> TreeMap k v -> Bool) -> Eq (TreeMap k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq v, Eq k) => TreeMap k v -> TreeMap k v -> Bool
/= :: TreeMap k v -> TreeMap k v -> Bool
$c/= :: forall k v. (Eq v, Eq k) => TreeMap k v -> TreeMap k v -> Bool
== :: TreeMap k v -> TreeMap k v -> Bool
$c== :: forall k v. (Eq v, Eq k) => TreeMap k v -> TreeMap k v -> Bool
Eq)

{- |
Convert the given list of values into a 'TreeMap'.

For example,
@
fromList [[A, B, C], [A, B], [A, C, D], [Z]]
@
would become
@
TreeMap
  { value = Nothing
  , children = Map.fromList
      [ ("A", TreeMap
          { value = Nothing
          , children = Map.fromList
              [ ("B", TreeMap
                  { value = Just ...
                  , children = Map.fromList
                      ("C", [ TreeMap
                          { value = Just ...
                          , children = Map.empty
                          }
                      ])
                  })
              , ("C", TreeMap
                  { value = Nothing
                  , children = Map.fromList
                      [ ("D", TreeMap
                          { value = Just ...
                          , children = Map.empty
                          })
                      ]
                  })
              ]
          })
    , ("Z", TreeMap
        { value = Just ...
        , children = Map.empty
        })
    ]
  }
@
-}
fromList :: Ord k => [([k], v)] -> TreeMap k v
fromList :: [([k], v)] -> TreeMap k v
fromList = (([k], v) -> TreeMap k v -> TreeMap k v)
-> TreeMap k v -> [([k], v)] -> TreeMap k v
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (([k] -> v -> TreeMap k v -> TreeMap k v)
-> ([k], v) -> TreeMap k v -> TreeMap k v
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [k] -> v -> TreeMap k v -> TreeMap k v
forall k v. Ord k => [k] -> v -> TreeMap k v -> TreeMap k v
insert) TreeMap k v
forall k v. TreeMap k v
empty

empty :: TreeMap k v
empty :: TreeMap k v
empty = Maybe v -> Map k (TreeMap k v) -> TreeMap k v
forall k v. Maybe v -> Map k (TreeMap k v) -> TreeMap k v
TreeMap Maybe v
forall a. Maybe a
Nothing Map k (TreeMap k v)
forall k a. Map k a
Map.empty

insert :: Ord k => [k] -> v -> TreeMap k v -> TreeMap k v
insert :: [k] -> v -> TreeMap k v -> TreeMap k v
insert [k]
originalKeys v
v = [k] -> TreeMap k v -> TreeMap k v
forall k. Ord k => [k] -> TreeMap k v -> TreeMap k v
go [k]
originalKeys
  where
    go :: [k] -> TreeMap k v -> TreeMap k v
go [k]
ks TreeMap k v
treeMap =
      case [k]
ks of
        [] -> TreeMap k v
treeMap{value :: Maybe v
value = v -> Maybe v
forall a. a -> Maybe a
Just v
v}
        k
k : [k]
ks' -> TreeMap k v
treeMap{children :: Map k (TreeMap k v)
children = (Maybe (TreeMap k v) -> Maybe (TreeMap k v))
-> k -> Map k (TreeMap k v) -> Map k (TreeMap k v)
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter (TreeMap k v -> Maybe (TreeMap k v)
forall a. a -> Maybe a
Just (TreeMap k v -> Maybe (TreeMap k v))
-> (Maybe (TreeMap k v) -> TreeMap k v)
-> Maybe (TreeMap k v)
-> Maybe (TreeMap k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [k] -> TreeMap k v -> TreeMap k v
go [k]
ks' (TreeMap k v -> TreeMap k v)
-> (Maybe (TreeMap k v) -> TreeMap k v)
-> Maybe (TreeMap k v)
-> TreeMap k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TreeMap k v -> Maybe (TreeMap k v) -> TreeMap k v
forall a. a -> Maybe a -> a
fromMaybe TreeMap k v
forall k v. TreeMap k v
empty) k
k (TreeMap k v -> Map k (TreeMap k v)
forall k v. TreeMap k v -> Map k (TreeMap k v)
children TreeMap k v
treeMap)}

foldTreeMap :: (Maybe v -> Map k r -> r) -> TreeMap k v -> r
foldTreeMap :: (Maybe v -> Map k r -> r) -> TreeMap k v -> r
foldTreeMap Maybe v -> Map k r -> r
f = TreeMap k v -> r
go
  where
    go :: TreeMap k v -> r
go TreeMap{Maybe v
Map k (TreeMap k v)
children :: Map k (TreeMap k v)
value :: Maybe v
children :: forall k v. TreeMap k v -> Map k (TreeMap k v)
value :: forall k v. TreeMap k v -> Maybe v
..} = Maybe v -> Map k r -> r
f Maybe v
value (TreeMap k v -> r
go (TreeMap k v -> r) -> Map k (TreeMap k v) -> Map k r
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map k (TreeMap k v)
children)