Network/DNS/DomainMap.hs

   1 module Network.DNS.DomainMap
   2     ( DomainMap
   3
   4     , empty
   5     , null
   6     , singleton
   7     , size
   8
   9     , fromList
  10     , toList
  11     , toListBy
  12
  13     , nearest
  14     , lookup
  15
  16     , insert
  17     , delete
  18     )
  19     where
  20
  21 import qualified Data.List as L hiding (null)
  22 import           Data.Maybe
  23 import qualified Data.Trie as T
  24 import           Data.Trie (Trie)
  25 import           Network.DNS.Message
  26 import           Prelude hiding (null, lookup)
  27
  28
  29 newtype DomainMap a = DM (Node a)
  30
  31 data Node a = Node (Maybe a) (Trie (Node a))
  32
  33
  34 empty :: DomainMap a
  35 empty = DM emptyNode
  36
  37
  38 emptyNode :: Node a
  39 emptyNode = Node Nothing T.empty
  40
  41
  42 null :: DomainMap a -> Bool
  43 null (DM n) = nodeNull n
  44
  45
  46 nodeNull :: Node a -> Bool
  47 nodeNull (Node (Just _) _) = False
  48 nodeNull (Node Nothing  t) = all nodeNull $ T.toListBy (\ _ a -> a) t
  49
  50
  51 singleton :: DomainName -> a -> DomainMap a
  52 singleton key value
  53     = insert key value empty
  54
  55
  56 size :: forall a. DomainMap a -> Int
  57 size (DM n) = nodeSize n
  58
  59
  60 nodeSize :: Node a -> Int
  61 nodeSize (Node m t)
  62     = let x  = if isJust m then
  63                    1
  64                else
  65                    0
  66           xs = map nodeSize $ T.toListBy (\ _ a -> a) t
  67       in
  68         L.foldl' (+) x xs
  69
  70
  71 fromList :: forall a. [(DomainName, a)] -> DomainMap a
  72 fromList []            = empty
  73 fromList ((k, v) : xs) = insert k v $ fromList xs
  74
  75
  76 toList :: forall a. DomainMap a -> [(DomainName, a)]
  77 toList (DM root) = toList' rootName root
  78     where
  79       toList' :: DomainName -> Node a -> [(DomainName, a)]
  80       toList' soFar (Node m t)
  81           = let x  = case m of
  82                        Just v  -> [(soFar, v)]
  83                        Nothing -> []
  84                 xs = concat $ map toList'' $ T.toList t
  85             in
  86               x ++ xs
  87           where
  88             toList'' :: (DomainLabel, Node a) -> [(DomainName, a)]
  89             toList'' (l, n) = toList' (consLabel l soFar) n
  90
  91
  92 toListBy :: forall a b. DomainMap a -> (DomainName -> a -> b) -> [b]
  93 toListBy dm f = map (uncurry f) $ toList dm
  94
  95
  96 nearest :: forall a. DomainName -> DomainMap a -> Maybe a
  97 nearest key (DM root) = nearest' key' root
  98     where
  99       key' :: [DomainLabel]
 100       key' = reverse $ nameToLabels key
 101
 102       nearest' :: [DomainLabel] -> Node a -> Maybe a
 103       nearest' []     _          = error ("Illegal key: " ++ show key)
 104       nearest' (_:[]) (Node m _) = m
 105       nearest' (x:xs) (Node m t) = let nearer = do subNode' <- T.lookup x t
 106                                                    nearest' xs subNode'
 107                                    in
 108                                      if isJust nearer then
 109                                          nearer
 110                                      else
 111                                          m
 112
 113
 114 lookup :: forall a. DomainName -> DomainMap a -> Maybe a
 115 lookup key (DM root) = lookup' key' root
 116     where
 117       key' :: [DomainLabel]
 118       key' = reverse $ nameToLabels key
 119
 120       lookup' :: [DomainLabel] -> Node a -> Maybe a
 121       lookup' []     _          = error ("Illegal key: " ++ show key)
 122       lookup' (_:[]) (Node m _) = m
 123       lookup' (x:xs) (Node _ t) = do subNode' <- T.lookup x t
 124                                      lookup' xs subNode'
 125
 126
 127 insert :: forall a. DomainName -> a -> DomainMap a -> DomainMap a
 128 insert key value (DM root) = DM $ insert' key' root
 129     where
 130       key' :: [DomainLabel]
 131       key' = reverse $ nameToLabels key
 132
 133       insert' :: [DomainLabel] -> Node a -> Node a
 134       insert' []     _          = error ("Illegal key: " ++ show key)
 135       insert' (_:[]) (Node _ t) = Node (Just value) t
 136       insert' (x:xs) (Node m t) = let subNode' = case T.lookup x t of
 137                                                    Just subNode -> insert' xs subNode
 138                                                    Nothing      -> insert' xs emptyNode
 139                                       trie'    = T.insert x subNode' t
 140                                   in
 141                                     Node m trie'
 142
 143 delete :: forall a. DomainName -> DomainMap a -> DomainMap a
 144 delete key (DM root) = DM $ delete' key' root
 145     where
 146       key' :: [DomainLabel]
 147       key' = reverse $ nameToLabels key
 148
 149       delete' :: [DomainLabel] -> Node a -> Node a
 150       delete' []     _          = error ("Illegal key: " ++ show key)
 151       delete' (_:[]) (Node _ t) = Node Nothing t
 152       delete' (x:xs) (Node m t) = let subNode' = case T.lookup x t of
 153                                                    Just subNode -> delete' xs subNode
 154                                                    Nothing      -> emptyNode
 155                                       trie'    = if nodeNull subNode' then
 156                                                      T.delete x t
 157                                                  else
 158                                                      T.insert x subNode' t
 159                                   in
 160                                     Node m trie'