module Data.IP.Mask where

import Data.Bits
import Data.IP.Addr
import Data.Word

maskIPv4 :: Int -> IPv4
maskIPv4 len =
    IP4 $ complement $ 0xffffffff `shift` (-len)

maskIPv6 :: Int -> IPv6
maskIPv6 len =
    IP6 $
        toIP6Addr $
            bimapTup complement $
                (0xffffffffffffffff, 0xffffffffffffffff) `shift128` (-len)
  where
    bimapTup f (x, y) = (f x, f y)

shift128 :: (Word64, Word64) -> Int -> (Word64, Word64)
shift128 x i
    | i < 0 = x `shiftR128` (-i)
    | i > 0 = x `shiftL128` i
    | otherwise = x

shiftL128 :: (Word64, Word64) -> Int -> (Word64, Word64)
shiftL128 (h, l) i =
    ((h `shiftL` i) .|. (l `shift` (i - 64)), (l `shiftL` i))

shiftR128 :: (Word64, Word64) -> Int -> (Word64, Word64)
shiftR128 (h, l) i =
    (h `shiftR` i, (l `shiftR` i) .|. h `shift` (64 - i))

fromIP6Addr :: IPv6Addr -> (Word64, Word64)
fromIP6Addr (w3, w2, w1, w0) =
    ( (fromIntegral w3 `shiftL` 32) .|. fromIntegral w2
    , (fromIntegral w1 `shiftL` 32) .|. fromIntegral w0
    )

toIP6Addr :: (Word64, Word64) -> IPv6Addr
toIP6Addr (h, l) =
    ( fromIntegral $ (h `shiftR` 32) .&. m
    , fromIntegral $ h .&. m
    , fromIntegral $ (l `shiftR` 32) .&. m
    , fromIntegral $ l .&. m
    )
  where
    m = 0xffffffff
