Word and Endian classes for folding over Bits and Bytes

I was deciding whether I was going to write my virtual image data in either Little- or Big- Endian format, and while doing so, I reached for some trusty old friends that are useful for grabbing a specific byte out of an arbitrarily-sized byte-addressible word, and placing them side by side, refactoring them slightly to express their similarity:

-- NOTE: Logical shifts in Bits use Little-Endian / LSB bit order
-- They are unaffected by architecture endianness
leByte :: (Integral i, Bits i) => Int -> Int -> i -> Word8
leByte _ 0 i = fromIntegral (i .&. 0xFF) -- shiftR i (8 * 0) probably compiles to this anyway
leByte _ n i = fromIntegral ((shiftR i (8 * n)) .&. 0xFF)

beByte :: (Integral i, Bits i) => Int -> Int -> i -> Word8
beByte sz n i = fromIntegral ((shiftR i (8 * ((sz - n) - 1))) .&. 0xFF)

In doing so, I realized I had accidentally solved one of the last problems in the memalloc library, which is the thorny issue of surfacing endianness in a way that unifies it with addressability - that is, how endianness affects bit and byte order, which is something that affects bit-shifts, which affects Bitwise, which is a core class in memalloc, because a group / compound word’s boolean instances comes from the instances of its subwords, showing how the two are united.

With this problem unsolved, memalloc has sort of sat there as I worked on other things that I knew I could make progress on, because I still have limited energy that I have to spend wisely due to my ongoing recovery. (I am continuing to improve, and have at least doubled the amount of time I can spend on code work lately).


What I had realized with my leByte and beByte was I needed to focus on expressing how chunks or subwords are extracted from a compound word, because that would allow me to completely ignore how they were actually ordered, thus abstracting the concept.

First, we’re going to enable a few friends:

