mirror of
https://github.com/hasura/graphql-engine.git
synced 2024-12-16 09:51:59 +03:00
100 lines
3.1 KiB
Haskell
100 lines
3.1 KiB
Haskell
|
-- | 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
|
||
|
{ trieMap :: M.HashMap k (Trie k v),
|
||
|
trieData :: Maybe v
|
||
|
}
|
||
|
deriving stock (Eq, Show, Ord, 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 m0 v0 <> Trie m1 v1 = Trie (M.unionWith (<>) m0 m1) (v0 <> v1)
|
||
|
|
||
|
instance (Eq k, Hashable k, Semigroup v) => Monoid (Trie k v) where
|
||
|
mempty = empty
|
||
|
|
||
|
instance (ToJSONKey a, ToJSON v) => ToJSON (Trie a v)
|
||
|
|
||
|
-------------------------------------------------------------------------------
|
||
|
|
||
|
-- | Construct an empty trie.
|
||
|
empty :: Trie k v
|
||
|
empty = Trie M.empty 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 ps v = foldr (\p t -> Trie (M.singleton p t) Nothing) (Trie M.empty (Just v)) ps
|
||
|
|
||
|
-------------------------------------------------------------------------------
|
||
|
|
||
|
-- | Find a value at the given path, if any.
|
||
|
lookup :: (Eq k, Hashable k) => [k] -> Trie k v -> Maybe v
|
||
|
lookup [] (Trie _ value) = value
|
||
|
lookup (p : ps) (Trie tmap _) = lookup ps =<< M.lookup p 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 = insertWith 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 fun path newValue t = go t path
|
||
|
where
|
||
|
go (Trie tmap value) = \case
|
||
|
[] -> Trie tmap $
|
||
|
Just $ case value of
|
||
|
Nothing -> newValue
|
||
|
Just oldValue -> fun newValue oldValue
|
||
|
(p : ps) -> Trie (M.alter (step ps) p tmap) value
|
||
|
step ps = \case
|
||
|
-- this path did not exist and must be created
|
||
|
Nothing -> Just $ go empty ps
|
||
|
-- we found the path
|
||
|
Just st -> Just $ go st ps
|
||
|
|
||
|
-- | Extract all values of the trie, discarding any path information.
|
||
|
elems :: Trie k v -> [v]
|
||
|
elems (Trie m v) =
|
||
|
let subElems = concatMap elems (M.elems m)
|
||
|
in case v of
|
||
|
Nothing -> subElems
|
||
|
Just val -> val : subElems
|