-- | 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
[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)

-- | Creates an empty trie.
empty :: Trie a
empty :: Trie a
empty = Maybe a -> Map Word8 (Trie a) -> Trie a
forall a. Maybe a -> Map Word8 (Trie a) -> Trie a
Trie Maybe a
forall a. Maybe a
Nothing Map Word8 (Trie a)
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 :: [(BS, a)] -> Trie a
fromList = (Trie a -> (BS, a) -> Trie a) -> Trie a -> [(BS, a)] -> Trie a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Trie a
t (BS
x,a
y) -> BS -> a -> Trie a -> Trie a
forall a. BS -> a -> Trie a -> Trie a
insert BS
x a
y Trie a
t) Trie a
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 :: Trie a -> BS -> Maybe (BS, a, BS)
match Trie a
t BS
s =
  case Maybe (Int, a) -> Int -> Trie a -> BS -> Maybe (Int, a)
forall a b.
Num a =>
Maybe (a, b) -> a -> Trie b -> BS -> Maybe (a, b)
go Maybe (Int, a)
forall a. Maybe a
Nothing Int
0 Trie a
t BS
s of
    Maybe (Int, a)
Nothing -> Maybe (BS, a, BS)
forall a. Maybe a
Nothing
    Just (Int
n, a
v) -> (BS, a, BS) -> Maybe (BS, a, BS)
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 (a, b) -> a -> Trie b -> BS -> Maybe (a, b)
go Maybe (a, b)
a a
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 (a, b)
a
        Just (Word8
w, BS
t) ->
          case Word8 -> Map Word8 (Trie b) -> Maybe (Trie b)
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 (a, b)
a
            Just b :: Trie b
b@(Trie (Just b
v2) Map Word8 (Trie b)
_) ->
              Maybe (a, b) -> a -> Trie b -> BS -> Maybe (a, b)
go ((a, b) -> Maybe (a, b)
forall a. a -> Maybe a
Just (a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
1, b
v2)) (a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) Trie b
b BS
t
            Just Trie b
b ->
              Maybe (a, b) -> a -> Trie b -> BS -> Maybe (a, b)
go Maybe (a, b)
a (a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) Trie b
b BS
t

-- | Inserts a value into a trie.
insert :: BS -> a -> Trie a -> Trie a
insert :: BS -> a -> Trie a -> Trie a
insert BS
x = [Word8] -> a -> Trie a -> Trie a
forall a. [Word8] -> a -> Trie a -> Trie a
insertWords ([Word8] -> a -> Trie a -> Trie a)
-> [Word8] -> a -> Trie a -> Trie a
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 :: [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) =
      Maybe a -> Map Word8 (Trie a) -> Trie a
forall a. Maybe a -> Map Word8 (Trie a) -> Trie a
Trie (a -> Maybe a
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 Word8 -> Map Word8 (Trie a) -> Maybe (Trie a)
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 ->
          Maybe a -> Map Word8 (Trie a) -> Trie a
forall a. Maybe a -> Map Word8 (Trie a) -> Trie a
Trie Maybe a
v (Map Word8 (Trie a) -> Trie a) -> Map Word8 (Trie a) -> Trie a
forall a b. (a -> b) -> a -> b
$ Word8 -> Trie a -> Map Word8 (Trie a) -> Map Word8 (Trie a)
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
forall a. Trie a
empty) Map Word8 (Trie a)
m
        Just Trie a
b ->
          Maybe a -> Map Word8 (Trie a) -> Trie a
forall a. Maybe a -> Map Word8 (Trie a) -> Trie a
Trie Maybe a
v (Map Word8 (Trie a) -> Trie a) -> Map Word8 (Trie a) -> Trie a
forall a b. (a -> b) -> a -> b
$ Word8 -> Trie a -> Map Word8 (Trie a) -> Map Word8 (Trie a)
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