iproute-1.7.15: Data/IP/RouteTable/Internal.hs
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveGeneric #-}
-- |
-- IP routing table is a tree of 'AddrRange'
-- to search one of them on the longest
-- match base. It is a kind of TRIE with one
-- way branching removed. Both IPv4 and IPv6
-- are supported.
module Data.IP.RouteTable.Internal where
import Control.Applicative hiding (empty)
import qualified Control.Applicative as A (empty)
import Control.Monad
import Data.Bits
import Data.Foldable (Foldable (..))
import Data.IP.Addr
import Data.IP.Op
import Data.IP.Range
import Data.IntMap (IntMap, (!))
import qualified Data.IntMap as IM (fromList)
import Data.Monoid hiding ((<>))
import Data.Semigroup
import Data.Traversable
import Data.Word
import GHC.Generics (Generic, Generic1)
import Prelude hiding (lookup)
-- $setup
-- >>> :set -XOverloadedStrings
----------------------------------------------------------------
-- |
-- A class to contain IPv4 and IPv6.
class Addr a => Routable a where
-- |
-- The 'intToTBit' function takes 'Int' and returns an 'Routable' address
-- whose only n-th bit is set.
intToTBit :: Int -> a
-- |
-- The 'isZero' function takes an 'Routable' address and an test bit
-- 'Routable' address and returns 'True' is the bit is unset,
-- otherwise returns 'False'.
isZero :: a -> a -> Bool
instance Routable IPv4 where
intToTBit = intToTBitIPv4
isZero a b = a `masked` b == IP4 0
instance Routable IPv6 where
intToTBit = intToTBitIPv6
isZero a b = a `masked` b == IP6 (0, 0, 0, 0)
----------------------------------------------------------------
--
-- Test Bit
--
intToTBitIPv4 :: Int -> IPv4
intToTBitIPv4 len = IP4 (intToTBitsIPv4 ! len)
intToTBitIPv6 :: Int -> IPv6
intToTBitIPv6 len = IP6 (intToTBitsIPv6 ! len)
intToTBitsWord32 :: [Word32]
intToTBitsWord32 = iterate (`shift` (-1)) 0x80000000
intToTBitsIPv4 :: IntMap IPv4Addr
intToTBitsIPv4 = IM.fromList $ zip [0 .. 32] intToTBitsWord32
intToTBitsIPv6 :: IntMap IPv6Addr
intToTBitsIPv6 = IM.fromList $ zip [0 .. 128] bs
where
bs = b1 ++ b2 ++ b3 ++ b4 ++ b5
b1 = map (\vbit -> (vbit, all0, all0, all0)) intToTBits
b2 = map (\vbit -> (all0, vbit, all0, all0)) intToTBits
b3 = map (\vbit -> (all0, all0, vbit, all0)) intToTBits
b4 = map (\vbit -> (all0, all0, all0, vbit)) intToTBits
b5 = [(all0, all0, all0, all0)]
intToTBits = take 32 intToTBitsWord32
all0 = 0x00000000
----------------------------------------------------------------
-- |
-- The Tree structure for IP routing table based on TRIE with
-- one way branching removed. This is an abstract data type,
-- so you cannot touch its inside. Please use 'insert' or 'lookup', instead.
data IPRTable k a
= Nil
| Node !(AddrRange k) !k !(Maybe a) !(IPRTable k a) !(IPRTable k a)
deriving (Eq, Generic, Generic1, Show)
----------------------------------------------------------------
-- |
-- The 'empty' function returns an empty IP routing table.
--
-- >>> (empty :: IPRTable IPv4 ()) == fromList []
-- True
empty :: Routable k => IPRTable k a
empty = Nil
instance Functor (IPRTable k) where
fmap _ Nil = Nil
fmap f (Node r a mv b1 b2) = Node r a (f <$> mv) (fmap f b1) (fmap f b2)
instance Foldable (IPRTable k) where
foldMap _ Nil = mempty
foldMap f (Node _ _ mv b1 b2) = foldMap f mv `mappend` foldMap f b1 `mappend` foldMap f b2
instance Traversable (IPRTable k) where
traverse _ Nil = pure Nil
traverse f (Node r a mv b1 b2) = Node r a <$> traverse f mv <*> traverse f b1 <*> traverse f b2
-- | Note that Semigroup and Monoid instances are right-biased.
-- That is, if both arguments have the same key, the value from the right
-- argument will be used.
-- Since: 1.7.5
instance Routable k => Semigroup (IPRTable k a) where
a <> b = foldlWithKey (\rt k v -> insert k v rt) a b
stimes = stimesIdempotent
-- | Since: 1.7.5
instance Routable k => Monoid (IPRTable k a) where
mempty = empty
mappend = (<>)
----------------------------------------------------------------
-- |
-- The 'insert' function inserts a value with a key of 'AddrRange' to 'IPRTable'
-- and returns a new 'IPRTable'.
--
-- >>> (insert ("127.0.0.1" :: AddrRange IPv4) () empty) == fromList [("127.0.0.1",())]
-- True
insert :: Routable k => AddrRange k -> a -> IPRTable k a -> IPRTable k a
insert k1 v1 Nil = Node k1 tb1 (Just v1) Nil Nil
where
tb1 = keyToTestBit k1
insert k1 v1 s@(Node k2 tb2 v2 l r)
| k1 == k2 = Node k1 tb1 (Just v1) l r
| k2 >:> k1 =
if isLeft k1 tb2
then
Node k2 tb2 v2 (insert k1 v1 l) r
else
Node k2 tb2 v2 l (insert k1 v1 r)
| k1 >:> k2 =
if isLeft k2 tb1
then
Node k1 tb1 (Just v1) s Nil
else
Node k1 tb1 (Just v1) Nil s
| otherwise =
let n = Node k1 tb1 (Just v1) Nil Nil
in link n s
where
tb1 = keyToTestBit k1
link :: Routable k => IPRTable k a -> IPRTable k a -> IPRTable k a
link s1@(Node k1 _ _ _ _) s2@(Node k2 _ _ _ _)
| isLeft k1 tbg = Node kg tbg Nothing s1 s2
| otherwise = Node kg tbg Nothing s2 s1
where
kg = glue 0 k1 k2
tbg = keyToTestBit kg
link _ _ = error "link"
glue :: Routable k => Int -> AddrRange k -> AddrRange k -> AddrRange k
glue n k1 k2
| addr k1 `masked` mk == addr k2 `masked` mk = glue (n + 1) k1 k2
| otherwise = makeAddrRange (addr k1) (n - 1)
where
mk = intToMask n
keyToTestBit :: Routable k => AddrRange k -> k
keyToTestBit = intToTBit . mlen
isLeft :: Routable k => AddrRange k -> k -> Bool
isLeft adr = isZero (addr adr)
----------------------------------------------------------------
-- |
-- The 'delete' function deletes a value by a key of 'AddrRange' from 'IPRTable'
-- and returns a new 'IPRTable'.
--
-- >>> delete "127.0.0.1" (insert "127.0.0.1" () empty) == (empty :: IPRTable IPv4 ())
-- True
delete :: Routable k => AddrRange k -> IPRTable k a -> IPRTable k a
delete _ Nil = Nil
delete k1 s@(Node k2 tb2 v2 l r)
| k1 == k2 = node k2 tb2 Nothing l r
| k2 >:> k1 =
if isLeft k1 tb2
then
node k2 tb2 v2 (delete k1 l) r
else
node k2 tb2 v2 l (delete k1 r)
| otherwise = s
node
:: Routable k
=> AddrRange k -> k -> Maybe a -> IPRTable k a -> IPRTable k a -> IPRTable k a
node _ _ Nothing Nil r = r
node _ _ Nothing l Nil = l
node k tb v l r = Node k tb v l r
----------------------------------------------------------------
-- |
-- The 'lookup' function looks up 'IPRTable' with a key of 'AddrRange'.
-- If a routing information in 'IPRTable' matches the key, its value
-- is returned.
--
-- >>> let v4 = ["133.4.0.0/16","133.5.0.0/16","133.5.16.0/24","133.5.23.0/24"] :: [AddrRange IPv4]
-- >>> let rt = fromList $ zip v4 v4
-- >>> lookup "127.0.0.1" rt
-- Nothing
-- >>> lookup "133.3.0.1" rt
-- Nothing
-- >>> lookup "133.4.0.0" rt
-- Just 133.4.0.0/16
-- >>> lookup "133.4.0.1" rt
-- Just 133.4.0.0/16
-- >>> lookup "133.5.16.0" rt
-- Just 133.5.16.0/24
-- >>> lookup "133.5.16.1" rt
-- Just 133.5.16.0/24
lookup :: Routable k => AddrRange k -> IPRTable k a -> Maybe a
lookup k s = fmap snd (search k s Nothing)
-- |
-- The 'lookupKeyValue' function looks up 'IPRTable' with a key of 'AddrRange'.
-- If a routing information in 'IPRTable' matches the key, both key and value
-- are returned.
--
-- >>> :set -XOverloadedStrings
-- >>> let rt = fromList ([("192.168.0.0/24", 1), ("10.10.0.0/16", 2)] :: [(AddrRange IPv4, Int)])
-- >>> lookupKeyValue "127.0.0.1" rt
-- Nothing
-- >>> lookupKeyValue "192.168.0.1" rt
-- Just (192.168.0.0/24,1)
-- >>> lookupKeyValue "10.10.0.1" rt
-- Just (10.10.0.0/16,2)
lookupKeyValue
:: Routable k => AddrRange k -> IPRTable k a -> Maybe (AddrRange k, a)
lookupKeyValue k s = search k s Nothing
search
:: Routable k
=> AddrRange k
-> IPRTable k a
-> Maybe (AddrRange k, a)
-> Maybe (AddrRange k, a)
search _ Nil res = res
search k1 (Node k2 tb2 Nothing l r) res
| k1 == k2 = res
| k2 >:> k1 =
if isLeft k1 tb2
then
search k1 l res
else
search k1 r res
| otherwise = res
search k1 (Node k2 tb2 (Just vl) l r) res
| k1 == k2 = Just (k1, vl)
| k2 >:> k1 =
if isLeft k1 tb2
then
search k1 l $ Just (k2, vl)
else
search k1 r $ Just (k2, vl)
| otherwise = res
-- |
-- 'lookupAll' is a version of 'lookup' that returns all entries matching the
-- given key, not just the longest match.
--
-- >>> :set -XOverloadedStrings
-- >>> let rt = fromList ([("192.168.0.0/24", 1), ("10.10.0.0/16", 2), ("10.0.0.0/8", 3)] :: [(AddrRange IPv4, Int)])
-- >>> lookupAll "127.0.0.1" rt
-- []
-- >>> lookupAll "192.168.0.1" rt
-- [(192.168.0.0/24,1)]
-- >>> lookupAll "10.10.0.1" rt
-- [(10.10.0.0/16,2),(10.0.0.0/8,3)]
lookupAll :: Routable k => AddrRange k -> IPRTable k a -> [(AddrRange k, a)]
lookupAll range = go []
where
go acc Nil = acc
go acc (Node k tb Nothing l r)
| k == range = acc
| k >:> range = go acc $ if isLeft range tb then l else r
| otherwise = acc
go acc (Node k tb (Just v) l r)
| k == range = (k, v) : acc
| k >:> range = go ((k, v) : acc) $ if isLeft range tb then l else r
| otherwise = acc
----------------------------------------------------------------
-- |
-- The 'findMatch' function looks up 'IPRTable' with a key of 'AddrRange'.
-- If the key matches routing informations in 'IPRTable', they are
-- returned.
--
-- >>> let v4 = ["133.4.0.0/16","133.5.0.0/16","133.5.16.0/24","133.5.23.0/24"] :: [AddrRange IPv4]
-- >>> let rt = fromList $ zip v4 $ repeat ()
-- >>> findMatch "133.4.0.0/15" rt :: [(AddrRange IPv4,())]
-- [(133.4.0.0/16,()),(133.5.0.0/16,()),(133.5.16.0/24,()),(133.5.23.0/24,())]
findMatch
:: Alternative m => Routable k => AddrRange k -> IPRTable k a -> m (AddrRange k, a)
findMatch _ Nil = A.empty
findMatch k1 (Node k2 _ Nothing l r)
| k1 >:> k2 = findMatch k1 l <|> findMatch k1 r
| k2 >:> k1 = findMatch k1 l <|> findMatch k1 r
| otherwise = A.empty
findMatch k1 (Node k2 _ (Just vl) l r)
| k1 >:> k2 = pure (k2, vl) <|> findMatch k1 l <|> findMatch k1 r
| k2 >:> k1 = findMatch k1 l <|> findMatch k1 r
| otherwise = A.empty
----------------------------------------------------------------
-- |
-- The 'fromList' function creates a new IP routing table from
-- a list of a pair of 'IPrange' and value.
fromList :: Routable k => [(AddrRange k, a)] -> IPRTable k a
fromList = foldl' (\s (k, v) -> insert k v s) empty
-- |
-- The 'toList' function creates a list of a pair of 'AddrRange' and
-- value from an IP routing table.
toList :: Routable k => IPRTable k a -> [(AddrRange k, a)]
toList = foldt toL []
where
toL Nil xs = xs
toL (Node _ _ Nothing _ _) xs = xs
toL (Node k _ (Just a) _ _) xs = (k, a) : xs
----------------------------------------------------------------
foldt :: (IPRTable k a -> b -> b) -> b -> IPRTable k a -> b
foldt _ v Nil = v
foldt func v rt@(Node _ _ _ l r) = foldt func (foldt func (func rt v) l) r
-- | /O(n)/. Fold the keys and values in the IPRTable using the given
-- left-associative binary operator.
-- This function is equivalent to Data.Map.foldlWithKey with necessary to
-- IPRTable changes.
-- Since: 1.7.5
foldlWithKey :: (b -> AddrRange k -> a -> b) -> b -> IPRTable k a -> b
foldlWithKey f zr = go zr
where
go z Nil = z
go z (Node _ _ Nothing l r) = go (go z l) r
go z (Node n _ (Just v) l r) = go (f (go z l) n v) r
{-# INLINE foldlWithKey #-}
-- | /O(n)/. Fold the keys and values in the IPRTable using the given
-- right-associative binary operator.
-- This function is equivalent to Data.Map.foldrWithKey with necessary to
-- IPRTable changes.
-- Since: 1.7.5
foldrWithKey :: (AddrRange k -> a -> b -> b) -> b -> IPRTable k a -> b
foldrWithKey f zr = go zr
where
go z Nil = z
go z (Node _ _ Nothing l r) = go (go z r) l
go z (Node n _ (Just v) l r) = go (f n v (go z r)) l
{-# INLINE foldrWithKey #-}