{-# LANGUAGE
GeneralizedNewtypeDeriving
, DeriveFunctor
, DeriveFoldable
, DeriveTraversable
, DeriveGeneric
, DeriveDataTypeable
, MultiParamTypeClasses
, FlexibleInstances
, FlexibleContexts
, UndecidableInstances
, TypeFamilies
, TupleSections
#-}
module Data.Trie.HashMap where
import Data.Trie.Class (Trie (..))
import Data.Semigroup (Semigroup)
import Data.Monoid (First (..), Last (..), (<>))
import Data.Hashable (Hashable)
import Data.Maybe (fromMaybe, fromJust)
import qualified Data.Foldable as F
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NE
import qualified Data.HashMap.Lazy as HM
import qualified Data.Key as K
import Control.Monad (replicateM)
import Data.Data (Data, Typeable)
import GHC.Generics (Generic)
import Control.DeepSeq (NFData)
import Prelude hiding (lookup, null)
import Test.QuickCheck (Arbitrary (..), resize, choose, sized, scale)
import Test.QuickCheck.Instances ()
data HashMapChildren c p a = HashMapChildren
{ hashMapNode :: !(Maybe a)
, hashMapChildren :: !(Maybe (c p a))
} deriving (Show, Eq, Functor, Foldable, Traversable, Generic, Data, Typeable)
instance ( NFData (c p a)
, NFData p
, NFData a
) => NFData (HashMapChildren c p a)
instance ( Arbitrary a
, Arbitrary p
, Arbitrary (c p a)
) => Arbitrary (HashMapChildren c p a) where
arbitrary = HashMapChildren <$> arbitrary <*> scale (`div` 2) arbitrary
instance ( Semigroup (c p a)
) => Semigroup (HashMapChildren c p a) where
(HashMapChildren mx mxs) <> (HashMapChildren my mys) =
HashMapChildren (getLast (Last mx <> Last my))
(mxs <> mys)
instance ( Monoid (c p a)
) => Monoid (HashMapChildren c p a) where
mempty = HashMapChildren Nothing Nothing
newtype HashMapStep c p a = HashMapStep
{ unHashMapStep :: HM.HashMap p (HashMapChildren c p a)
} deriving (Show, Eq, Functor, Foldable, Traversable, Generic, Data, Typeable)
instance ( NFData (c p a)
, NFData p
, NFData a
) => NFData (HashMapStep c p a)
instance ( Arbitrary a
, Arbitrary p
, Arbitrary (c p a)
, Hashable p
, Eq p
) => Arbitrary (HashMapStep c p a) where
arbitrary = sized go
where
go n = do
i <- choose (0,n)
xs <- replicateM i $ (,) <$> arbitrary <*> resize (n `div` 2) arbitrary
pure (HashMapStep (HM.fromList xs))
instance ( Hashable p
, Eq p
, Trie NonEmpty p c
) => Trie NonEmpty p (HashMapStep c) where
lookup (p:|ps) (HashMapStep xs)
| F.null ps = hashMapNode
=<< HM.lookup p xs
| otherwise = lookup (NE.fromList ps)
=<< hashMapChildren
=<< HM.lookup p xs
delete (p:|ps) (HashMapStep xs)
| F.null ps = let mxs = hashMapChildren =<< HM.lookup p xs
in HashMapStep (HM.insert p (HashMapChildren Nothing mxs) xs)
| otherwise = let (HashMapChildren mx mxs) =
fromMaybe (HashMapChildren Nothing Nothing)
(HM.lookup p xs)
in HashMapStep (HM.insert p
(HashMapChildren mx (delete (NE.fromList ps) <$> mxs))
xs)
insert :: ( Hashable p
, Eq p
, Trie NonEmpty p c
, Monoid (c p a)
) => NonEmpty p -> a -> HashMapStep c p a -> HashMapStep c p a
insert (p:|ps) x (HashMapStep xs)
| F.null ps = let mxs = hashMapChildren =<< HM.lookup p xs
in HashMapStep (HM.insert p
(HashMapChildren (Just x) $! mxs)
xs)
| otherwise = let mx = hashMapNode =<< HM.lookup p xs
xs' = fromMaybe mempty (hashMapChildren =<< HM.lookup p xs)
in HashMapStep (HM.insert p
(HashMapChildren mx
(Just (Data.Trie.Class.insert (NE.fromList ps) x xs')))
xs)
{-# INLINEABLE insert #-}
instance ( Hashable p
, Eq p
, Semigroup (c p a)
) => Semigroup (HashMapStep c p a) where
(HashMapStep xs) <> (HashMapStep ys) =
HashMapStep (HM.unionWith (<>) xs ys)
instance ( Hashable p
, Eq p
, Monoid (c p a)
) => Monoid (HashMapStep c p a) where
mempty = empty
empty :: HashMapStep c p a
empty = HashMapStep HM.empty
singleton :: Hashable p => p -> a -> HashMapStep c p a
singleton p x = HashMapStep (HM.singleton p (HashMapChildren (Just x) Nothing))
{-# INLINEABLE singleton #-}
newtype HashMapTrie p a = HashMapTrie
{ unHashMapTrie :: HashMapStep HashMapTrie p a
} deriving (Show, Eq, Functor, Foldable, Traversable, Semigroup, Monoid, Arbitrary)
instance (Hashable p, Eq p) => Trie NonEmpty p HashMapTrie where
lookup ts (HashMapTrie xs) = lookup ts xs
delete ts (HashMapTrie xs) = HashMapTrie (delete ts xs)
insert ts x (HashMapTrie xs) = HashMapTrie (Data.Trie.HashMap.insert ts x xs)
type instance K.Key (HashMapTrie p) = NonEmpty p
instance ( Hashable p
, Eq p
) => K.Lookup (HashMapTrie p) where
lookup = lookup
keys :: ( Hashable p
, Eq p
) => HashMapTrie p a -> [NonEmpty p]
keys (HashMapTrie (HashMapStep xs)) =
let ks = HM.keys xs
in F.concatMap go ks
where
go k = let (HashMapChildren _ mxs) = fromJust (HM.lookup k xs)
in case mxs of
Nothing -> []
Just xs' -> NE.cons k <$> keys xs'
{-# INLINEABLE keys #-}
elems :: HashMapTrie p a -> [a]
elems = F.toList
subtrie :: ( Hashable p
, Eq p
) => NonEmpty p -> HashMapTrie p a -> Maybe (HashMapTrie p a)
subtrie (p:|ps) (HashMapTrie (HashMapStep xs))
| F.null ps = hashMapChildren =<< HM.lookup p xs
| otherwise = subtrie (NE.fromList ps) =<< hashMapChildren =<< HM.lookup p xs
{-# INLINEABLE subtrie #-}
match :: ( Hashable p
, Eq p
) => NonEmpty p -> HashMapTrie p a -> Maybe (NonEmpty p, a, [p])
match (p:|ps) (HashMapTrie (HashMapStep xs)) = do
(HashMapChildren mx mxs) <- HM.lookup p xs
let mFoundHere = (p:|[],,ps) <$> mx
if F.null ps
then mFoundHere
else getFirst $ First (do (pre,y,suff) <- match (NE.fromList ps) =<< mxs
pure (NE.cons p pre, y, suff))
<> First mFoundHere
{-# INLINEABLE match #-}
matches :: ( Hashable p
, Eq p
) => NonEmpty p -> HashMapTrie p a -> [(NonEmpty p, a, [p])]
matches (p:|ps) (HashMapTrie (HashMapStep xs)) =
let (HashMapChildren mx mxs) = fromMaybe mempty (HM.lookup p xs)
foundHere = case mx of
Nothing -> []
Just x -> [(p:|[],x,ps)]
in if F.null ps
then foundHere
else let rs = case mxs of
Nothing -> []
Just xs' -> matches (NE.fromList ps) xs'
in foundHere ++ (prependAncestry <$> rs)
where prependAncestry (pre,x,suff) = (NE.cons p pre,x,suff)
{-# INLINEABLE matches #-}