module Data.Trie.HashMap where
import Data.Trie.Class
import Data.Monoid
import Data.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
import Data.Data
import GHC.Generics
import Control.DeepSeq
import Prelude hiding (lookup, null)
import Test.QuickCheck
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 (\n -> floor $ fromIntegral n / 2) arbitrary
instance ( Monoid (c p a)
) => Monoid (HashMapChildren c p a) where
mempty = HashMapChildren Nothing Nothing
mappend (HashMapChildren mx mxs) (HashMapChildren my mys) =
HashMapChildren (getLast $ Last mx <> Last my)
(mxs <> mys)
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 (floor $ fromIntegral n / 2) arbitrary
return $ 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
instance ( Hashable p
, Eq p
, Monoid (c p a)
) => Monoid (HashMapStep c p a) where
mempty = empty
mappend (HashMapStep xs) (HashMapStep ys) =
HashMapStep $ HM.unionWith (<>) xs ys
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
newtype HashMapTrie p a = HashMapTrie
{ unHashMapTrie :: HashMapStep HashMapTrie p a
} deriving (Show, Eq, Functor, Foldable, Traversable, 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 map (k :|) $ fromMaybe [] $ do
xs' <- mxs
return $ NE.toList <$> keys xs'
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
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 (p:|NE.toList pre, y, suff))
<> First mFoundHere
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 = fromMaybe [] $ (\x -> [(p:|[],x,ps)]) <$> mx
in if F.null ps
then foundHere
else let rs = fromMaybe [] $ matches (NE.fromList ps) <$> mxs
in foundHere ++ (prependAncestry <$> rs)
where prependAncestry (pre,x,suff) = (p:| NE.toList pre,x,suff)