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

-- | Implementation of prefix tree for TOML AST.

module Toml.PrefixTree
       ( -- * Core types
         Piece (..)
       , Key (..)
       , Prefix
       , pattern (:||)

         -- * Key difference
       , KeysDiff (..)
       , keysDiff

         -- * Non-empty prefix tree
       , PrefixTree (..)
       , (<|)
       , singleT
       , insertT
       , lookupT
       , toListT

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

import Prelude hiding (lookup)

import Control.DeepSeq (NFData)
import Data.Bifunctor (first)
import Data.Coerce (coerce)
import Data.Foldable (foldl')
import Data.Hashable (Hashable)
import Data.HashMap.Strict (HashMap)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Semigroup (Semigroup (..))
import Data.String (IsString (..))
import Data.Text (Text)
import GHC.Generics (Generic)

import qualified Data.HashMap.Strict as HashMap
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Text as Text


-- | Represents the key piece of some layer.
newtype Piece = Piece
    { Piece -> Text
unPiece :: Text
    } deriving stock ((forall x. Piece -> Rep Piece x)
-> (forall x. Rep Piece x -> Piece) -> Generic Piece
forall x. Rep Piece x -> Piece
forall x. Piece -> Rep Piece x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Piece x -> Piece
$cfrom :: forall x. Piece -> Rep Piece x
Generic)
      deriving newtype (Int -> Piece -> ShowS
[Piece] -> ShowS
Piece -> String
(Int -> Piece -> ShowS)
-> (Piece -> String) -> ([Piece] -> ShowS) -> Show Piece
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Piece] -> ShowS
$cshowList :: [Piece] -> ShowS
show :: Piece -> String
$cshow :: Piece -> String
showsPrec :: Int -> Piece -> ShowS
$cshowsPrec :: Int -> Piece -> ShowS
Show, Piece -> Piece -> Bool
(Piece -> Piece -> Bool) -> (Piece -> Piece -> Bool) -> Eq Piece
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Piece -> Piece -> Bool
$c/= :: Piece -> Piece -> Bool
== :: Piece -> Piece -> Bool
$c== :: Piece -> Piece -> Bool
Eq, Eq Piece
Eq Piece =>
(Piece -> Piece -> Ordering)
-> (Piece -> Piece -> Bool)
-> (Piece -> Piece -> Bool)
-> (Piece -> Piece -> Bool)
-> (Piece -> Piece -> Bool)
-> (Piece -> Piece -> Piece)
-> (Piece -> Piece -> Piece)
-> Ord Piece
Piece -> Piece -> Bool
Piece -> Piece -> Ordering
Piece -> Piece -> Piece
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Piece -> Piece -> Piece
$cmin :: Piece -> Piece -> Piece
max :: Piece -> Piece -> Piece
$cmax :: Piece -> Piece -> Piece
>= :: Piece -> Piece -> Bool
$c>= :: Piece -> Piece -> Bool
> :: Piece -> Piece -> Bool
$c> :: Piece -> Piece -> Bool
<= :: Piece -> Piece -> Bool
$c<= :: Piece -> Piece -> Bool
< :: Piece -> Piece -> Bool
$c< :: Piece -> Piece -> Bool
compare :: Piece -> Piece -> Ordering
$ccompare :: Piece -> Piece -> Ordering
$cp1Ord :: Eq Piece
Ord, Int -> Piece -> Int
Piece -> Int
(Int -> Piece -> Int) -> (Piece -> Int) -> Hashable Piece
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: Piece -> Int
$chash :: Piece -> Int
hashWithSalt :: Int -> Piece -> Int
$chashWithSalt :: Int -> Piece -> Int
Hashable, String -> Piece
(String -> Piece) -> IsString Piece
forall a. (String -> a) -> IsString a
fromString :: String -> Piece
$cfromString :: String -> Piece
IsString, Piece -> ()
(Piece -> ()) -> NFData Piece
forall a. (a -> ()) -> NFData a
rnf :: Piece -> ()
$crnf :: Piece -> ()
NFData)

{- | Key of value in @key = val@ pair. Represents as non-empty list of key
components — 'Piece's. Key like

@
site."google.com"
@

is represented like

@
Key (Piece "site" :| [Piece "\\"google.com\\""])
@
-}
newtype Key = Key
    { Key -> NonEmpty Piece
unKey :: NonEmpty Piece
    } deriving stock ((forall x. Key -> Rep Key x)
-> (forall x. Rep Key x -> Key) -> Generic Key
forall x. Rep Key x -> Key
forall x. Key -> Rep Key x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Key x -> Key
$cfrom :: forall x. Key -> Rep Key x
Generic)
      deriving newtype (Int -> Key -> ShowS
[Key] -> ShowS
Key -> String
(Int -> Key -> ShowS)
-> (Key -> String) -> ([Key] -> ShowS) -> Show Key
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Key] -> ShowS
$cshowList :: [Key] -> ShowS
show :: Key -> String
$cshow :: Key -> String
showsPrec :: Int -> Key -> ShowS
$cshowsPrec :: Int -> Key -> ShowS
Show, Key -> Key -> Bool
(Key -> Key -> Bool) -> (Key -> Key -> Bool) -> Eq Key
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Key -> Key -> Bool
$c/= :: Key -> Key -> Bool
== :: Key -> Key -> Bool
$c== :: Key -> Key -> Bool
Eq, Eq Key
Eq Key =>
(Key -> Key -> Ordering)
-> (Key -> Key -> Bool)
-> (Key -> Key -> Bool)
-> (Key -> Key -> Bool)
-> (Key -> Key -> Bool)
-> (Key -> Key -> Key)
-> (Key -> Key -> Key)
-> Ord Key
Key -> Key -> Bool
Key -> Key -> Ordering
Key -> Key -> Key
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Key -> Key -> Key
$cmin :: Key -> Key -> Key
max :: Key -> Key -> Key
$cmax :: Key -> Key -> Key
>= :: Key -> Key -> Bool
$c>= :: Key -> Key -> Bool
> :: Key -> Key -> Bool
$c> :: Key -> Key -> Bool
<= :: Key -> Key -> Bool
$c<= :: Key -> Key -> Bool
< :: Key -> Key -> Bool
$c< :: Key -> Key -> Bool
compare :: Key -> Key -> Ordering
$ccompare :: Key -> Key -> Ordering
$cp1Ord :: Eq Key
Ord, Int -> Key -> Int
Key -> Int
(Int -> Key -> Int) -> (Key -> Int) -> Hashable Key
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: Key -> Int
$chash :: Key -> Int
hashWithSalt :: Int -> Key -> Int
$chashWithSalt :: Int -> Key -> Int
Hashable, Key -> ()
(Key -> ()) -> NFData Key
forall a. (a -> ()) -> NFData a
rnf :: Key -> ()
$crnf :: Key -> ()
NFData, b -> Key -> Key
NonEmpty Key -> Key
Key -> Key -> Key
(Key -> Key -> Key)
-> (NonEmpty Key -> Key)
-> (forall b. Integral b => b -> Key -> Key)
-> Semigroup Key
forall b. Integral b => b -> Key -> Key
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: b -> Key -> Key
$cstimes :: forall b. Integral b => b -> Key -> Key
sconcat :: NonEmpty Key -> Key
$csconcat :: NonEmpty Key -> Key
<> :: Key -> Key -> Key
$c<> :: Key -> Key -> Key
Semigroup)

