module Data.Chatty.TST where
import Control.Arrow
import Data.Maybe
import Data.Chatty.None
data TST a = EmptyTST | TST Char (Maybe a) (TST a) (TST a) (TST a)
tstInsert :: String -> a -> TST a -> TST a
tstInsert (c:[]) a EmptyTST = TST c (Just a) EmptyTST EmptyTST EmptyTST
tstInsert (c:cs) a EmptyTST = TST c Nothing (tstInsert cs a EmptyTST) EmptyTST EmptyTST
tstInsert cs a (TST c1 a1 f l r)
| head cs < c1 = TST c1 a1 f (tstInsert cs a l) r
| head cs > c1 = TST c1 a1 f l (tstInsert cs a r)
| cs == [c1] = TST c1 (Just a) f l r
| otherwise = TST c1 a1 (tstInsert (tail cs) a f) l r
tstLookup :: String -> TST a -> Maybe a
tstLookup _ EmptyTST = Nothing
tstLookup cs (TST c1 a1 f l r)
| head cs < c1 = tstLookup cs l
| head cs > c1 = tstLookup cs r
| cs == [c1] = a1
| otherwise = tstLookup (tail cs) f
tstContains :: String -> TST a -> Bool
tstContains cs = isJust . tstLookup cs
tstTraverse :: TST a -> [(String, a)]
tstTraverse EmptyTST = []
tstTraverse (TST c1 Nothing l f r) = tstTraverse l ++ map (first (c1:)) (tstTraverse f) ++ tstTraverse r
tstTraverse (TST c1 (Just a) l f r) = tstTraverse l ++ [([c1],a)] ++ map (first (c1:)) (tstTraverse f) ++ tstTraverse r
instance None (TST a) where
none = EmptyTST