-- | A basic trie over byte strings.
module Zenacy.HTML.Internal.Trie
  ( Trie
  , empty
  , fromList
  , insert
  , insertWords
  , match
  ) where

import Zenacy.HTML.Internal.BS
import Data.Map
  ( Map
  )
import qualified Data.Map as Map
  ( empty
  , fromList
  , insert
  , lookup
  )
import Data.Word
  ( Word8
  )

-- | Defines the tree.
data Trie a = Trie (Maybe a) (Map Word8 (Trie a)) deriving (Int -> Trie a -> ShowS
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)

-- | Creates an empty trie.
empty :: Trie a
empty :: forall a. Trie a
empty = forall a. Maybe a -> Map Word8 (Trie a) -> Trie a
Trie forall a. Maybe a
Nothing forall k a. Map k a
Map.empty

-- | Creates a trie from a list of tuples containing key and value.
fromList :: [(BS,a)] -> Trie a
fromList :: forall a. [(BS, a)] -> Trie a
fromList = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Trie a
t (BS
x,a
y) -> forall a. BS -> a -> Trie a -> Trie a
insert BS
x a
y Trie a
t) forall a. Trie a
empty

-- | Finds the longest prefix with a value in the trie and returns
-- the prefix, the value, and the remaining string.
match :: Trie a -> BS -> Maybe (BS, a, BS)
match :: forall a. Trie a -> BS -> Maybe (BS, a, BS)
match Trie a
t BS
s =
  case forall {t} {b}.
Num t =>
Maybe (t, b) -> t -> Trie b -> BS -> Maybe (t, b)
go forall a. Maybe a
Nothing Int
0 Trie a
t BS
s of
    Maybe (Int, a)
Nothing -> forall a. Maybe a
Nothing
    Just (Int
n, a
v) -> forall a. a -> Maybe a
Just (Int -> BS -> BS
bsTake Int
n BS
s, a
v, Int -> BS -> BS
bsDrop Int
n BS
s)
  where
    go :: Maybe (t, b) -> t -> Trie b -> BS -> Maybe (t, b)
go Maybe (t, b)
a t
n (Trie Maybe b
v Map Word8 (Trie b)
m) BS
s =
      case BS -> Maybe (Word8, BS)
bsUncons BS
s of
        Maybe (Word8, BS)
Nothing -> Maybe (t, b)
a
        Just (Word8
w, BS
t) ->
          case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Word8
w Map Word8 (Trie b)
m of
            Maybe (Trie b)
Nothing -> Maybe (t, b)
a
            Just b :: Trie b
b@(Trie (Just b
v2) Map Word8 (Trie b)
_) ->
              Maybe (t, b) -> t -> Trie b -> BS -> Maybe (t, b)
go (forall a. a -> Maybe a
Just (t
n forall a. Num a => a -> a -> a
+ t
1, b
v2)) (t
n forall a. Num a => a -> a -> a
+ t
1) Trie b
b BS
t
            Just Trie b
b ->
              Maybe (t, b) -> t -> Trie b -> BS -> Maybe (t, b)
go Maybe (t, b)
a (t
n forall a. Num a => a -> a -> a
+ t
1) Trie b
b BS
t

-- | Inserts a value into a trie.
insert :: BS -> a -> Trie a -> Trie a
insert :: forall a. BS -> a -> Trie a -> Trie a
insert BS
x = forall a. [Word8] -> a -> Trie a -> Trie a
insertWords forall a b. (a -> b) -> a -> b
$ BS -> [Word8]
bsUnpack BS
x

-- | Inserts a value into a trie.
insertWords :: [Word8] -> a -> Trie a -> Trie a
insertWords :: forall a. [Word8] -> a -> Trie a -> Trie a
insertWords [Word8]
x a
y = [Word8] -> Trie a -> Trie a
go [Word8]
x
  where
    go :: [Word8] -> Trie a -> Trie a
go [] (Trie Maybe a
v Map Word8 (Trie a)
m) =
      forall a. Maybe a -> Map Word8 (Trie a) -> Trie a
Trie (forall a. a -> Maybe a
Just a
y) Map Word8 (Trie a)
m
    go (Word8
w:[Word8]
ws) (Trie Maybe a
v Map Word8 (Trie a)
m) =
      case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Word8
w Map Word8 (Trie a)
m of
        Maybe (Trie a)
Nothing ->
          forall a. Maybe a -> Map Word8 (Trie a) -> Trie a
Trie Maybe a
v forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Word8
w ([Word8] -> Trie a -> Trie a
go [Word8]
ws forall a. Trie a
empty) Map Word8 (Trie a)
m
        Just Trie a
b ->
          forall a. Maybe a -> Map Word8 (Trie a) -> Trie a
Trie Maybe a
v forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Word8
w ([Word8] -> Trie a -> Trie a
go [Word8]
ws Trie a
b) Map Word8 (Trie a)
m