{-# 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
  { forall k v. TreeMap k v -> Maybe v
value :: Maybe v
  , forall k v. TreeMap k v -> Map k (TreeMap k v)
children :: Map k (TreeMap k v)
  }
  deriving (Int -> TreeMap k v -> ShowS
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
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 :: forall k v. Ord k => [([k], v)] -> TreeMap k v
fromList = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall k v. Ord k => [k] -> v -> TreeMap k v -> TreeMap k v
insert) forall k v. TreeMap k v
empty

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

insert :: (Ord k) => [k] -> v -> TreeMap k v -> TreeMap k v
insert :: forall k v. Ord k => [k] -> v -> TreeMap k v -> TreeMap k v
insert [k]
originalKeys v
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 = forall a. a -> Maybe a
Just v
v}
        k
k : [k]
ks' -> TreeMap k v
treeMap{children :: Map k (TreeMap k v)
children = forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter (forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. [k] -> TreeMap k v -> TreeMap k v
go [k]
ks' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a -> a
fromMaybe forall k v. TreeMap k v
empty) k
k (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 :: forall v k r. (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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map k (TreeMap k v)
children)