Subnet.hs

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds       #-}
{-# LANGUAGE PolyKinds       #-}
{-# LANGUAGE TypeFamilies    #-}
{-# LANGUAGE TypeOperators   #-}
{-# LANGUAGE ViewPatterns    #-}

module Subnet
    ( KnownBits
    , bitmask

    , Netmask (..)
    , netmask

    , Subnet  (..)
    , hostbits
    , netbits
    , cidr
    , hosts

    , divide
    , join

    , prefix
    ) where

import Data.Word (Word32)

import GHC.TypeLits (CmpNat, KnownNat, natVal, type (+), type (-), Nat)

import Numeric (showInt)

import Terrafomo.Syntax.IP

import qualified Data.Bits     as Bit
import qualified Data.Foldable as Fold

-- FIXME: smart constructors for safe ip range / prefixes.

type KnownBits n = (KnownNat n, CmpNat n 33 ~ 'LT)

bitmask :: KnownBits n => proxy n -> Bits
bitmask = toEnum . fromEnum . natVal

newtype Netmask = Netmask Word32
    deriving (Eq)

instance Show Netmask where
    showsPrec _ (Netmask w) =
        let showBits n = showInt (fromEnum (Bit.shiftR w n Bit..&. 0xff))
         in showBits 0o30
          . showChar '.'
          . showBits 0o20
          . showChar '.'
          . showBits 0o10
          . showChar '.'
          . showBits 0o00

netmask :: KnownBits n => Subnet n -> Netmask
netmask s = Netmask (Bit.shiftL 1 32 - Bit.shiftL 1 (hostbits s))

newtype Subnet (n :: Nat) = Subnet IP

instance KnownBits n => Show (Subnet n) where
    showsPrec _ (cidr -> ip :/ mask) =
          shows ip
        . showString "/"
        . shows (fromEnum mask)

hostbits :: KnownBits n => Subnet n -> Int
hostbits s = 32 - netbits s

netbits :: KnownBits n => Subnet n -> Int
netbits = fromIntegral . natVal

cidr :: KnownBits n => Subnet n -> CIDR
cidr s@(Subnet ip) = ip :/ bitmask s

hosts :: KnownBits n => Subnet n -> Int
hosts s = 2 ^ (32 - fromEnum (bitmask s)) - 2

-- @
-- divide subnet == uncurry join (divide subnet)
-- @
divide
    :: (KnownBits n, CmpNat n 32 ~ 'LT)
    => Subnet n
    -> (Subnet (n + 1), Subnet (n + 1))
divide s@(Subnet (IPv4 w)) =
    let bits = 32 - (netbits s + 1)
     in ( Subnet (IPv4 w)
        , Subnet (IPv4 (Bit.setBit w bits))
        )

-- This hasn't had a lot of applied thought.
join :: KnownBits n => Subnet n -> Subnet n -> Subnet (n - 1)
join s@(Subnet ip) _ = Subnet (prefix ip (bitmask s))

-- If @prefix ip bits /= ip@, then it's not a subnet boundary.
prefix :: IP -> Bits -> IP
prefix (IPv4 w) mask = IPv4 (Fold.foldl' go w bits)
  where
    go x i = x Bit..&. Bit.complement (Bit.shiftL 1 i)
    bits   = [0 .. 31 - fromEnum mask]