graphql-engine/server/src-lib/Data/Trie.hs
Antoine Leblanc 0e3beb028d Extract generic containers from the codebase
### Description

There were several places in the codebase where we would either implement a generic container, or express the need for one. This PR extracts / creates all relevant containers, and adapts the relevant parts of the code to make use of said new generic containers. More specifically, it introduces the following modules:
- `Data.Set.Extended`, for new functions on `Data.Set`
- `Data.HashMap.Strict.Multi`, for hash maps that accept multiple values
- `Data.HashMap.Strict.NonEmpty`, for hash maps that can never be constructed as empty
- `Data.Trie`, for a generic implementation of a prefix tree

This PR makes use of those new containers in the following parts of the code:
- `Hasura.GraphQL.Execute.RemoteJoin.Types`
- `Hasura.RQL.Types.Endpoint*`

PR-URL: https://github.com/hasura/graphql-engine-mono/pull/3828
GitOrigin-RevId: e6c1b971bcb3f5ab66bc91d0fa4d0e9df7a0c6c6
2022-03-01 16:04:22 +00:00

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