-- | Prefix trees on arbitrary keys.
module Data.Trie
  ( -- * Type
    Trie (..),

    -- * Construction
    empty,
    singleton,

    -- * Basic interface
    lookup,
    insert,
    insertWith,
    elems,
  )
where

import Data.Aeson (ToJSON, ToJSONKey)
import Data.HashMap.Strict qualified as M
import Data.Hashable (Hashable)
import GHC.Generics (Generic)
import Prelude hiding (lookup)

-------------------------------------------------------------------------------

-- | Data structure for storing a value @v@ keyed on a sequence of @k@s
data Trie k v = Trie
  { Trie k v -> HashMap k (Trie k v)
trieMap :: M.HashMap k (Trie k v),
    Trie k v -> Maybe v
trieData :: Maybe v
  }
  deriving stock (Trie k v -> Trie k v -> Bool
(Trie k v -> Trie k v -> Bool)
-> (Trie k v -> Trie k v -> Bool) -> Eq (Trie k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => Trie k v -> Trie k v -> Bool
/= :: Trie k v -> Trie k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => Trie k v -> Trie k v -> Bool
== :: Trie k v -> Trie k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => Trie k v -> Trie k v -> Bool
Eq, Int -> Trie k v -> ShowS
[Trie k v] -> ShowS
Trie k v -> String
(Int -> Trie k v -> ShowS)
-> (Trie k v -> String) -> ([Trie k v] -> ShowS) -> Show (Trie k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> Trie k v -> ShowS
forall k v. (Show k, Show v) => [Trie k v] -> ShowS
forall k v. (Show k, Show v) => Trie k v -> String
showList :: [Trie k v] -> ShowS
$cshowList :: forall k v. (Show k, Show v) => [Trie k v] -> ShowS
show :: Trie k v -> String
$cshow :: forall k v. (Show k, Show v) => Trie k v -> String
showsPrec :: Int -> Trie k v -> ShowS
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> Trie k v -> ShowS
Show, Eq (Trie k v)
Eq (Trie k v)
-> (Trie k v -> Trie k v -> Ordering)
-> (Trie k v -> Trie k v -> Bool)
-> (Trie k v -> Trie k v -> Bool)
-> (Trie k v -> Trie k v -> Bool)
-> (Trie k v -> Trie k v -> Bool)
-> (Trie k v -> Trie k v -> Trie k v)
-> (Trie k v -> Trie k v -> Trie k v)
-> Ord (Trie k v)
Trie k v -> Trie k v -> Bool
Trie k v -> Trie k v -> Ordering
Trie k v -> Trie k v -> Trie k v
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall k v. (Ord k, Ord v) => Eq (Trie k v)
forall k v. (Ord k, Ord v) => Trie k v -> Trie k v -> Bool
forall k v. (Ord k, Ord v) => Trie k v -> Trie k v -> Ordering
forall k v. (Ord k, Ord v) => Trie k v -> Trie k v -> Trie k v
min :: Trie k v -> Trie k v -> Trie k v
$cmin :: forall k v. (Ord k, Ord v) => Trie k v -> Trie k v -> Trie k v
max :: Trie k v -> Trie k v -> Trie k v
$cmax :: forall k v. (Ord k, Ord v) => Trie k v -> Trie k v -> Trie k v
>= :: Trie k v -> Trie k v -> Bool
$c>= :: forall k v. (Ord k, Ord v) => Trie k v -> Trie k v -> Bool
> :: Trie k v -> Trie k v -> Bool
$c> :: forall k v. (Ord k, Ord v) => Trie k v -> Trie k v -> Bool
<= :: Trie k v -> Trie k v -> Bool
$c<= :: forall k v. (Ord k, Ord v) => Trie k v -> Trie k v -> Bool
< :: Trie k v -> Trie k v -> Bool
$c< :: forall k v. (Ord k, Ord v) => Trie k v -> Trie k v -> Bool
compare :: Trie k v -> Trie k v -> Ordering
$ccompare :: forall k v. (Ord k, Ord v) => Trie k v -> Trie k v -> Ordering
$cp1Ord :: forall k v. (Ord k, Ord v) => Eq (Trie k v)
Ord, (forall x. Trie k v -> Rep (Trie k v) x)
-> (forall x. Rep (Trie k v) x -> Trie k v) -> Generic (Trie k v)
forall x. Rep (Trie k v) x -> Trie k v
forall x. Trie k v -> Rep (Trie k v) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall k v x. Rep (Trie k v) x -> Trie k v
forall k v x. Trie k v -> Rep (Trie k v) x
$cto :: forall k v x. Rep (Trie k v) x -> Trie k v
$cfrom :: forall k v x. Trie k v -> Rep (Trie k v) x
Generic)

-- | Semigroup via union.
-- The resulting 'Trie' will contain all paths present in either tries. If both
-- tries contain a value at a given path, we use the value's semigroup instance
-- to compute the resulting value.
instance (Eq k, Hashable k, Semigroup v) => Semigroup (Trie k v) where
  Trie HashMap k (Trie k v)
m0 Maybe v
v0 <> :: Trie k v -> Trie k v -> Trie k v
<> Trie HashMap k (Trie k v)
m1 Maybe v
v1 = HashMap k (Trie k v) -> Maybe v -> Trie k v
forall k v. HashMap k (Trie k v) -> Maybe v -> Trie k v
Trie ((Trie k v -> Trie k v -> Trie k v)
-> HashMap k (Trie k v)
-> HashMap k (Trie k v)
-> HashMap k (Trie k v)
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
M.unionWith Trie k v -> Trie k v -> Trie k v
forall a. Semigroup a => a -> a -> a
(<>) HashMap k (Trie k v)
m0 HashMap k (Trie k v)
m1) (Maybe v
v0 Maybe v -> Maybe v -> Maybe v
forall a. Semigroup a => a -> a -> a
<> Maybe v
v1)

instance (Eq k, Hashable k, Semigroup v) => Monoid (Trie k v) where
  mempty :: Trie k v
mempty = Trie k v
forall k v. Trie k v
empty

instance (ToJSONKey a, ToJSON v) => ToJSON (Trie a v)

-------------------------------------------------------------------------------

-- | Construct an empty trie.
empty :: Trie k v
empty :: Trie k v
empty = HashMap k (Trie k v) -> Maybe v -> Trie k v
forall k v. HashMap k (Trie k v) -> Maybe v -> Trie k v
Trie HashMap k (Trie k v)
forall k v. HashMap k v
M.empty Maybe v
forall a. Maybe a
Nothing

-- | Creates a trie from a path and a value
--
-- >>> singleton ["a", "b"] 5
-- Trie (fromList [("a", Trie (fromList [("b", Trie (fromList []) (Just 5))]) Nothing)]) Nothing
--
-- >>> singleton [] 5
-- Trie (fromList []) (Just 5)
singleton :: (Hashable k) => [k] -> v -> Trie k v
singleton :: [k] -> v -> Trie k v
singleton [k]
ps v
v = (k -> Trie k v -> Trie k v) -> Trie k v -> [k] -> Trie k v
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\k
p Trie k v
t -> HashMap k (Trie k v) -> Maybe v -> Trie k v
forall k v. HashMap k (Trie k v) -> Maybe v -> Trie k v
Trie (k -> Trie k v -> HashMap k (Trie k v)
forall k v. Hashable k => k -> v -> HashMap k v
M.singleton k
p Trie k v
t) Maybe v
forall a. Maybe a
Nothing) (HashMap k (Trie k v) -> Maybe v -> Trie k v
forall k v. HashMap k (Trie k v) -> Maybe v -> Trie k v
Trie HashMap k (Trie k v)
forall k v. HashMap k v
M.empty (v -> Maybe v
forall a. a -> Maybe a
Just v
v)) [k]
ps

-------------------------------------------------------------------------------

-- | Find a value at the given path, if any.
lookup :: (Eq k, Hashable k) => [k] -> Trie k v -> Maybe v
lookup :: [k] -> Trie k v -> Maybe v
lookup [] (Trie HashMap k (Trie k v)
_ Maybe v
value) = Maybe v
value
lookup (k
p : [k]
ps) (Trie HashMap k (Trie k v)
tmap Maybe v
_) = [k] -> Trie k v -> Maybe v
forall k v. (Eq k, Hashable k) => [k] -> Trie k v -> Maybe v
lookup [k]
ps (Trie k v -> Maybe v) -> Maybe (Trie k v) -> Maybe v
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< k -> HashMap k (Trie k v) -> Maybe (Trie k v)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup k
p HashMap k (Trie k v)
tmap

-- | Insert the given value at the given path.
--
-- If there's already a value at the given path, it is replaced.
insert :: (Eq k, Hashable k) => [k] -> v -> Trie k v -> Trie k v
insert :: [k] -> v -> Trie k v -> Trie k v
insert = (v -> v -> v) -> [k] -> v -> Trie k v -> Trie k v
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> [k] -> v -> Trie k v -> Trie k v
insertWith v -> v -> v
forall a b. a -> b -> a
const

-- | Insert the value at the given path.
--
-- If there's already a value at the given path, the old value is replaced by
-- the result of applying the given function to the new and old value.
insertWith :: (Eq k, Hashable k) => (v -> v -> v) -> [k] -> v -> Trie k v -> Trie k v
insertWith :: (v -> v -> v) -> [k] -> v -> Trie k v -> Trie k v
insertWith v -> v -> v
fun [k]
path v
newValue Trie k v
t = Trie k v -> [k] -> Trie k v
go Trie k v
t [k]
path
  where
    go :: Trie k v -> [k] -> Trie k v
go (Trie HashMap k (Trie k v)
tmap Maybe v
value) = \case
      [] -> HashMap k (Trie k v) -> Maybe v -> Trie k v
forall k v. HashMap k (Trie k v) -> Maybe v -> Trie k v
Trie HashMap k (Trie k v)
tmap (Maybe v -> Trie k v) -> Maybe v -> Trie k v
forall a b. (a -> b) -> a -> b
$
        v -> Maybe v
forall a. a -> Maybe a
Just (v -> Maybe v) -> v -> Maybe v
forall a b. (a -> b) -> a -> b
$ case Maybe v
value of
          Maybe v
Nothing -> v
newValue
          Just v
oldValue -> v -> v -> v
fun v
newValue v
oldValue
      (k
p : [k]
ps) -> HashMap k (Trie k v) -> Maybe v -> Trie k v
forall k v. HashMap k (Trie k v) -> Maybe v -> Trie k v
Trie ((Maybe (Trie k v) -> Maybe (Trie k v))
-> k -> HashMap k (Trie k v) -> HashMap k (Trie k v)
forall k v.
(Eq k, Hashable k) =>
(Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
M.alter ([k] -> Maybe (Trie k v) -> Maybe (Trie k v)
step [k]
ps) k
p HashMap k (Trie k v)
tmap) Maybe v
value
    step :: [k] -> Maybe (Trie k v) -> Maybe (Trie k v)
step [k]
ps = \case
      -- this path did not exist and must be created
      Maybe (Trie k v)
Nothing -> Trie k v -> Maybe (Trie k v)
forall a. a -> Maybe a
Just (Trie k v -> Maybe (Trie k v)) -> Trie k v -> Maybe (Trie k v)
forall a b. (a -> b) -> a -> b
$ Trie k v -> [k] -> Trie k v
go Trie k v
forall k v. Trie k v
empty [k]
ps
      -- we found the path
      Just Trie k v
st -> Trie k v -> Maybe (Trie k v)
forall a. a -> Maybe a
Just (Trie k v -> Maybe (Trie k v)) -> Trie k v -> Maybe (Trie k v)
forall a b. (a -> b) -> a -> b
$ Trie k v -> [k] -> Trie k v
go Trie k v
st [k]
ps

-- | Extract all values of the trie, discarding any path information.
elems :: Trie k v -> [v]
elems :: Trie k v -> [v]
elems (Trie HashMap k (Trie k v)
m Maybe v
v) =
  let subElems :: [v]
subElems = (Trie k v -> [v]) -> [Trie k v] -> [v]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Trie k v -> [v]
forall k v. Trie k v -> [v]
elems (HashMap k (Trie k v) -> [Trie k v]
forall k v. HashMap k v -> [v]
M.elems HashMap k (Trie k v)
m)
   in case Maybe v
v of
        Maybe v
Nothing -> [v]
subElems
        Just v
val -> v
val v -> [v] -> [v]
forall a. a -> [a] -> [a]
: [v]
subElems