module Graphics.LambdaCube.Tree where
import Data.ByteString.Char8 (ByteString)
import Data.Trie (Trie)
import qualified Data.ByteString.Char8 as B
import qualified Data.Trie as T
import qualified Data.Trie.Convenience as T
data IndexedTree a = IT String (Trie a) (Trie ByteString) deriving Show
separator :: ByteString
separator = B.pack "\0"
sep :: Char
sep = B.head separator
root :: String -> a -> IndexedTree a
root name dat = IT "\1" (T.singleton nn dat) (T.singleton n nn)
where
n = B.pack name
nn = B.pack "\1"
addNode :: String -> String -> a -> IndexedTree a -> Maybe (IndexedTree a)
addNode parent name dat it@(IT nn dm nm) = case T.member n nm of
True -> Just (updateNode (const dat) name it)
False -> fmap add (T.lookup p nm)
where
p = B.pack parent
n = B.pack name
nn' = nextName nn
add path = IT nn' (T.insert path' dat dm) (T.insert n path' nm)
where
path' = B.intercalate separator [path, B.pack nn']
deleteNode :: String -> IndexedTree a -> IndexedTree a
deleteNode name it@(IT nn dm nm) = case T.lookup n nm of
Nothing -> it
Just path -> IT nn dm' nm'
where
subkeys = T.keys (T.submap path dm)
subnames = map (snd . B.breakEnd (==sep)) subkeys
dm' = T.disunion dm (T.fromListS (map (\k -> (k,undefined)) subkeys))
nm' = T.disunion nm (T.fromListS (map (\k -> (k,B.pack "")) subnames))
where
n = B.pack name
updateNode :: (a -> a) -> String -> IndexedTree a -> IndexedTree a
updateNode f name it@(IT nn dm nm) = case T.lookup n nm of
Nothing -> it
Just path -> IT nn (T.adjust f path dm) nm
where
n = B.pack name
getDataAt :: String -> IndexedTree a -> Maybe a
getDataAt name (IT _ dm nm) = flip T.lookup dm =<< T.lookup n nm
where
n = B.pack name
flattenTree :: (a -> b -> a) -> (a -> b -> r) -> a -> IndexedTree b -> [r]
flattenTree f g acc (IT _ dm _) = go (T.toList dm) [] [acc]
where
go [] _ _ = []
go ((path,dat):pds') [] as@(acc:_) = g acc dat : go pds' [path] (f acc dat:as)
go pds@((path,dat):pds') ps@(p:ps') as@(acc:as')
| prefix = g acc dat : go pds' (path:ps) (f acc dat:as)
| otherwise = go pds ps' as'
where
prefix = (p `B.isPrefixOf` path) && (B.index path (B.length p) == sep)
nextName :: String -> String
nextName "" = "\1"
nextName ('\255':s) = '\1' : nextName s
nextName (c:s) = succ c : s