{-# OPTIONS_GHC -Wall -fwarn-tabs #-}
{-# LANGUAGE NoImplicitPrelude, CPP #-}
#if __GLASGOW_HASKELL__ >= 701
{-# LANGUAGE Safe #-}
#endif
----------------------------------------------------------------
--                                                  ~ 2021.12.14
-- |
-- Module      :  Data.Trie
-- Copyright   :  2008--2021 wren romano
-- License     :  BSD-3-Clause
-- Maintainer  :  wren@cpan.org
-- Stability   :  provisional
-- Portability :  portable
--
-- An efficient implementation of finite maps from strings to values.
-- The implementation is based on /big-endian patricia trees/, like
-- "Data.IntMap". We first trie on the elements of "Data.ByteString"
-- and then trie on the big-endian bit representation of those
-- elements. For further details on the latter, see
--
--    * Chris Okasaki and Andy Gill,  \"/Fast Mergeable Integer Maps/\",
--    Workshop on ML, September 1998, pages 77-86,
--    <http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.37.5452>
--
--    * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve/
--    /Information Coded In Alphanumeric/\", Journal of the ACM, 15(4),
--    October 1968, pages 514-534.
--
-- This module aims to provide an austere interface, while being
-- detailed enough for most users. For an extended interface with
-- many additional functions, see "Data.Trie.Convenience". For
-- functions that give more detailed (potentially abstraction-breaking)
-- access to the data strucuture, or for experimental functions
-- which aren't quite ready for the public API, see "Data.Trie.Internal".
----------------------------------------------------------------

module Data.Trie
    (
    -- * Data type
      Trie()

    -- * Basic functions
    , empty, null, singleton, size

    -- * Conversion functions
    , fromList, toListBy, toList, keys, elems

    -- * Query functions
    , lookupBy, lookup, member, submap, match, minMatch, matches

    -- * Simple modification
    , insert, adjust, adjustBy, alterBy, delete, deleteSubmap

    -- * Combining tries
    , mergeBy, unionL, unionR
    , intersectBy, intersectL, intersectR

    -- * Mapping functions
    , mapBy, filterMap
    -- TODO: re-export 'filter', 'filterA', 'wither', etc?
    ) where

import Prelude hiding     (null, lookup)

import Data.Trie.Internal
import Data.ByteString    (ByteString)
import qualified Data.ByteString as S
import Data.Maybe         (isJust)
----------------------------------------------------------------
----------------------------------------------------------------


{---------------------------------------------------------------
-- Conversion functions
---------------------------------------------------------------}

-- | Return all keys in the trie, in sorted order.
--
-- __Warning__: This function suffers <Data-Trie-Internal.html#bug25 Bug #25>.
keys :: Trie a -> [ByteString]
{-# INLINE keys #-}
keys :: Trie a -> [ByteString]
keys  = (ByteString -> a -> ByteString) -> Trie a -> [ByteString]
forall a b. (ByteString -> a -> b) -> Trie a -> [b]
toListBy ByteString -> a -> ByteString
forall a b. a -> b -> a
const



{---------------------------------------------------------------
-- Query functions (just recurse)
---------------------------------------------------------------}

-- | Generic function to find a value (if it exists) and the subtrie
-- rooted at the prefix.
lookupBy :: (Maybe a -> Trie a -> b) -> ByteString -> Trie a -> b
{-# INLINE lookupBy #-}
lookupBy :: (Maybe a -> Trie a -> b) -> ByteString -> Trie a -> b
lookupBy Maybe a -> Trie a -> b
f = (a -> Trie a -> b)
-> (Trie a -> b) -> b -> ByteString -> Trie a -> b
forall a b.
(a -> Trie a -> b)
-> (Trie a -> b) -> b -> ByteString -> Trie a -> b
lookupBy_ (Maybe a -> Trie a -> b
f (Maybe a -> Trie a -> b) -> (a -> Maybe a) -> a -> Trie a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe a
forall a. a -> Maybe a
Just) (Maybe a -> Trie a -> b
f Maybe a
forall a. Maybe a
Nothing) (Maybe a -> Trie a -> b
f Maybe a
forall a. Maybe a
Nothing Trie a
forall a. Trie a
empty)

-- | Return the value associated with a query string if it exists.
lookup :: ByteString -> Trie a -> Maybe a
{-# INLINE lookup #-}
lookup :: ByteString -> Trie a -> Maybe a
lookup = (Maybe a -> Trie a -> Maybe a) -> ByteString -> Trie a -> Maybe a
forall a b. (Maybe a -> Trie a -> b) -> ByteString -> Trie a -> b
lookupBy Maybe a -> Trie a -> Maybe a
forall a b. a -> b -> a
const

-- | Does a string have a value in the trie?
member :: ByteString -> Trie a -> Bool
{-# INLINE member #-}
member :: ByteString -> Trie a -> Bool
member ByteString
q = Maybe a -> Bool
forall a. Maybe a -> Bool
isJust (Maybe a -> Bool) -> (Trie a -> Maybe a) -> Trie a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Trie a -> Maybe a
forall a. ByteString -> Trie a -> Maybe a
lookup ByteString
q

getMatch :: ByteString -> (Int, a) -> (ByteString, a, ByteString)
{-# INLINE getMatch #-}
getMatch :: ByteString -> (Int, a) -> (ByteString, a, ByteString)
getMatch ByteString
q (Int
n,a
x) =
    case Int -> ByteString -> (ByteString, ByteString)
S.splitAt Int
n ByteString
q of
    (ByteString
p,ByteString
q') -> (ByteString
p, a
x, ByteString
q')

-- | Given a query, find the longest prefix with an associated value
-- in the trie, and return that prefix, it's value, and the remainder
-- of the query.
--
-- @since 0.2.4
match :: Trie a -> ByteString -> Maybe (ByteString, a, ByteString)
{-# INLINE match #-}
match :: Trie a -> ByteString -> Maybe (ByteString, a, ByteString)
match Trie a
t ByteString
q = ByteString -> (Int, a) -> (ByteString, a, ByteString)
forall a. ByteString -> (Int, a) -> (ByteString, a, ByteString)
getMatch ByteString
q ((Int, a) -> (ByteString, a, ByteString))
-> Maybe (Int, a) -> Maybe (ByteString, a, ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Trie a -> ByteString -> Maybe (Int, a)
forall a. Trie a -> ByteString -> Maybe (Int, a)
match_ Trie a
t ByteString
q

-- | Given a query, find the shortest prefix with an associated value
-- in the trie, and return that prefix, it's value, and the remainder
-- of the query.
--
-- @since 0.2.6
minMatch :: Trie a -> ByteString -> Maybe (ByteString, a, ByteString)
{-# INLINE minMatch #-}
minMatch :: Trie a -> ByteString -> Maybe (ByteString, a, ByteString)
minMatch Trie a
t ByteString
q =
    case Trie a -> ByteString -> [(ByteString, a, ByteString)]
forall a. Trie a -> ByteString -> [(ByteString, a, ByteString)]
matches Trie a
t ByteString
q of
    []  -> Maybe (ByteString, a, ByteString)
forall a. Maybe a
Nothing
    (ByteString, a, ByteString)
x:[(ByteString, a, ByteString)]
_ -> (ByteString, a, ByteString) -> Maybe (ByteString, a, ByteString)
forall a. a -> Maybe a
Just (ByteString, a, ByteString)
x

-- | Given a query, find all prefixes with associated values in the
-- trie, and return their (prefix, value, remainder) triples in
-- order from shortest prefix to longest.  This function is a good
-- producer for list fusion.
--
-- @since 0.2.4
matches :: Trie a -> ByteString -> [(ByteString, a, ByteString)]
{-# INLINE matches #-}
matches :: Trie a -> ByteString -> [(ByteString, a, ByteString)]
matches Trie a
t ByteString
q = ByteString -> (Int, a) -> (ByteString, a, ByteString)
forall a. ByteString -> (Int, a) -> (ByteString, a, ByteString)
getMatch ByteString
q ((Int, a) -> (ByteString, a, ByteString))
-> [(Int, a)] -> [(ByteString, a, ByteString)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Trie a -> ByteString -> [(Int, a)]
forall a. Trie a -> ByteString -> [(Int, a)]
matches_ Trie a
t ByteString
q


{---------------------------------------------------------------
-- Simple modification functions (recurse and clone spine)
---------------------------------------------------------------}

-- | Insert a new key. If the key is already present, overrides the
-- old value
insert :: ByteString -> a -> Trie a -> Trie a
{-# INLINE insert #-}
insert :: ByteString -> a -> Trie a -> Trie a
insert = (ByteString -> a -> Maybe a -> Maybe a)
-> ByteString -> a -> Trie a -> Trie a
forall a.
(ByteString -> a -> Maybe a -> Maybe a)
-> ByteString -> a -> Trie a -> Trie a
alterBy (\ByteString
_ a
x Maybe a
_ -> a -> Maybe a
forall a. a -> Maybe a
Just a
x)

-- | Alter the value associated with a given key. If the key is not
-- present, then the trie is returned unaltered. See 'alterBy' if
-- you are interested in inserting new keys or deleting old keys.
-- Because this function does not need to worry about changing the
-- trie structure, it is somewhat faster than 'alterBy'.
--
-- @since 0.2.6
-- __Note__: Prior to version 0.2.6 this function was exported from
-- "Data.Trie.Internal" instead.
adjustBy :: (ByteString -> a -> a -> a)
         -> ByteString -> a -> Trie a -> Trie a
{-# INLINE adjustBy #-}
adjustBy :: (ByteString -> a -> a -> a) -> ByteString -> a -> Trie a -> Trie a
adjustBy ByteString -> a -> a -> a
f ByteString
q a
x = (a -> a) -> ByteString -> Trie a -> Trie a
forall a. (a -> a) -> ByteString -> Trie a -> Trie a
adjust (ByteString -> a -> a -> a
f ByteString
q a
x) ByteString
q

-- | Remove the value stored at a key.
delete :: ByteString -> Trie a -> Trie a
{-# INLINE delete #-}
delete :: ByteString -> Trie a -> Trie a
delete = (Maybe a -> Trie a -> (Maybe a, Trie a))
-> ByteString -> Trie a -> Trie a
forall a.
(Maybe a -> Trie a -> (Maybe a, Trie a))
-> ByteString -> Trie a -> Trie a
alterBy_ (\Maybe a
_ Trie a
t -> (Maybe a
forall a. Maybe a
Nothing, Trie a
t))

-- | Remove all keys beginning with a prefix.
--
-- @since 0.2.6
deleteSubmap :: ByteString -> Trie a -> Trie a
{-# INLINE deleteSubmap #-}
deleteSubmap :: ByteString -> Trie a -> Trie a
deleteSubmap = (Maybe a -> Trie a -> (Maybe a, Trie a))
-> ByteString -> Trie a -> Trie a
forall a.
(Maybe a -> Trie a -> (Maybe a, Trie a))
-> ByteString -> Trie a -> Trie a
alterBy_ (\Maybe a
_ Trie a
_ -> (Maybe a
forall a. Maybe a
Nothing, Trie a
forall a. Trie a
empty))


{---------------------------------------------------------------
-- Trie-combining functions
---------------------------------------------------------------}

-- | Take the union of two tries, resolving conflicts by choosing
-- the value from the left trie.
unionL :: Trie a -> Trie a -> Trie a
{-# INLINE unionL #-}
unionL :: Trie a -> Trie a -> Trie a
unionL = (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
forall a. (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
mergeBy (\a
x a
_ -> a -> Maybe a
forall a. a -> Maybe a
Just a
x)

-- | Take the union of two tries, resolving conflicts by choosing
-- the value from the right trie.
unionR :: Trie a -> Trie a -> Trie a
{-# INLINE unionR #-}
unionR :: Trie a -> Trie a -> Trie a
unionR = (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
forall a. (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
mergeBy (\a
_ a
y -> a -> Maybe a
forall a. a -> Maybe a
Just a
y)

-- | Take the intersection of two tries, with values from the left trie.
--
-- @since 0.2.6
intersectL :: Trie a -> Trie b -> Trie a
{-# INLINE intersectL #-}
intersectL :: Trie a -> Trie b -> Trie a
intersectL = (a -> b -> Maybe a) -> Trie a -> Trie b -> Trie a
forall a b c. (a -> b -> Maybe c) -> Trie a -> Trie b -> Trie c
intersectBy (\a
x b
_ -> a -> Maybe a
forall a. a -> Maybe a
Just a
x)

-- | Take the intersection of two tries, with values from the right trie.
--
-- @since 0.2.6
intersectR :: Trie a -> Trie b -> Trie b
{-# INLINE intersectR #-}
intersectR :: Trie a -> Trie b -> Trie b
intersectR = (a -> b -> Maybe b) -> Trie a -> Trie b -> Trie b
forall a b c. (a -> b -> Maybe c) -> Trie a -> Trie b -> Trie c
intersectBy (\a
_ b
y -> b -> Maybe b
forall a. a -> Maybe a
Just b
y)

----------------------------------------------------------------
----------------------------------------------------------- fin.