{-# LANGUAGE DeriveAnyClass  #-}
{-# LANGUAGE PatternSynonyms #-}

{- |
Module                  : Toml.Type.PrefixTree
Copyright               : (c) 2018-2022 Kowainik
SPDX-License-Identifier : MPL-2.0
Maintainer              : Kowainik <xrom.xkov@gmail.com>
Stability               : Stable
Portability             : Portable

Implementation of prefix tree for TOML AST.

@since 0.0.0
-}

module Toml.Type.PrefixTree
    (
      -- * Non-empty prefix tree
      PrefixTree (..)
    , singleT
    , insertT
    , lookupT
    , toListT
    , addPrefixT
    , differenceWithT

      -- * Prefix map that stores roots of 'PrefixTree'
    , PrefixMap
    , single
    , insert
    , lookup
    , fromList
    , toList
    , differenceWith
    ) where

import Prelude hiding (lookup)

import Control.DeepSeq (NFData)
import Data.Bifunctor (first)
import Data.Foldable (foldl')
import Data.HashMap.Strict (HashMap)
import GHC.Generics (Generic)

import Toml.Type.Key (pattern (:||), Key, KeysDiff (..), Piece, Prefix, keysDiff, (<|))

import qualified Data.HashMap.Strict as HashMap


{- | Map of layer names and corresponding 'PrefixTree's.

@since 0.0.0
-}
type PrefixMap a = HashMap Piece (PrefixTree a)

{- | Data structure to represent table tree for @toml@.

@since 0.0.0
-}
data PrefixTree a
    = Leaf             -- ^ End of a key.
        !Key           -- ^ End piece of the key.
        !a             -- ^ Value at the end.
    | Branch           -- ^ Values along pieces of a key.
        !Prefix        -- ^ Greatest common key prefix.
        !(Maybe a)     -- ^ Possible value at that point.
        !(PrefixMap a) -- ^ Values at suffixes of the prefix.
    deriving stock (Int -> PrefixTree a -> ShowS
[PrefixTree a] -> ShowS
PrefixTree a -> String
(Int -> PrefixTree a -> ShowS)
-> (PrefixTree a -> String)
-> ([PrefixTree a] -> ShowS)
-> Show (PrefixTree a)
forall a. Show a => Int -> PrefixTree a -> ShowS
forall a. Show a => [PrefixTree a] -> ShowS
forall a. Show a => PrefixTree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> PrefixTree a -> ShowS
showsPrec :: Int -> PrefixTree a -> ShowS
$cshow :: forall a. Show a => PrefixTree a -> String
show :: PrefixTree a -> String
$cshowList :: forall a. Show a => [PrefixTree a] -> ShowS
showList :: [PrefixTree a] -> ShowS
Show, PrefixTree a -> PrefixTree a -> Bool
(PrefixTree a -> PrefixTree a -> Bool)
-> (PrefixTree a -> PrefixTree a -> Bool) -> Eq (PrefixTree a)
forall a. Eq a => PrefixTree a -> PrefixTree a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => PrefixTree a -> PrefixTree a -> Bool
== :: PrefixTree a -> PrefixTree a -> Bool
$c/= :: forall a. Eq a => PrefixTree a -> PrefixTree a -> Bool
/= :: PrefixTree a -> PrefixTree a -> Bool
Eq, (forall x. PrefixTree a -> Rep (PrefixTree a) x)
-> (forall x. Rep (PrefixTree a) x -> PrefixTree a)
-> Generic (PrefixTree a)
forall x. Rep (PrefixTree a) x -> PrefixTree a
forall x. PrefixTree a -> Rep (PrefixTree a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (PrefixTree a) x -> PrefixTree a
forall a x. PrefixTree a -> Rep (PrefixTree a) x
$cfrom :: forall a x. PrefixTree a -> Rep (PrefixTree a) x
from :: forall x. PrefixTree a -> Rep (PrefixTree a) x
$cto :: forall a x. Rep (PrefixTree a) x -> PrefixTree a
to :: forall x. Rep (PrefixTree a) x -> PrefixTree a
Generic)
    deriving anyclass (PrefixTree a -> ()
(PrefixTree a -> ()) -> NFData (PrefixTree a)
forall a. NFData a => PrefixTree a -> ()
forall a. (a -> ()) -> NFData a
$crnf :: forall a. NFData a => PrefixTree a -> ()
rnf :: PrefixTree a -> ()
NFData)

-- | @since 0.3
instance Semigroup (PrefixTree a) where
    PrefixTree a
a <> :: PrefixTree a -> PrefixTree a -> PrefixTree a
<> PrefixTree a
b = (PrefixTree a -> (Key, a) -> PrefixTree a)
-> PrefixTree a -> [(Key, a)] -> PrefixTree a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\PrefixTree a
tree (Key
k, a
v) -> Key -> a -> PrefixTree a -> PrefixTree a
forall a. Key -> a -> PrefixTree a -> PrefixTree a
insertT Key
k a
v PrefixTree a
tree) PrefixTree a
a (PrefixTree a -> [(Key, a)]
forall a. PrefixTree a -> [(Key, a)]
toListT PrefixTree a
b)

{- | Push 'Prefix' inside the given 'PrefixTree'.

@since 1.3.2.0
-}
addPrefixT :: Prefix -> PrefixTree a -> PrefixTree a
addPrefixT :: forall a. Key -> PrefixTree a -> PrefixTree a
addPrefixT Key
pref = \case
    Leaf Key
k a
a -> Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
Leaf (Key
pref Key -> Key -> Key
forall a. Semigroup a => a -> a -> a
<> Key
k) a
a
    Branch Key
k Maybe a
ma PrefixMap a
pma -> Key -> Maybe a -> PrefixMap a -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch (Key
pref Key -> Key -> Key
forall a. Semigroup a => a -> a -> a
<> Key
k) Maybe a
ma PrefixMap a
pma

{- | Convert branches to 'Leaf' or remove them at all.

@since 1.3.2.0
-}
compressTree :: PrefixTree a -> Maybe (PrefixTree a)
compressTree :: forall a. PrefixTree a -> Maybe (PrefixTree a)
compressTree = \case
    l :: PrefixTree a
l@(Leaf Key
_ a
_) -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
l
    b :: PrefixTree a
b@(Branch Key
p Maybe a
ma PrefixMap a
pma) -> case PrefixMap a -> [(Piece, PrefixTree a)]
forall k v. HashMap k v -> [(k, v)]
HashMap.toList PrefixMap a
pma of
        [] -> Maybe a
ma Maybe a -> (a -> Maybe (PrefixTree a)) -> Maybe (PrefixTree a)
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
a -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just (Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
Leaf Key
p a
a)
        [(Piece
_, PrefixTree a
child)] -> case Maybe a
ma of
            Just a
_ -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
b
            Maybe a
Nothing -> PrefixTree a -> Maybe (PrefixTree a)
forall a. PrefixTree a -> Maybe (PrefixTree a)
compressTree (PrefixTree a -> Maybe (PrefixTree a))
-> PrefixTree a -> Maybe (PrefixTree a)
forall a b. (a -> b) -> a -> b
$ Key -> PrefixTree a -> PrefixTree a
forall a. Key -> PrefixTree a -> PrefixTree a
addPrefixT Key
p PrefixTree a
child
        (Piece, PrefixTree a)
_ : (Piece, PrefixTree a)
_ : [(Piece, PrefixTree a)]
_ -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
b

{- | Creates a 'PrefixTree' of one key-value element.

@since 0.0.0
-}
singleT :: Key -> a -> PrefixTree a
singleT :: forall a. Key -> a -> PrefixTree a
singleT = Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
Leaf
{-# INLINE singleT #-}

{- | Creates a 'PrefixMap' of one key-value element.

@since 0.0.0
-}
single :: Key -> a -> PrefixMap a
single :: forall a. Key -> a -> PrefixMap a
single k :: Key
k@(Piece
p :|| [Piece]
_) = Piece -> PrefixTree a -> HashMap Piece (PrefixTree a)
forall k v. Hashable k => k -> v -> HashMap k v
HashMap.singleton Piece
p (PrefixTree a -> HashMap Piece (PrefixTree a))
-> (a -> PrefixTree a) -> a -> HashMap Piece (PrefixTree a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
singleT Key
k

{- | Inserts key-value element into the given 'PrefixTree'.

@since 0.0.0
-}
insertT :: Key -> a -> PrefixTree a -> PrefixTree a
insertT :: forall a. Key -> a -> PrefixTree a -> PrefixTree a
insertT Key
newK a
newV (Leaf Key
k a
v) =
    case Key -> Key -> KeysDiff
keysDiff Key
k Key
newK of
        KeysDiff
Equal -> Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
Leaf Key
k a
newV
        KeysDiff
NoPrefix -> String -> PrefixTree a
forall a. HasCallStack => String -> a
error String
"Algorithm error: can't be equal prefixes in insertT:Leaf case"
        FstIsPref Key
rK -> Key -> Maybe a -> PrefixMap a -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
k (a -> Maybe a
forall a. a -> Maybe a
Just a
v) (PrefixMap a -> PrefixTree a) -> PrefixMap a -> PrefixTree a
forall a b. (a -> b) -> a -> b
$ Key -> a -> PrefixMap a
forall a. Key -> a -> PrefixMap a
single Key
rK a
newV
        SndIsPref Key
lK -> Key -> Maybe a -> PrefixMap a -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
newK (a -> Maybe a
forall a. a -> Maybe a
Just a
newV) (PrefixMap a -> PrefixTree a) -> PrefixMap a -> PrefixTree a
forall a b. (a -> b) -> a -> b
$ Key -> a -> PrefixMap a
forall a. Key -> a -> PrefixMap a
single Key
lK a
v
        Diff Key
p k1 :: Key
k1@(Piece
f :|| [Piece]
_) k2 :: Key
k2@(Piece
s :|| [Piece]
_) ->
          Key -> Maybe a -> PrefixMap a -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
p Maybe a
forall a. Maybe a
Nothing (PrefixMap a -> PrefixTree a) -> PrefixMap a -> PrefixTree a
forall a b. (a -> b) -> a -> b
$ [(Piece, PrefixTree a)] -> PrefixMap a
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList [(Piece
f, Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
Leaf Key
k1 a
v), (Piece
s, Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
Leaf Key
k2 a
newV)]
insertT Key
newK a
newV (Branch Key
pref Maybe a
mv PrefixMap a
prefMap) =
    case Key -> Key -> KeysDiff
keysDiff Key
pref Key
newK of
        KeysDiff
Equal    -> Key -> Maybe a -> PrefixMap a -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
pref (a -> Maybe a
forall a. a -> Maybe a
Just a
newV) PrefixMap a
prefMap
        KeysDiff
NoPrefix -> String -> PrefixTree a
forall a. HasCallStack => String -> a
error String
"Algorithm error: can't be equal prefixes in insertT:Branch case"
        FstIsPref Key
rK -> Key -> Maybe a -> PrefixMap a -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
pref Maybe a
mv (PrefixMap a -> PrefixTree a) -> PrefixMap a -> PrefixTree a
forall a b. (a -> b) -> a -> b
$ Key -> a -> PrefixMap a -> PrefixMap a
forall a. Key -> a -> PrefixMap a -> PrefixMap a
insert Key
rK a
newV PrefixMap a
prefMap
        SndIsPref lK :: Key
lK@(Piece
piece :|| [Piece]
_) ->
            Key -> Maybe a -> PrefixMap a -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
newK (a -> Maybe a
forall a. a -> Maybe a
Just a
newV) (PrefixMap a -> PrefixTree a) -> PrefixMap a -> PrefixTree a
forall a b. (a -> b) -> a -> b
$ Piece -> PrefixTree a -> PrefixMap a
forall k v. Hashable k => k -> v -> HashMap k v
HashMap.singleton Piece
piece (Key -> Maybe a -> PrefixMap a -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
lK Maybe a
mv PrefixMap a
prefMap)
        Diff Key
p k1 :: Key
k1@(Piece
f :|| [Piece]
_) k2 :: Key
k2@(Piece
s :|| [Piece]
_) ->
            Key -> Maybe a -> PrefixMap a -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
p Maybe a
forall a. Maybe a
Nothing (PrefixMap a -> PrefixTree a) -> PrefixMap a -> PrefixTree a
forall a b. (a -> b) -> a -> b
$ [(Piece, PrefixTree a)] -> PrefixMap a
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList [ (Piece
f, Key -> Maybe a -> PrefixMap a -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
k1 Maybe a
mv PrefixMap a
prefMap)
                                                , (Piece
s, Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
Leaf Key
k2 a
newV)
                                                ]

{- | Inserts key-value element into the given 'PrefixMap'.

@since 0.0.0
-}
insert :: Key -> a -> PrefixMap a -> PrefixMap a
insert :: forall a. Key -> a -> PrefixMap a -> PrefixMap a
insert k :: Key
k@(Piece
p :|| [Piece]
_) a
v PrefixMap a
prefMap = case Piece -> PrefixMap a -> Maybe (PrefixTree a)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Piece
p PrefixMap a
prefMap of
    Just PrefixTree a
tree -> Piece -> PrefixTree a -> PrefixMap a -> PrefixMap a
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert Piece
p (Key -> a -> PrefixTree a -> PrefixTree a
forall a. Key -> a -> PrefixTree a -> PrefixTree a
insertT Key
k a
v PrefixTree a
tree) PrefixMap a
prefMap
    Maybe (PrefixTree a)
Nothing   -> Piece -> PrefixTree a -> PrefixMap a -> PrefixMap a
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert Piece
p (Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
singleT Key
k a
v) PrefixMap a
prefMap

{- | Looks up the value at a key in the 'PrefixTree'.

@since 0.0.0
-}
lookupT :: Key -> PrefixTree a -> Maybe a
lookupT :: forall a. Key -> PrefixTree a -> Maybe a
lookupT Key
lk (Leaf Key
k a
v) = if Key
lk Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
k then a -> Maybe a
forall a. a -> Maybe a
Just a
v else Maybe a
forall a. Maybe a
Nothing
lookupT Key
lk (Branch Key
pref Maybe a
mv PrefixMap a
prefMap) =
    case Key -> Key -> KeysDiff
keysDiff Key
pref Key
lk of
        KeysDiff
Equal       -> Maybe a
mv
        KeysDiff
NoPrefix    -> Maybe a
forall a. Maybe a
Nothing
        Diff Key
_ Key
_ Key
_  -> Maybe a
forall a. Maybe a
Nothing
        SndIsPref Key
_ -> Maybe a
forall a. Maybe a
Nothing
        FstIsPref Key
k -> Key -> PrefixMap a -> Maybe a
forall a. Key -> PrefixMap a -> Maybe a
lookup Key
k PrefixMap a
prefMap

{- | Looks up the value at a key in the 'PrefixMap'.

@since 0.0.0
-}
lookup :: Key -> PrefixMap a -> Maybe a
lookup :: forall a. Key -> PrefixMap a -> Maybe a
lookup k :: Key
k@(Piece
p :|| [Piece]
_) PrefixMap a
prefMap = Piece -> PrefixMap a -> Maybe (PrefixTree a)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Piece
p PrefixMap a
prefMap Maybe (PrefixTree a) -> (PrefixTree a -> Maybe a) -> Maybe a
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Key -> PrefixTree a -> Maybe a
forall a. Key -> PrefixTree a -> Maybe a
lookupT Key
k

{- | Constructs 'PrefixMap' structure from the given list of 'Key' and value pairs.

@since 0.0.0
-}
fromList :: [(Key, a)] -> PrefixMap a
fromList :: forall a. [(Key, a)] -> PrefixMap a
fromList = (PrefixMap a -> (Key, a) -> PrefixMap a)
-> PrefixMap a -> [(Key, a)] -> PrefixMap a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' PrefixMap a -> (Key, a) -> PrefixMap a
forall a. PrefixMap a -> (Key, a) -> PrefixMap a
insertPair PrefixMap a
forall a. Monoid a => a
mempty
  where
    insertPair :: PrefixMap a -> (Key, a) -> PrefixMap a
    insertPair :: forall a. PrefixMap a -> (Key, a) -> PrefixMap a
insertPair PrefixMap a
prefMap (Key
k, a
v) = Key -> a -> PrefixMap a -> PrefixMap a
forall a. Key -> a -> PrefixMap a -> PrefixMap a
insert Key
k a
v PrefixMap a
prefMap

{- | Converts 'PrefixTree' to the list of pairs.

@since 0.0.0
-}
toListT :: PrefixTree a -> [(Key, a)]
toListT :: forall a. PrefixTree a -> [(Key, a)]
toListT (Leaf Key
k a
v) = [(Key
k, a
v)]
toListT (Branch Key
pref Maybe a
ma PrefixMap a
prefMap) = case Maybe a
ma of
    Just a
a  -> (:) (Key
pref, a
a)
    Maybe a
Nothing -> [(Key, a)] -> [(Key, a)]
forall a. a -> a
id
    ([(Key, a)] -> [(Key, a)]) -> [(Key, a)] -> [(Key, a)]
forall a b. (a -> b) -> a -> b
$ ((Key, a) -> (Key, a)) -> [(Key, a)] -> [(Key, a)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Key
k, a
v) -> (Key
pref Key -> Key -> Key
forall a. Semigroup a => a -> a -> a
<> Key
k, a
v)) ([(Key, a)] -> [(Key, a)]) -> [(Key, a)] -> [(Key, a)]
forall a b. (a -> b) -> a -> b
$ PrefixMap a -> [(Key, a)]
forall a. PrefixMap a -> [(Key, a)]
toList PrefixMap a
prefMap

{- | Converts 'PrefixMap' to the list of pairs.

@since 0.0.0
-}
toList :: PrefixMap a -> [(Key, a)]
toList :: forall a. PrefixMap a -> [(Key, a)]
toList = ((Piece, PrefixTree a) -> [(Key, a)])
-> [(Piece, PrefixTree a)] -> [(Key, a)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(Piece
p, PrefixTree a
tr) -> (Key -> Key) -> (Key, a) -> (Key, a)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Piece
p Piece -> Key -> Key
<|) ((Key, a) -> (Key, a)) -> [(Key, a)] -> [(Key, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PrefixTree a -> [(Key, a)]
forall a. PrefixTree a -> [(Key, a)]
toListT PrefixTree a
tr) ([(Piece, PrefixTree a)] -> [(Key, a)])
-> (PrefixMap a -> [(Piece, PrefixTree a)])
-> PrefixMap a
-> [(Key, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrefixMap a -> [(Piece, PrefixTree a)]
forall k v. HashMap k v -> [(k, v)]
HashMap.toList

{- | Difference of two 'PrefixMap's. Returns elements of the first 'PrefixMap'
that are not existing in the second one.

@since 1.3.2.0
-}
differenceWith :: (a -> b -> Maybe a) -> PrefixMap a -> PrefixMap b -> PrefixMap a
differenceWith :: forall a b.
(a -> b -> Maybe a) -> PrefixMap a -> PrefixMap b -> PrefixMap a
differenceWith a -> b -> Maybe a
f = (PrefixTree a -> PrefixTree b -> Maybe (PrefixTree a))
-> HashMap Piece (PrefixTree a)
-> HashMap Piece (PrefixTree b)
-> HashMap Piece (PrefixTree a)
forall k v w.
(Eq k, Hashable k) =>
(v -> w -> Maybe v) -> HashMap k v -> HashMap k w -> HashMap k v
HashMap.differenceWith ((a -> b -> Maybe a)
-> PrefixTree a -> PrefixTree b -> Maybe (PrefixTree a)
forall a b.
(a -> b -> Maybe a)
-> PrefixTree a -> PrefixTree b -> Maybe (PrefixTree a)
differenceWithT a -> b -> Maybe a
f)

{- | Difference of two 'PrefixTree's. Returns elements of the first 'PrefixTree'
that are not existing in the second one.

@since 1.3.2.0
-}
differenceWithT :: (a -> b -> Maybe a) -> PrefixTree a -> PrefixTree b -> Maybe (PrefixTree a)
differenceWithT :: forall a b.
(a -> b -> Maybe a)
-> PrefixTree a -> PrefixTree b -> Maybe (PrefixTree a)
differenceWithT a -> b -> Maybe a
f PrefixTree a
pt1 PrefixTree b
pt2 = case (PrefixTree a
pt1, PrefixTree b
pt2) of
    (Leaf Key
k1 a
a, Leaf Key
k2 b
b)
        | Key
k1 Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
k2 -> a -> b -> Maybe a
f a
a b
b Maybe a -> (a -> Maybe (PrefixTree a)) -> Maybe (PrefixTree a)
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
aNew -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just (Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
Leaf Key
k1 a
aNew)
        | Bool
otherwise -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just (Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
Leaf Key
k1 a
a)

    (l :: PrefixTree a
l@(Leaf Key
k a
a), Branch Key
p Maybe b
mb PrefixMap b
pmb) -> case Key -> Key -> KeysDiff
keysDiff Key
k Key
p of
        KeysDiff
Equal -> Maybe b
mb Maybe b -> (b -> Maybe a) -> Maybe a
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> b -> Maybe a
f a
a Maybe a -> (a -> Maybe (PrefixTree a)) -> Maybe (PrefixTree a)
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
aNew -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just (Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
Leaf Key
k a
aNew)
        KeysDiff
NoPrefix -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
l
        FstIsPref Key
_ -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
l
        SndIsPref Key
kSuf -> case HashMap Piece (PrefixTree a) -> [(Piece, PrefixTree a)]
forall k v. HashMap k v -> [(k, v)]
HashMap.toList (HashMap Piece (PrefixTree a) -> [(Piece, PrefixTree a)])
-> HashMap Piece (PrefixTree a) -> [(Piece, PrefixTree a)]
forall a b. (a -> b) -> a -> b
$ (a -> b -> Maybe a)
-> HashMap Piece (PrefixTree a)
-> PrefixMap b
-> HashMap Piece (PrefixTree a)
forall a b.
(a -> b -> Maybe a) -> PrefixMap a -> PrefixMap b -> PrefixMap a
differenceWith a -> b -> Maybe a
f (Key -> a -> HashMap Piece (PrefixTree a)
forall a. Key -> a -> PrefixMap a
single Key
kSuf a
a) PrefixMap b
pmb of
            -- zero elements
            [] -> Maybe (PrefixTree a)
forall a. Maybe a
Nothing
            -- our single key
            [(Piece
_, PrefixTree a
aNew)] -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just (PrefixTree a -> Maybe (PrefixTree a))
-> PrefixTree a -> Maybe (PrefixTree a)
forall a b. (a -> b) -> a -> b
$ Key -> PrefixTree a -> PrefixTree a
forall a. Key -> PrefixTree a -> PrefixTree a
addPrefixT Key
k PrefixTree a
aNew
            -- shouldn't happen, but for some reasons
            (Piece, PrefixTree a)
_ : (Piece, PrefixTree a)
_ : [(Piece, PrefixTree a)]
_ -> Maybe (PrefixTree a)
forall a. Maybe a
Nothing
        Diff Key
_ Key
_ Key
_ -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
l

    (br :: PrefixTree a
br@(Branch Key
p Maybe a
ma HashMap Piece (PrefixTree a)
pma), Leaf Key
k b
b) -> case Key -> Key -> KeysDiff
keysDiff Key
p Key
k of
        KeysDiff
Equal -> PrefixTree a -> Maybe (PrefixTree a)
forall a. PrefixTree a -> Maybe (PrefixTree a)
compressTree (PrefixTree a -> Maybe (PrefixTree a))
-> PrefixTree a -> Maybe (PrefixTree a)
forall a b. (a -> b) -> a -> b
$ Key -> Maybe a -> HashMap Piece (PrefixTree a) -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
p (Maybe a
ma Maybe a -> (a -> Maybe a) -> Maybe a
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
a -> a -> b -> Maybe a
f a
a b
b) HashMap Piece (PrefixTree a)
pma
        KeysDiff
NoPrefix -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
br
        FstIsPref Key
kSuf -> PrefixTree a -> Maybe (PrefixTree a)
forall a. PrefixTree a -> Maybe (PrefixTree a)
compressTree (PrefixTree a -> Maybe (PrefixTree a))
-> PrefixTree a -> Maybe (PrefixTree a)
forall a b. (a -> b) -> a -> b
$ Key -> Maybe a -> HashMap Piece (PrefixTree a) -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
p Maybe a
ma ((a -> b -> Maybe a)
-> HashMap Piece (PrefixTree a)
-> PrefixMap b
-> HashMap Piece (PrefixTree a)
forall a b.
(a -> b -> Maybe a) -> PrefixMap a -> PrefixMap b -> PrefixMap a
differenceWith a -> b -> Maybe a
f HashMap Piece (PrefixTree a)
pma (PrefixMap b -> HashMap Piece (PrefixTree a))
-> PrefixMap b -> HashMap Piece (PrefixTree a)
forall a b. (a -> b) -> a -> b
$ Key -> b -> PrefixMap b
forall a. Key -> a -> PrefixMap a
single Key
kSuf b
b)
        SndIsPref Key
_ -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
br
        Diff Key
_ Key
_ Key
_ -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
br

    (b1 :: PrefixTree a
b1@(Branch Key
p1 Maybe a
ma HashMap Piece (PrefixTree a)
pma), Branch Key
p2 Maybe b
mb PrefixMap b
pmb) -> case Key -> Key -> KeysDiff
keysDiff Key
p1 Key
p2 of
        KeysDiff
Equal -> PrefixTree a -> Maybe (PrefixTree a)
forall a. PrefixTree a -> Maybe (PrefixTree a)
compressTree (PrefixTree a -> Maybe (PrefixTree a))
-> PrefixTree a -> Maybe (PrefixTree a)
forall a b. (a -> b) -> a -> b
$
            Key -> Maybe a -> HashMap Piece (PrefixTree a) -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
p1 (Maybe a
ma Maybe a -> (a -> Maybe a) -> Maybe a
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
a -> Maybe b
mb Maybe b -> (b -> Maybe a) -> Maybe a
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
b -> a -> b -> Maybe a
f a
a b
b) ((a -> b -> Maybe a)
-> HashMap Piece (PrefixTree a)
-> PrefixMap b
-> HashMap Piece (PrefixTree a)
forall a b.
(a -> b -> Maybe a) -> PrefixMap a -> PrefixMap b -> PrefixMap a
differenceWith a -> b -> Maybe a
f HashMap Piece (PrefixTree a)
pma PrefixMap b
pmb)
        KeysDiff
NoPrefix -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
b1
        FstIsPref p2Suf :: Key
p2Suf@(Piece
p2Head :|| [Piece]
_) -> PrefixTree a -> Maybe (PrefixTree a)
forall a. PrefixTree a -> Maybe (PrefixTree a)
compressTree (PrefixTree a -> Maybe (PrefixTree a))
-> PrefixTree a -> Maybe (PrefixTree a)
forall a b. (a -> b) -> a -> b
$
            Key -> Maybe a -> HashMap Piece (PrefixTree a) -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
p1 Maybe a
ma ((a -> b -> Maybe a)
-> HashMap Piece (PrefixTree a)
-> PrefixMap b
-> HashMap Piece (PrefixTree a)
forall a b.
(a -> b -> Maybe a) -> PrefixMap a -> PrefixMap b -> PrefixMap a
differenceWith a -> b -> Maybe a
f HashMap Piece (PrefixTree a)
pma (PrefixMap b -> HashMap Piece (PrefixTree a))
-> PrefixMap b -> HashMap Piece (PrefixTree a)
forall a b. (a -> b) -> a -> b
$ Piece -> PrefixTree b -> PrefixMap b
forall k v. Hashable k => k -> v -> HashMap k v
HashMap.singleton Piece
p2Head (PrefixTree b -> PrefixMap b) -> PrefixTree b -> PrefixMap b
forall a b. (a -> b) -> a -> b
$ Key -> Maybe b -> PrefixMap b -> PrefixTree b
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
p2Suf Maybe b
mb PrefixMap b
pmb)
        SndIsPref p1Suf :: Key
p1Suf@(Piece
p1Head :|| [Piece]
_) -> case Piece -> PrefixMap b -> Maybe (PrefixTree b)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Piece
p1Head PrefixMap b
pmb of
            Maybe (PrefixTree b)
Nothing -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
b1
            Just PrefixTree b
ch -> Key -> PrefixTree a -> PrefixTree a
forall a. Key -> PrefixTree a -> PrefixTree a
addPrefixT Key
p2 (PrefixTree a -> PrefixTree a)
-> Maybe (PrefixTree a) -> Maybe (PrefixTree a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> b -> Maybe a)
-> PrefixTree a -> PrefixTree b -> Maybe (PrefixTree a)
forall a b.
(a -> b -> Maybe a)
-> PrefixTree a -> PrefixTree b -> Maybe (PrefixTree a)
differenceWithT a -> b -> Maybe a
f (Key -> Maybe a -> HashMap Piece (PrefixTree a) -> PrefixTree a
forall a. Key -> Maybe a -> PrefixMap a -> PrefixTree a
Branch Key
p1Suf Maybe a
ma HashMap Piece (PrefixTree a)
pma) PrefixTree b
ch
        Diff Key
_ Key
_ Key
_ -> PrefixTree a -> Maybe (PrefixTree a)
forall a. a -> Maybe a
Just PrefixTree a
b1