{- | Split a dot-separated string into 'Key'. Empty string turns into a 'Key'
with single element — empty 'Piece'.

This instance is not safe for now. Use carefully. If you try to use as a key
string like this @site.\"google.com\"@ you will have list of three components
instead of desired two.
-}
instance IsString Key where
    fromString :: String -> Key
    fromString :: String -> Key
fromString = \case
        "" -> NonEmpty Piece -> Key
Key ("" Piece -> [Piece] -> NonEmpty Piece
forall a. a -> [a] -> NonEmpty a
:| [])
        s :: String
s  -> case Text -> Text -> [Text]
Text.splitOn "." (String -> Text
forall a. IsString a => String -> a
fromString String
s) of
            []   -> String -> Key
forall a. HasCallStack => String -> a
error "Text.splitOn returned empty string"  -- can't happen
            x :: Text
x:xs :: [Text]
xs -> NonEmpty Text -> Key
forall a b. Coercible a b => a -> b
coerce @(NonEmpty Text) @Key (Text
x Text -> [Text] -> NonEmpty Text
forall a. a -> [a] -> NonEmpty a
:| [Text]
xs)

{- | Bidirectional pattern synonym for constructing and deconstructing 'Key's.
-}
pattern (:||) :: Piece -> [Piece] -> Key
pattern x $b:|| :: Piece -> [Piece] -> Key
$m:|| :: forall r. Key -> (Piece -> [Piece] -> r) -> (Void# -> r) -> r
:|| xs <- Key (x :| xs)
  where
    x :: Piece
x :|| xs :: [Piece]
xs = NonEmpty Piece -> Key
Key (Piece
x Piece -> [Piece] -> NonEmpty Piece
forall a. a -> [a] -> NonEmpty a
:| [Piece]
xs)

{-# COMPLETE (:||) #-}

-- | Type synonym for 'Key'.
type Prefix = Key

-- | Map of layer names and corresponding 'PrefixTree's.
type PrefixMap a = HashMap Piece (PrefixTree a)

-- | Data structure to represent table tree for @toml@.
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
showList :: [PrefixTree a] -> ShowS
$cshowList :: forall a. Show a => [PrefixTree a] -> ShowS
show :: PrefixTree a -> String
$cshow :: forall a. Show a => PrefixTree a -> String
showsPrec :: Int -> PrefixTree a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> 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
/= :: PrefixTree a -> PrefixTree a -> Bool
$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
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
$cto :: forall a x. Rep (PrefixTree a) x -> PrefixTree a
$cfrom :: forall a x. PrefixTree a -> Rep (PrefixTree a) x
Generic)
    deriving anyclass (PrefixTree a -> ()
(PrefixTree a -> ()) -> NFData (PrefixTree a)
forall a. NFData a => PrefixTree a -> ()
forall a. (a -> ()) -> NFData a
rnf :: PrefixTree a -> ()
$crnf :: forall a. NFData a => PrefixTree a -> ()
NFData)

instance Semigroup (PrefixTree a) where
    a :: PrefixTree a
a <> :: PrefixTree a -> PrefixTree a -> PrefixTree a
<> b :: PrefixTree a
b = (PrefixTree a -> (Key, a) -> PrefixTree a)
-> PrefixTree a -> [(Key, a)] -> PrefixTree a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\tree :: PrefixTree a
tree (k :: Key
k, v :: 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)

-- | Data structure to compare keys.
data KeysDiff
    = Equal      -- ^ Keys are equal
    | NoPrefix   -- ^ Keys don't have any common part.
    | FstIsPref  -- ^ The first key is the prefix of the second one.
        !Key     -- ^ Rest of the second key.
    | SndIsPref  -- ^ The second key is the prefix of the first one.
        !Key     -- ^ Rest of the first key.
    | Diff       -- ^ Key have a common prefix.
        !Key     -- ^ Common prefix.
        !Key     -- ^ Rest of the first key.
        !Key     -- ^ Rest of the second key.
    deriving stock (Int -> KeysDiff -> ShowS
[KeysDiff] -> ShowS
KeysDiff -> String
(Int -> KeysDiff -> ShowS)
-> (KeysDiff -> String) -> ([KeysDiff] -> ShowS) -> Show KeysDiff
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [KeysDiff] -> ShowS
$cshowList :: [KeysDiff] -> ShowS
show :: KeysDiff -> String
$cshow :: KeysDiff -> String
showsPrec :: Int -> KeysDiff -> ShowS
$cshowsPrec :: Int -> KeysDiff -> ShowS
Show, KeysDiff -> KeysDiff -> Bool
(KeysDiff -> KeysDiff -> Bool)
-> (KeysDiff -> KeysDiff -> Bool) -> Eq KeysDiff
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: KeysDiff -> KeysDiff -> Bool
$c/= :: KeysDiff -> KeysDiff -> Bool
== :: KeysDiff -> KeysDiff -> Bool
$c== :: KeysDiff -> KeysDiff -> Bool
Eq)

-- | Compares two keys
keysDiff :: Key -> Key -> KeysDiff
keysDiff :: Key -> Key -> KeysDiff
keysDiff (x :: Piece
x :|| xs :: [Piece]
xs) (y :: Piece
y :|| ys :: [Piece]
ys)
    | Piece
x Piece -> Piece -> Bool
forall a. Eq a => a -> a -> Bool
== Piece
y    = [Piece] -> [Piece] -> [Piece] -> KeysDiff
listSame [Piece]
xs [Piece]
ys []
    | Bool
otherwise = KeysDiff
NoPrefix
  where
    listSame :: [Piece] -> [Piece] -> [Piece] -> KeysDiff
    listSame :: [Piece] -> [Piece] -> [Piece] -> KeysDiff
listSame [] []     _ = KeysDiff
Equal
    listSame [] (s :: Piece
s:ss :: [Piece]
ss) _ = Key -> KeysDiff
FstIsPref (Key -> KeysDiff) -> Key -> KeysDiff
forall a b. (a -> b) -> a -> b
$ Piece
s Piece -> [Piece] -> Key
:|| [Piece]
ss
    listSame (f :: Piece
f:fs :: [Piece]
fs) [] _ = Key -> KeysDiff
SndIsPref (Key -> KeysDiff) -> Key -> KeysDiff
forall a b. (a -> b) -> a -> b
$ Piece
f Piece -> [Piece] -> Key
:|| [Piece]
fs
    listSame (f :: Piece
f:fs :: [Piece]
fs) (s :: Piece
s:ss :: [Piece]
ss) pr :: [Piece]
pr =
        if Piece
f Piece -> Piece -> Bool
forall a. Eq a => a -> a -> Bool
== Piece
s
        then [Piece] -> [Piece] -> [Piece] -> KeysDiff
listSame [Piece]
fs [Piece]
ss ([Piece]
pr [Piece] -> [Piece] -> [Piece]
forall a. [a] -> [a] -> [a]
++ [Piece
f])
        else Key -> Key -> Key -> KeysDiff
Diff (Piece
x Piece -> [Piece] -> Key
:|| [Piece]
pr) (Piece
f Piece -> [Piece] -> Key
:|| [Piece]
fs) (Piece
s Piece -> [Piece] -> Key
:|| [Piece]
ss)

-- | Prepends 'Piece' to the beginning of the 'Key'.
(<|) :: Piece -> Key -> Key
<| :: Piece -> Key -> Key
(<|) p :: Piece
p k :: Key
k = NonEmpty Piece -> Key
Key (Piece
p Piece -> NonEmpty Piece -> NonEmpty Piece
forall a. a -> NonEmpty a -> NonEmpty a
NonEmpty.<| Key -> NonEmpty Piece
unKey Key
k)
{-# INLINE (<|) #-}

-- | Creates a 'PrefixTree' of one key-value element.
singleT :: Key -> a -> PrefixTree a
singleT :: 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.
single :: Key -> a -> PrefixMap a
single :: Key -> a -> PrefixMap a
single k :: Key
k@(p :: Piece
p :|| _) = Piece -> PrefixTree a -> PrefixMap a
forall k v. Hashable k => k -> v -> HashMap k v
HashMap.singleton Piece
p (PrefixTree a -> PrefixMap a)
-> (a -> PrefixTree a) -> a -> PrefixMap 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'.
insertT :: Key -> a -> PrefixTree a -> PrefixTree a
insertT :: Key -> a -> PrefixTree a -> PrefixTree a
insertT newK :: Key
newK newV :: a
newV (Leaf k :: Key
k v :: a
v) =
    case Key -> Key -> KeysDiff
keysDiff Key
k Key
newK of
        Equal -> Key -> a -> PrefixTree a
forall a. Key -> a -> PrefixTree a
Leaf Key
k a
newV
        NoPrefix -> String -> PrefixTree a
forall a. HasCallStack => String -> a
error "Algorithm error: can't be equal prefixes in insertT:Leaf case"
        FstIsPref rK :: 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 lK :: 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 p :: Key
p k1 :: Key
k1@(f :: Piece
f :|| _) k2 :: Key
k2@(s :: Piece
s :|| _) ->
          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 newK :: Key
newK newV :: a
newV (Branch pref :: Key
pref mv :: Maybe a
mv prefMap :: PrefixMap a
prefMap) =
    case Key -> Key -> KeysDiff
keysDiff Key
pref Key
newK of
        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
        NoPrefix -> String -> PrefixTree a
forall a. HasCallStack => String -> a
error "Algorithm error: can't be equal prefixes in insertT:Branch case"
        FstIsPref rK :: 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 p :: Key
p k1 :: Key
k1@(f :: Piece
f :|| _) k2 :: Key
k2@(s :: Piece
s :|| _) ->
            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'.
insert :: Key -> a -> PrefixMap a -> PrefixMap a
insert :: Key -> a -> PrefixMap a -> PrefixMap a
insert k :: Key
k@(p :: Piece
p :|| _) v :: a
v prefMap :: 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 tree :: 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
    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'.
lookupT :: Key -> PrefixTree a -> Maybe a
lookupT :: Key -> PrefixTree a -> Maybe a
lookupT lk :: Key
lk (Leaf k :: Key
k v :: 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 lk :: Key
lk (Branch pref :: Key
pref mv :: Maybe a
mv prefMap :: PrefixMap a
prefMap) =
    case Key -> Key -> KeysDiff
keysDiff Key
pref Key
lk of
        Equal       -> Maybe a
mv
        NoPrefix    -> Maybe a
forall a. Maybe a
Nothing
        Diff _ _ _  -> Maybe a
forall a. Maybe a
Nothing
        SndIsPref _ -> Maybe a
forall a. Maybe a
Nothing
        FstIsPref k :: 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'.
lookup :: Key -> PrefixMap a -> Maybe a
lookup :: Key -> PrefixMap a -> Maybe a
lookup k :: Key
k@(p :: Piece
p :|| _) prefMap :: 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 (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.
fromList :: [(Key, a)] -> PrefixMap a
fromList :: [(Key, a)] -> PrefixMap a
fromList = (PrefixMap a -> (Key, a) -> PrefixMap a)
-> PrefixMap a -> [(Key, a)] -> PrefixMap a
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 :: PrefixMap a -> (Key, a) -> PrefixMap a
insertPair prefMap :: PrefixMap a
prefMap (k :: Key
k, v :: 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.
toListT :: PrefixTree a -> [(Key, a)]
toListT :: PrefixTree a -> [(Key, a)]
toListT (Leaf k :: Key
k v :: a
v) = [(Key
k, a
v)]
toListT (Branch pref :: Key
pref ma :: Maybe a
ma prefMap :: PrefixMap a
prefMap) = case Maybe a
ma of
    Just a :: a
a  -> (:) (Key
pref, a
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 (\(k :: Key
k, v :: 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.
toList :: PrefixMap a -> [(Key, a)]
toList :: 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 (\(p :: Piece
p, tr :: PrefixTree a
tr) -> (Key -> Key) -> (Key, a) -> (Key, a)
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