{-# LANGUAGE Safe #-}

{-
  This module is part of Chatty.
  Copyleft (c) 2014 Marvin Cohrs

  All wrongs reversed. Sharing is an act of love, not crime.
  Please share Chatty with everyone you like.

  Chatty is free software: you can redistribute it and/or modify
  it under the terms of the GNU Affero General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  Chatty is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  GNU Affero General Public License for more details.

  You should have received a copy of the GNU Affero General Public License
  along with Chatty. If not, see <http://www.gnu.org/licenses/>.
-}

-- | Provides a ternary search trie
module Data.Chatty.TST where

import Control.Arrow
import Data.Maybe
import Data.Chatty.None

-- | A ternary search trie
data TST a = EmptyTST | TST Char (Maybe a) (TST a) (TST a) (TST a)

-- | Insert something into the TST
tstInsert :: String -> a -> TST a -> TST a
tstInsert :: String -> a -> TST a -> TST a
tstInsert (Char
c:[]) a
a TST a
EmptyTST = Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
forall a. Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
TST Char
c (a -> Maybe a
forall a. a -> Maybe a
Just a
a) TST a
forall a. TST a
EmptyTST TST a
forall a. TST a
EmptyTST TST a
forall a. TST a
EmptyTST
tstInsert (Char
c:String
cs) a
a TST a
EmptyTST = Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
forall a. Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
TST Char
c Maybe a
forall a. Maybe a
Nothing (String -> a -> TST a -> TST a
forall a. String -> a -> TST a -> TST a
tstInsert String
cs a
a TST a
forall a. TST a
EmptyTST) TST a
forall a. TST a
EmptyTST TST a
forall a. TST a
EmptyTST
tstInsert String
cs a
a (TST Char
c1 Maybe a
a1 TST a
f TST a
l TST a
r)
  | String -> Char
forall a. [a] -> a
head String
cs Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
< Char
c1 = Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
forall a. Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
TST Char
c1 Maybe a
a1 TST a
f (String -> a -> TST a -> TST a
forall a. String -> a -> TST a -> TST a
tstInsert String
cs a
a TST a
l) TST a
r
  | String -> Char
forall a. [a] -> a
head String
cs Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
> Char
c1 = Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
forall a. Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
TST Char
c1 Maybe a
a1 TST a
f TST a
l (String -> a -> TST a -> TST a
forall a. String -> a -> TST a -> TST a
tstInsert String
cs a
a TST a
r)
  | String
cs String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== [Char
c1] = Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
forall a. Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
TST Char
c1 (a -> Maybe a
forall a. a -> Maybe a
Just a
a) TST a
f TST a
l TST a
r
  | Bool
otherwise = Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
forall a. Char -> Maybe a -> TST a -> TST a -> TST a -> TST a
TST Char
c1 Maybe a
a1 (String -> a -> TST a -> TST a
forall a. String -> a -> TST a -> TST a
tstInsert (String -> String
forall a. [a] -> [a]
tail String
cs) a
a TST a
f) TST a
l TST a
r

-- | Lookup some string in the TST
tstLookup :: String -> TST a -> Maybe a
tstLookup :: String -> TST a -> Maybe a
tstLookup String
_ TST a
EmptyTST = Maybe a
forall a. Maybe a
Nothing
tstLookup String
cs (TST Char
c1 Maybe a
a1 TST a
f TST a
l TST a
r)
  | String -> Char
forall a. [a] -> a
head String
cs Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
< Char
c1 = String -> TST a -> Maybe a
forall a. String -> TST a -> Maybe a
tstLookup String
cs TST a
l
  | String -> Char
forall a. [a] -> a
head String
cs Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
> Char
c1 = String -> TST a -> Maybe a
forall a. String -> TST a -> Maybe a
tstLookup String
cs TST a
r
  | String
cs String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== [Char
c1] = Maybe a
a1
  | Bool
otherwise = String -> TST a -> Maybe a
forall a. String -> TST a -> Maybe a
tstLookup (String -> String
forall a. [a] -> [a]
tail String
cs) TST a
f

-- | Check if the TST contains the given key
tstContains :: String -> TST a -> Bool
tstContains :: String -> TST a -> Bool
tstContains String
cs = Maybe a -> Bool
forall a. Maybe a -> Bool
isJust (Maybe a -> Bool) -> (TST a -> Maybe a) -> TST a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> TST a -> Maybe a
forall a. String -> TST a -> Maybe a
tstLookup String
cs

-- | Traverse TST
tstTraverse :: TST a -> [(String, a)]
tstTraverse :: TST a -> [(String, a)]
tstTraverse TST a
EmptyTST = []
tstTraverse (TST Char
c1 Maybe a
Nothing TST a
l TST a
f TST a
r) = TST a -> [(String, a)]
forall a. TST a -> [(String, a)]
tstTraverse TST a
l [(String, a)] -> [(String, a)] -> [(String, a)]
forall a. [a] -> [a] -> [a]
++ ((String, a) -> (String, a)) -> [(String, a)] -> [(String, a)]
forall a b. (a -> b) -> [a] -> [b]
map ((String -> String) -> (String, a) -> (String, a)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (Char
c1Char -> String -> String
forall a. a -> [a] -> [a]
:)) (TST a -> [(String, a)]
forall a. TST a -> [(String, a)]
tstTraverse TST a
f) [(String, a)] -> [(String, a)] -> [(String, a)]
forall a. [a] -> [a] -> [a]
++ TST a -> [(String, a)]
forall a. TST a -> [(String, a)]
tstTraverse TST a
r
tstTraverse (TST Char
c1 (Just a
a) TST a
l TST a
f TST a
r) = TST a -> [(String, a)]
forall a. TST a -> [(String, a)]
tstTraverse TST a
l [(String, a)] -> [(String, a)] -> [(String, a)]
forall a. [a] -> [a] -> [a]
++ [([Char
c1],a
a)] [(String, a)] -> [(String, a)] -> [(String, a)]
forall a. [a] -> [a] -> [a]
++ ((String, a) -> (String, a)) -> [(String, a)] -> [(String, a)]
forall a b. (a -> b) -> [a] -> [b]
map ((String -> String) -> (String, a) -> (String, a)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (Char
c1Char -> String -> String
forall a. a -> [a] -> [a]
:)) (TST a -> [(String, a)]
forall a. TST a -> [(String, a)]
tstTraverse TST a
f) [(String, a)] -> [(String, a)] -> [(String, a)]
forall a. [a] -> [a] -> [a]
++ TST a -> [(String, a)]
forall a. TST a -> [(String, a)]
tstTraverse TST a
r

instance None (TST a) where
  none :: TST a
none = TST a
forall a. TST a
EmptyTST