module TrieIO where
import IOExts

data TrieIO a = TrieIO String a (IORef [TrieIO a])
     
initTrie :: String -> a -> IO (TrieIO a)
initTrie s a = do
 cref <- newIORef []
 return $ TrieIO s a cref

registerInTrie :: TrieIO a -> String -> a -> IO a
registerInTrie t p v = do
  (mbres,_) <- register t
  return $ maybe v id mbres
 where
  register t@(TrieIO s a ref) = do
   if p == s then return (Just a,True)
    else if s `prefix` p then do
     ts <- readIORef ref
     (mba,regd) <- regList ts
     if regd then return (mba,regd)
      else do
        cref <- newIORef []
        writeIORef ref (TrieIO p v cref:ts)
        return (Nothing,True)
    else return (Nothing,False)
  
  regList [] = return (Nothing,False)
  regList (t:ts) = do
    (mba,regd) <- register t
    if regd then return (mba,True)
     else regList ts


lookupTrie :: TrieIO a -> String -> IO (Maybe a)
lookupTrie t p = findOne t 
 where 
 findOne (TrieIO s a cref) =
  if s == p then return $ Just a
  else if s `prefix` p then do
    cs <- readIORef cref
    findList cs
  else
    return Nothing

 findList [] = return Nothing
 findList (t:ts) = do
   mbv <- findOne t
   maybe (findList ts) (return . Just) mbv

foldTrie :: TrieIO a -> IO [a]
foldTrie (TrieIO s a cref) = do 
  cs <- readIORef cref
  vs <- fmap concat $ mapM foldTrie cs
  return $ a:vs


deleteFromTrie :: String -> TrieIO a -> IO (Bool,Bool,[a])
deleteFromTrie p t =  delete t
 where
  delete t@(TrieIO s a cref) = do
   if s == p then do
      vs <- foldTrie t 
      return (True,True,vs)
    else if s `prefix` p then do
      cs <- readIORef cref
      (isdead, siblings,children) <- delList cs
      writeIORef cref siblings
      return (isdead,False,children)
    else
      return (False,False,[])

  delList [] = return (False,[],[])
  delList (t:ts) = do
    (isdead, medead, children) <- delete t
    if isdead then
      return (True,(if medead then [] else [t]) ++ ts,children)
     else do
      (isdead,siblings,children) <- delList ts
      return (isdead,t:siblings,children)
      
      

prefix :: String -> String -> Bool
prefix "" _ = True
prefix (c:tag) c_tag' =
   case c_tag' of
	c':tag' | c==c' -> prefix tag tag'
	_               -> False