{-# LANGUAGE
    DataKinds
,   AllowAmbiguousTypes
#-}

You see, if you squint, beBytes and leBytes are almost identical. Not only that, the functions using them are also almost identical:

leByte :: (Integral i, Bits i) => Int -> Int -> i -> Word8
leByte _ n i = fromIntegral ((shiftR i (8 * n)) .&. 0xFF)

leBytes :: (Integral i, Bits i) => Int -> i -> [Word8] -> [Word8]
leBytes sz d rbs = foldr (\ n acc -> leByte sz n d : acc) rbs [ 0 .. (sz - 1) ]

beByte :: (Integral i, Bits i) => Int -> Int -> i -> Word8
beByte sz n i = fromIntegral ((shiftR i (8 * ((sz - n) - 1))) .&. 0xFF)

beBytes :: (Integral i, Bits i) => Int -> i -> [Word8] -> [Word8]
beBytes sz d rbs = foldr (\ n acc -> beByte sz n d : acc) rbs [ 0 .. (sz - 1) ]

NOTE: I am going to be ignoring architecture endianness for brevity, but it can be used to give us even more zero-cost abstractions.

I had adapted them to this form to make it easier to use them with a type RawS = [Word8] -> [Word8] byte-builder, thus placing the accumulator as the last argument (rbs stands for remaining bytes). I had just never placed them together in this way before that so clearly expressed this, because the generalized form is obvious (and no surprise).

foldrBytes f sz d rbs = foldr (\ n acc -> f sz n d : acc) rbs [ 0 .. (sz - 1) ]

What isn’t obvious is that it actually folds over the indices, and not the bytes - the f function is actually supposed to access the byte via its canonical index. This actually tells us all we need to know to properly implement an abstraction that supports not just big- and little-endian bitmappings, but also esoteric things like middle-endian with 3-4-1-2 format and stuff.

In particular, it tells us how to finally unify BitSize and ByteSize, two classes that have no real point in existing separately, because one is the other, just 8 times larger. It tells us our goal - we need to build a foldWords function - or more pointedly, a foldWordsEndian function.


When dealing with types having statically known bit and byte sizes you might be tempted to reach for a type family to deal with this:

type family BitSize w :: Nat
type instance BitSize Bool   = 1
type instance BitSize Word8  = 8
type instance BitSize Word16 = 16
type instance BitSize Word32 = 32
type instance BitSize Word64 = 64
type instance BitSize Int8   = 8
type instance BitSize Int16  = 16
type instance BitSize Int32  = 32
type instance BitSize Int64  = 64

type family ByteSize w :: Nat
type instance ByteSize Bool   = 1
type instance ByteSize Word8  = 1
type instance ByteSize Word16 = 2
type instance ByteSize Word32 = 4
type instance ByteSize Word64 = 8
type instance ByteSize Int8   = 1
type instance ByteSize Int16  = 2
type instance ByteSize Int32  = 4
type instance ByteSize Int64  = 8

This is fine - its type-level, so zero runtime overhead which is nice - I like that. There are a few libraries that define type families for static sizes - its easy and it works, but it keeps the concept of bits and bytes separate - I don’t like that.

We can actually now unify them, and keep our zero-cost abstractions! We start with a root Word class that has a WordSize ie its size is measurable in the number of some sub-word.

-- Compound-word, sub-word
class (KnownNat (WordSize cw sw)) => Word cw sw where

    type WordSize cw sw :: Nat

    proxyWordSize :: Proxy cw -> Proxy sw -> Int
    proxyWordSize _ _ = fromIntegral (natVal (Proxy @(WordSize cw sw)))

-- NOTE: Requires AllowAmbiguousTypes
wordSize :: forall cw sw . Word cw sw => Int
wordSize = proxyWordSize (Proxy @cw) (Proxy @sw)

Because it is a multi-parameter class, we can have instances for both Bit and Byte / Bool and Word8:

-- BitSize
instance Word Bool Bool where
    type WordSize Bool Bool = 1

instance Word Word8 Bool where
    type WordSize Word8 Bool = 8

instance Word Word16 Bool where
    type WordSize Word16 Bool = 16

-- ByteSize
instance Word Word8 Word8 where
    type WordSize Word8 Word8 = 1

instance Word Word16 Word8 where
    type WordSize Word16 Word8 = 2

NOTE: We could auto-generate these with some more type magic but its fine for now

Using these is easy - either supply a proxy manually or use the convenience function:

> proxyWordSize (Proxy @Word16) (Proxy @Bool)
16
> wordSize @Word16 @Word8
2

To talk about endian-ness, we need to give our word’s indices / addresses some sort of order:

class (Word cw sw) => WordOrder endian cw sw where
    -- The native endian-ness, aka a bitmapping
    type Endian cw sw :: *
    type Endian cw sw = BigEndian -- We're assuming big-endian to match literal order eg 0x01234567
    
    proxyWordAddr :: Proxy endian -> Proxy cw -> Proxy sw -> Int -> Int

wordAddrEndian :: forall endian cw sw . WordOrder endian cw sw => Int -> Int
wordAddrEndian = proxyWordAddr (Proxy @endian) (Proxy @cw) (Proxy @sw)

wordAddr :: forall cw sw . WordOrder (Endian cw sw) cw sw => Int -> Int
wordAddr = wordAddrEndian @(Endian cw sw) @cw @sw

This provides a way of transforming the index into an endian-specific address. On top of that, we layer the actual ability to fetch individual words:

class (WordOrder endian cw sw) => WordAccess endian cw sw where
    proxyWord :: Proxy endian -> Proxy cw -> Proxy sw -> cw -> Int -> sw

wordEndian :: forall endian cw sw . WordAccess endian cw sw => cw -> Int -> sw
wordEndian = proxyWord (Proxy @endian) (Proxy @cw) (Proxy @sw)

word :: forall cw sw . (WordAccess (Endian cw sw) cw sw) => cw -> Int -> sw
word = wordEndian @(Endian cw sw)

Then, we can define our endian bitmappings:

data BigEndian
data LittleEndian

They are just type tags! So, how does this work? We’ll define instances for Word8 and Word16 as a demonstration:

-- Word8

instance Word Word8 Word8 where
    type WordSize Word8 Word8 = 1

instance WordOrder BigEndian Word8 Word8 where
    proxyWordAddr _ _ _ idx = idx

instance WordAccess BigEndian Word8 Word8 where
    proxyWord _ pcw psw cw swi = sw where
        sz = proxyWordSize pcw psw
        sw = beByte sz swi cw

instance WordOrder LittleEndian Word8 Word8 where
    proxyWordAddr _ pcw psw idx = addr where
        sz = proxyWordSize pcw psw
        addr = (sz - 1) - idx

instance WordAccess LittleEndian Word8 Word8 where
    proxyWord _ pcw psw cw swi = sw where
        sz = proxyWordSize pcw psw
        sw = leByte sz swi cw

-- Word16
instance Word Word16 Word8 where
    type WordSize Word16 Word8 = 2

instance WordOrder BigEndian Word16 Word8 where
    proxyWordAddr _ _ _ idx = idx

instance WordAccess BigEndian Word16 Word8 where
    proxyWord _ pcw psw cw swi = sw where
        sz = proxyWordSize pcw psw
        sw = beByte sz swi cw

instance WordOrder LittleEndian Word16 Word8 where
    proxyWordAddr _ pcw psw idx = addr where
        sz = proxyWordSize pcw psw
        addr = (sz - 1) - idx

instance WordAccess LittleEndian Word16 Word8 where
    proxyWord _ pcw psw cw swi = sw where
        sz = proxyWordSize pcw psw
        sw = leByte sz swi cw

NOTE: I didn’t actually use proxyWordAddr because I already had leByte and beByte available. Otherwise the WordAccess instances would all be something like:

instance WordAccess BigEndian Word64 Word8 where
    proxyWord pe pcw psw cw swi = sw where
        addr = proxyWordAddr pe pcw psw swi
        sw = specializedGetWord8 addr cw

where for CPU-atomic words you might want specializedGetWord8 to be testByte addr i = fromIntegral ((shiftR i (8 * addr)) .&. 0xFF) but eg for a large array of bytes (say a Word1024) you might want to pluck it out directly via pointer magic instead of shifting eg hundreds of bytes for no reason.

You’ll notice, the code is the same for both types. Extremely easy to generate. And, extremely easy to use:

w = 0x0123456789ABCDEF :: Word64
-- Word64 literal is BigEndian by default
word @Word64 @Word8 w 0
> 1 -- 0x01
word @Word64 @Word8 w 7
> 239 -- 0xEF
wordEndian @LittleEndian @Word64 @Word8 w 0
> 239
wordEndian @LittleEndian @Word64 @Word8 w 7
> 1

In fact, we can take it even further, and implement a lovely endian-generic, word-generic equivalent of foldr:

foldrWordsEndian
    :: forall endian cw sw a . (WordAccess endian cw sw)
    => (sw -> a -> a) -> a -> cw -> a
foldrWordsEndian f z cw = foldr (\ i acc -> f (wordEndian @endian cw i) acc) z [0 .. wordSize @cw @sw - 1]

toWordsEndian :: forall endian cw sw . (WordAccess endian cw sw) => cw -> [sw]
toWordsEndian = foldrWordsEndian @endian @cw @sw (:) []

foldrWords :: forall cw sw a . (WordAccess (Endian cw sw) cw sw)
    => (sw -> a -> a) -> a -> cw -> a
foldrWords = foldrWordsEndian @(Endian cw sw) @cw @sw @a

toWords :: forall cw sw . (WordAccess (Endian cw sw) cw sw) => cw -> [sw]
toWords = foldrWords (:) []

We can test it out! I’ve omitted the Word64 instances for brevity but they’re there, and we can do cool things with them:

> foldrWordsEndian @BigEndian @Word64 @Word8 (:) [] w
[1,35,69,103,137,171,205,239]
> toWordsEndian @LittleEndian @Word64 @Word8 w
[239,205,171,137,103,69,35,1]
> toWords @Word64 @Word8 w
[1,35,69,103,137,171,205,239]

What’s great about this is, if I wanted to also implement Word Word64 Bool I haven’t blocked myself off by implementing Word Word64 Word8 - that is, this provides both Bit-addressability AND Byte-addressability, which is something that has eluded me for some time!

NOTE: This isn’t just strictly for binary integers, either - it can support other base word types, like Word Float4x16 Float16, or even Word Qubyte Qubit or more exotic architectures.

There’s still things to do, for example this is actually all more roughly analogous to FiniteBits and ye olde memory:Data.ByteArray.Sized classes for fixed-width data structures, because of the use of KnownNat - there will need to be a separate weaker set of classes for dynamic-width things like ByteString, but that’s exactly why leBytes and beBytes and foldrBytes all take a value-level size argument, we can just do the same things after querying the value for its size instead of the type.

But now memalloc is officially unblocked, and has endian and word classes congruent with the Bitwise and Boolean stuff, and now it is getting pulled into the new compiler as one of its base libraries! :partying_face:

5 Likes