Using RGB types to gain intuition about Type Level Programming and Generics

Learn
13

@tobz619 - Expanding upon my earlier example with some type-level stuff from the deep!

First, we need to turn on some more language features, because it’s going to get funky:

-- ...
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE NoStarIsType #-}
-- For defining fieldWidth and fieldOffset
{-# LANGUAGE AllowAmbiguousTypes #-}

And give our U type a few more instances:

-- ...
deriving newtype instance Enum (FindSmallest n HostStorage) => Enum (U n)
deriving newtype instance Bounded (FindSmallest n HostStorage) => Bounded (U n)
deriving newtype instance Bits (FindSmallest n HostStorage) => Bits (U n)
deriving newtype instance FiniteBits (FindSmallest n HostStorage) => FiniteBits (U n)
deriving newtype instance Real (FindSmallest n HostStorage) => Real (U n)
deriving newtype instance Integral (FindSmallest n HostStorage) => Integral (U n)

And we also need this handy typealias:

type KnownU (n :: Nat) =
    ( KnownNat n
    -- , Eq (U n)
    -- , Ord (U n)
    -- , Num (U n)
    -- , Show (U n )
    -- , Enum (U n)
    -- , Bounded (U n)
    , Bits (U n)
    -- , FiniteBits (U n)
    , Real (U n)
    , Integral (U n)
    )

Why have we done this? So we can do this:

-- A named and sized field
-- NOTE: We are using Symbol here, but we could also have used data kinds. 
newtype Field (name :: Symbol) (width :: Nat) = Field (U width)

-- An unpacked collection of fields
data Bitfield (fs :: [Type]) where
    Nil  :: Bitfield '[]
    Cons :: Field name width -> Bitfield fs -> Bitfield (Field name width : fs)

-- A class for packing things into fields
class (KnownU (PackedWidth fs)) => Packable (fs :: [Type]) where
    type PackedWidth fs :: Nat
    packFields :: Bitfield fs -> U (PackedWidth fs)

instance Packable '[] where
    type PackedWidth '[] = 0
    packFields Nil = 0

instance
    ( KnownU width, KnownU (width + PackedWidth fs), Packable fs
    ) => Packable (Field name width : fs) where
    type PackedWidth (Field name width : fs) = width + PackedWidth fs
    packFields (Cons (Field val) rest) = 
        let offset  = fromIntegral $ natVal (Proxy @(PackedWidth fs))
            bits    = fromIntegral val `shiftL` offset
        in  bits .|. fromIntegral (packFields rest)

-- A newtype for packed fields
newtype Packed (fs :: [Type]) = Packed { getPacked :: U (PackedWidth fs) }

deriving newtype instance Eq   (U (PackedWidth fs)) => Eq   (Packed fs)
deriving newtype instance Ord  (U (PackedWidth fs)) => Ord  (Packed fs)
deriving newtype instance Num  (U (PackedWidth fs)) => Num  (Packed fs)
deriving newtype instance Show (U (PackedWidth fs)) => Show (Packed fs)
deriving newtype instance Enum (U (PackedWidth fs)) => Enum (Packed fs)
deriving newtype instance Bounded (U (PackedWidth fs)) => Bounded (Packed fs)
deriving newtype instance Bits (U (PackedWidth fs)) => Bits (Packed fs)
deriving newtype instance FiniteBits (U (PackedWidth fs)) => FiniteBits (Packed fs)
deriving newtype instance Real (U (PackedWidth fs)) => Real (Packed fs)
deriving newtype instance Integral (U (PackedWidth fs)) => Integral (Packed fs)

-- A convenient function
pack :: (Packable fs) => Bitfield fs -> Packed fs
pack fs = Packed $ packFields fs

-- A color format
type R5G6B5 = '[Field "Red" 5, Field "Green" 6, Field "Blue" 5]

-- A packed color value
packR5G6B5 :: U 5 -> U 6 -> U 5 -> Packed R5G6B5
packR5G6B5 r g b = pack
    $ Cons (Field @"Red" r)
    $ Cons (Field @"Green" g)
    $ Cons (Field @"Blue" b)
    $ Nil

-- Example:
-- ghci> packR5G6B5 1 2 3
-- 2115

Alrighty, so that’s great for packing values into a bitfield, how about getting them out again? We use some type families and a typeclass to provide convenient access!

type family FieldOffset (name :: Symbol) (fs :: [Type]) :: Nat where
    FieldOffset name (Field name w : fs) = PackedWidth fs
    FieldOffset name (Field other w : fs) = FieldOffset name fs
    FieldOffset name '[] = TypeError
        (     'Text "FieldOffset: Field not found: "
        ':<>: 'ShowType name
        ) -- Or '0'

type family FieldWidth (name :: Symbol) (fs :: [Type]) :: Nat where
    FieldWidth name (Field name w : fs) = w
    FieldWidth name (Field other w : fs) = FieldWidth name fs
    FieldWidth name '[] = TypeError
        (     'Text "FieldWidth: Field not found: "
        ':<>: 'ShowType name
        ) -- Or '0'

fieldWidth :: forall name fs. (KnownU (FieldWidth name fs)) => Int
fieldWidth = fromIntegral $ natVal (Proxy @(FieldWidth name fs))

fieldOffset :: forall name fs. (KnownU (FieldOffset name fs)) => Int
fieldOffset = fromIntegral $ natVal (Proxy @(FieldOffset name fs))

-- A class for accessing specific fields from within packed data
class
    ( KnownU (FieldOffset name fs), KnownU (FieldWidth name fs), Packable fs
    ) => PluckField (name :: Symbol) (fs :: [Type]) where
    pluck :: Packed fs -> U (FieldWidth name fs)

instance
    ( KnownU (FieldWidth name fs), KnownU (FieldOffset name fs), Packable fs
    ) => PluckField name fs where
    pluck (Packed val) =
        let offset = fieldOffset @name @fs
            width  = fieldWidth @name @fs
            mask   = (1 `shiftL` width) - 1
            bits = val `shiftR` offset
        in fromIntegral (mask .&. bits)

pluckGreen :: Packed R5G6B5 -> U 6
pluckGreen = pluck @"Green"

-- Example:
-- ghci> pluckGreen $ packR5G6B5 1 2 3
-- 2

And viola! :partying_face: