{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
-- | A trie from 'Text' keys to values.
module Text.Seonbi.Trie
    ( Trie
    , elems
    , empty
    , fromList
    , insert
    , keys
    , lookup
    , member
    , mergeBy
    , null
    , singleton
    , size
    , toList
    , unionL
    , unionR
    ) where

import Prelude hiding (lookup, null)

import Control.Monad (ap)
import qualified GHC.Exts

import Data.ByteString (ByteString)
import Data.Text hiding (empty, null, singleton)
import Data.Text.Encoding (encodeUtf8, decodeUtf8)
import qualified Data.Trie as BTrie

-- | A trie from 'Text' keys to 'a' values.
newtype Trie a
  = Trie (BTrie.Trie a)
  deriving (Trie a -> Trie a -> Bool
(Trie a -> Trie a -> Bool)
-> (Trie a -> Trie a -> Bool) -> Eq (Trie a)
forall a. Eq a => Trie a -> Trie a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Trie a -> Trie a -> Bool
$c/= :: forall a. Eq a => Trie a -> Trie a -> Bool
== :: Trie a -> Trie a -> Bool
$c== :: forall a. Eq a => Trie a -> Trie a -> Bool
Eq, Int -> Trie a -> ShowS
[Trie a] -> ShowS
Trie a -> String
(Int -> Trie a -> ShowS)
-> (Trie a -> String) -> ([Trie a] -> ShowS) -> Show (Trie a)
forall a. Show a => Int -> Trie a -> ShowS
forall a. Show a => [Trie a] -> ShowS
forall a. Show a => Trie a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Trie a] -> ShowS
$cshowList :: forall a. Show a => [Trie a] -> ShowS
show :: Trie a -> String
$cshow :: forall a. Show a => Trie a -> String
showsPrec :: Int -> Trie a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Trie a -> ShowS
Show)

encodeKey :: Text -> ByteString
encodeKey :: Text -> ByteString
encodeKey = Text -> ByteString
encodeUtf8

decodeKey :: ByteString -> Text
decodeKey :: ByteString -> Text
decodeKey = ByteString -> Text
decodeUtf8

-- | The empty trie.
empty :: Trie a
empty :: Trie a
empty = Trie a -> Trie a
forall a. Trie a -> Trie a
Trie Trie a
forall a. Trie a
BTrie.empty

-- | Checks if the trie is empty.
null :: Trie a -> Bool
null :: Trie a -> Bool
null (Trie Trie a
btrie) = Trie a -> Bool
forall a. Trie a -> Bool
BTrie.null Trie a
btrie

-- | Constructs a singleton trie.
singleton :: Text -> a -> Trie a
singleton :: Text -> a -> Trie a
singleton Text
key a
value = Trie a -> Trie a
forall a. Trie a -> Trie a
Trie (Trie a -> Trie a) -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ ByteString -> a -> Trie a
forall a. ByteString -> a -> Trie a
BTrie.singleton (Text -> ByteString
encodeKey Text
key) a
value

-- | Gets the number of elements in the trie.
size :: Trie a -> Int
size :: Trie a -> Int
size (Trie Trie a
btrie) = Trie a -> Int
forall a. Trie a -> Int
BTrie.size Trie a
btrie

fromList' :: [(Text, a)] -> Trie a
fromList' :: [(Text, a)] -> Trie a
fromList' [(Text, a)]
list = Trie a -> Trie a
forall a. Trie a -> Trie a
Trie (Trie a -> Trie a) -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ [(ByteString, a)] -> Trie a
forall a. [(ByteString, a)] -> Trie a
BTrie.fromList [(Text -> ByteString
encodeKey Text
k, a
v) | (Text
k, a
v) <- [(Text, a)]
list]

toList' :: Trie a -> [(Text, a)]
toList' :: Trie a -> [(Text, a)]
toList' (Trie Trie a
btrie) = [(ByteString -> Text
decodeKey ByteString
k, a
v) | (ByteString
k, a
v) <- Trie a -> [(ByteString, a)]
forall a. Trie a -> [(ByteString, a)]
BTrie.toList Trie a
btrie]

-- | Converts a list of associated pairs into a trie.  For duplicate keys,
-- values earlier in the list shadow later ones.
fromList :: [(Text, a)] -> Trie a
fromList :: [(Text, a)] -> Trie a
fromList = [(Text, a)] -> Trie a
forall a. [(Text, a)] -> Trie a
fromList'

-- | Converts a trie into a list of associated pairs.  Keys will be ordered.
toList :: Trie a -> [(Text, a)]
toList :: Trie a -> [(Text, a)]
toList = Trie a -> [(Text, a)]
forall a. Trie a -> [(Text, a)]
toList'

-- | Lists all keys in the trie.  Keys will be ordered.
keys :: Trie a -> [Text]
keys :: Trie a -> [Text]
keys (Trie Trie a
btrie) = (ByteString -> Text) -> [ByteString] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
Prelude.map ByteString -> Text
decodeKey ([ByteString] -> [Text]) -> [ByteString] -> [Text]
forall a b. (a -> b) -> a -> b
$ Trie a -> [ByteString]
forall a. Trie a -> [ByteString]
BTrie.keys Trie a
btrie

-- | Lists all values in the trie.  Values are ordered by their associated keys.
elems :: Trie a -> [a]
elems :: Trie a -> [a]
elems (Trie Trie a
btrie) = Trie a -> [a]
forall a. Trie a -> [a]
BTrie.elems Trie a
btrie

-- | Gets the value associated with a key if it exists.
lookup :: Text -> Trie a -> Maybe a
lookup :: Text -> Trie a -> Maybe a
lookup Text
key (Trie Trie a
btrie) = ByteString -> Trie a -> Maybe a
forall a. ByteString -> Trie a -> Maybe a
BTrie.lookup (Text -> ByteString
encodeKey Text
key) Trie a
btrie

-- | Checks if a key has a value in the trie.
member :: Text -> Trie a -> Bool
member :: Text -> Trie a -> Bool
member Text
key (Trie Trie a
btrie) = ByteString -> Trie a -> Bool
forall a. ByteString -> Trie a -> Bool
BTrie.member (Text -> ByteString
encodeKey Text
key) Trie a
btrie

-- | Inserts a new key into the trie.
insert
    :: Text
    -- ^ A new key to insert.  If there is already the same key in the trie,
    -- the existing value is overwritten by the new value.
    -> a
    -- ^ A value associated to the key.
    -> Trie a
    -- ^ An existing trie.
    -> Trie a
    -- ^ The new trie with the inserted key.
insert :: Text -> a -> Trie a -> Trie a
insert Text
key a
value (Trie Trie a
btrie) = Trie a -> Trie a
forall a. Trie a -> Trie a
Trie (Trie a -> Trie a) -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ ByteString -> a -> Trie a -> Trie a
forall a. ByteString -> a -> Trie a -> Trie a
BTrie.insert (Text -> ByteString
encodeKey Text
key) a
value Trie a
btrie

-- | Combines two tries, using a function to resolve collisions.  This can only
-- define the space of functions between union and symmetric difference but,
-- with those two, all set operations can be defined (albeit inefficiently).
mergeBy :: (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
mergeBy :: (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
mergeBy a -> a -> Maybe a
f (Trie Trie a
a) (Trie Trie a
b) = Trie a -> Trie a
forall a. Trie a -> Trie a
Trie (Trie a -> Trie a) -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
forall a. (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
BTrie.mergeBy a -> a -> Maybe a
f Trie a
a Trie a
b

-- | Combines two tries, resolving conflicts by choosing the value from the
-- left (former) trie.
unionL :: Trie a -> Trie a -> Trie a
unionL :: Trie a -> Trie a -> Trie a
unionL (Trie Trie a
left) (Trie Trie a
right) = Trie a -> Trie a
forall a. Trie a -> Trie a
Trie (Trie a -> Trie a) -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ Trie a -> Trie a -> Trie a
forall a. Trie a -> Trie a -> Trie a
BTrie.unionL Trie a
left Trie a
right

-- | Combines two tries, resolving conflicts by choosing the value from the
-- right (latter) trie.
unionR :: Trie a -> Trie a -> Trie a
unionR :: Trie a -> Trie a -> Trie a
unionR (Trie Trie a
left) (Trie Trie a
right) = Trie a -> Trie a
forall a. Trie a -> Trie a
Trie (Trie a -> Trie a) -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ Trie a -> Trie a -> Trie a
forall a. Trie a -> Trie a -> Trie a
BTrie.unionR Trie a
left Trie a
right

instance Functor Trie where
    fmap :: (a -> b) -> Trie a -> Trie b
fmap a -> b
f (Trie Trie a
btrie) = Trie b -> Trie b
forall a. Trie a -> Trie a
Trie (Trie b -> Trie b) -> Trie b -> Trie b
forall a b. (a -> b) -> a -> b
$ (a -> b) -> Trie a -> Trie b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Trie a
btrie

instance Foldable Trie where
    foldMap :: (a -> m) -> Trie a -> m
foldMap a -> m
f (Trie Trie a
btrie) = (a -> m) -> Trie a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f Trie a
btrie

instance Traversable Trie where
    traverse :: (a -> f b) -> Trie a -> f (Trie b)
traverse a -> f b
f (Trie Trie a
btrie) = Trie b -> Trie b
forall a. Trie a -> Trie a
Trie (Trie b -> Trie b) -> f (Trie b) -> f (Trie b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Trie a -> f (Trie b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Trie a
btrie

instance Applicative Trie where
    pure :: a -> Trie a
pure = a -> Trie a
forall (m :: * -> *) a. Monad m => a -> m a
return
    <*> :: Trie (a -> b) -> Trie a -> Trie b
(<*>) = Trie (a -> b) -> Trie a -> Trie b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad Trie where
    return :: a -> Trie a
return = Text -> a -> Trie a
forall a. Text -> a -> Trie a
singleton Text
""
    Trie Trie a
btrie >>= :: Trie a -> (a -> Trie b) -> Trie b
>>= a -> Trie b
f = Trie b -> Trie b
forall a. Trie a -> Trie a
Trie (Trie b -> Trie b) -> Trie b -> Trie b
forall a b. (a -> b) -> a -> b
$ Trie a
btrie Trie a -> (a -> Trie b) -> Trie b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\ a
v -> case a -> Trie b
f a
v of { Trie Trie b
b -> Trie b
b })

instance (Semigroup a) => Semigroup (Trie a) where
    (Trie Trie a
a) <> :: Trie a -> Trie a -> Trie a
<> (Trie Trie a
b) = Trie a -> Trie a
forall a. Trie a -> Trie a
Trie (Trie a
a Trie a -> Trie a -> Trie a
forall a. Semigroup a => a -> a -> a
<> Trie a
b)

instance (Monoid a) => Monoid (Trie a) where
    mempty :: Trie a
mempty = Trie a -> Trie a
forall a. Trie a -> Trie a
Trie Trie a
forall a. Monoid a => a
mempty
    mappend :: Trie a -> Trie a -> Trie a
mappend (Trie Trie a
a) (Trie Trie a
b) = Trie a -> Trie a
forall a. Trie a -> Trie a
Trie (Trie a -> Trie a) -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ Trie a -> Trie a -> Trie a
forall a. Monoid a => a -> a -> a
mappend Trie a
a Trie a
b

instance GHC.Exts.IsList (Trie a) where
    type Item (Trie a) = (Text, a)
    fromList :: [Item (Trie a)] -> Trie a
fromList = [Item (Trie a)] -> Trie a
forall a. [(Text, a)] -> Trie a
fromList'
    toList :: Trie a -> [Item (Trie a)]
toList = Trie a -> [Item (Trie a)]
forall a. Trie a -> [(Text, a)]
toList'