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

{- |
Copyright: (c) 2018-2020 Kowainik
SPDX-License-Identifier: MPL-2.0
Maintainer: Kowainik <xrom.xkov@gmail.com>

Implementation of prefix tree for TOML AST.

@since 0.0.0
-}

module Toml.Type.PrefixTree
    (
      -- * 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.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
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)

-- | @since 0.3
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)


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

@since 0.0.0
-}
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.

@since 0.0.0
-}
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'.

@since 0.0.0
-}
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'.

@since 0.0.0
-}
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'.

@since 0.0.0
-}
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'.

@since 0.0.0
-}
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.

@since 0.0.0
-}
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.

@since 0.0.0
-}
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.

@since 0.0.0
-}
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