DomainMap: totally untested yet

[haskell-dns.git] / Network / DNS / DomainMap.hs

diff --git a/Network/DNS/DomainMap.hs b/Network/DNS/DomainMap.hs
new file mode 100644 (file)
index 0000000..193f2e0
--- /dev/null
+++ b/Network/DNS/DomainMap.hs
@@ -0,0 +1,160 @@
+module Network.DNS.DomainMap
+    ( DomainMap
+
+    , empty
+    , null
+    , singleton
+    , size
+
+    , fromList
+    , toList
+    , toListBy
+
+    , nearest
+    , lookup
+
+    , insert
+    , delete
+    )
+    where
+
+import qualified Data.List as L hiding (null)
+import           Data.Maybe
+import qualified Data.Trie as T
+import           Data.Trie (Trie)
+import           Network.DNS.Message
+import           Prelude hiding (null, lookup)
+
+
+newtype DomainMap a = DM (Node a)
+
+data Node a = Node (Maybe a) (Trie (Node a))
+
+
+empty :: DomainMap a
+empty = DM emptyNode
+
+
+emptyNode :: Node a
+emptyNode = Node Nothing T.empty
+
+
+null :: DomainMap a -> Bool
+null (DM n) = nodeNull n
+
+
+nodeNull :: Node a -> Bool
+nodeNull (Node (Just _) _) = False
+nodeNull (Node Nothing  t) = all nodeNull $ T.toListBy (\ _ a -> a) t
+
+
+singleton :: DomainName -> a -> DomainMap a
+singleton key value
+    = insert key value empty
+
+
+size :: forall a. DomainMap a -> Int
+size (DM n) = nodeSize n
+
+
+nodeSize :: Node a -> Int
+nodeSize (Node m t)
+    = let x  = if isJust m then
+                   1
+               else
+                   0
+          xs = map nodeSize $ T.toListBy (\ _ a -> a) t
+      in
+        L.foldl' (+) x xs
+
+
+fromList :: forall a. [(DomainName, a)] -> DomainMap a
+fromList []            = empty
+fromList ((k, v) : xs) = insert k v $ fromList xs
+
+
+toList :: forall a. DomainMap a -> [(DomainName, a)]
+toList (DM root) = toList' rootName root
+    where
+      toList' :: DomainName -> Node a -> [(DomainName, a)]
+      toList' soFar (Node m t)
+          = let x  = case m of
+                       Just v  -> [(soFar, v)]
+                       Nothing -> []
+                xs = concat $ map toList'' $ T.toList t
+            in
+              x ++ xs
+          where
+            toList'' :: (DomainLabel, Node a) -> [(DomainName, a)]
+            toList'' (l, n) = toList' (consLabel l soFar) n
+
+
+toListBy :: forall a b. DomainMap a -> (DomainName -> a -> b) -> [b]
+toListBy dm f = map (uncurry f) $ toList dm
+
+
+nearest :: forall a. DomainName -> DomainMap a -> Maybe a
+nearest key (DM root) = nearest' key' root
+    where
+      key' :: [DomainLabel]
+      key' = reverse $ nameToLabels key
+
+      nearest' :: [DomainLabel] -> Node a -> Maybe a
+      nearest' []     _          = error ("Illegal key: " ++ show key)
+      nearest' (_:[]) (Node m _) = m
+      nearest' (x:xs) (Node m t) = let nearer = do subNode' <- T.lookup x t
+                                                   nearest' xs subNode'
+                                   in
+                                     if isJust nearer then
+                                         nearer
+                                     else
+                                         m
+
+
+lookup :: forall a. DomainName -> DomainMap a -> Maybe a
+lookup key (DM root) = lookup' key' root
+    where
+      key' :: [DomainLabel]
+      key' = reverse $ nameToLabels key
+
+      lookup' :: [DomainLabel] -> Node a -> Maybe a
+      lookup' []     _          = error ("Illegal key: " ++ show key)
+      lookup' (_:[]) (Node m _) = m
+      lookup' (x:xs) (Node _ t) = do subNode' <- T.lookup x t
+                                     lookup' xs subNode'
+
+
+insert :: forall a. DomainName -> a -> DomainMap a -> DomainMap a
+insert key value (DM root) = DM $ insert' key' root
+    where
+      key' :: [DomainLabel]
+      key' = reverse $ nameToLabels key
+
+      insert' :: [DomainLabel] -> Node a -> Node a
+      insert' []     _          = error ("Illegal key: " ++ show key)
+      insert' (_:[]) (Node _ t) = Node (Just value) t
+      insert' (x:xs) (Node m t) = let subNode' = case T.lookup x t of
+                                                   Just subNode -> insert' xs subNode
+                                                   Nothing      -> insert' xs emptyNode
+                                      trie'    = T.insert x subNode' t
+                                  in
+                                    Node m trie'
+
+delete :: forall a. DomainName -> DomainMap a -> DomainMap a
+delete key (DM root) = DM $ delete' key' root
+    where
+      key' :: [DomainLabel]
+      key' = reverse $ nameToLabels key
+
+      delete' :: [DomainLabel] -> Node a -> Node a
+      delete' []     _          = error ("Illegal key: " ++ show key)
+      delete' (_:[]) (Node _ t) = Node Nothing t
+      delete' (x:xs) (Node m t) = let subNode' = case T.lookup x t of
+                                                   Just subNode -> delete' xs subNode
+                                                   Nothing      -> emptyNode
+                                      trie'    = if nodeNull subNode' then
+                                                     T.delete x t
+                                                 else
+                                                     T.insert x subNode' t
+                                  in
+                                    Node m trie'