{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE QuasiQuotes #-}

{- | Packing values into bits of a stream of machine words.

A layout is a sequence of fields, each a number of bits, read and written
most significant bit first; a field may cross a word boundary, and 'word'
groups the fields that make up exactly one word. 'BitGet' asks for the next
word only when a field reaches it, so the layout itself decides how many words
are read. A failure names where it happened through 'label' \/ '<?>'.
-}
module Wrench.Machine.BitPack (
    -- * Packing
    BitPack (..),
    packWords,
    unpackWords,

    -- * Monads
    BitGet,
    BitPut,

    -- * Fields
    getField,
    putField,
    getEnum,
    putEnum,
    getPad,
    putPad,

    -- * Structure
    Layout (..),
    (<?>),

    -- * Running
    runBitGetWith,
    runBitGet,
    runBitPut,
) where

import Control.Monad (ap, liftM)
import Control.Monad.Except (throwError)
import Data.Bits (finiteBitSize, shiftL, shiftR, (.&.), (.|.))
import Data.String.Interpolate (i)
import Relude
import Relude.Extra (safeToEnum)
import Wrench.Machine.Types

-- * Packing

-- | A type with a bit layout of its own, so layouts compose by type:
-- @tag <- unpack@ inside a larger layout.
class BitPack a where
    unpack :: (IsWord w) => BitGet w a
    pack :: (IsWord w) => a -> BitPut w ()

-- | Pack a value into whole words.
packWords :: forall w a. (BitPack a, IsWord w) => a -> Either Text [w]
packWords :: forall w a. (BitPack a, IsWord w) => a -> Either Text [w]
packWords = BitPut w () -> Either Text [w]
forall w. IsWord w => BitPut w () -> Either Text [w]
runBitPut (BitPut w () -> Either Text [w])
-> (a -> BitPut w ()) -> a -> Either Text [w]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> BitPut w ()
forall w. IsWord w => a -> BitPut w ()
forall a w. (BitPack a, IsWord w) => a -> BitPut w ()
pack

-- | Unpack a value from exactly these words; fails if words are left over.
unpackWords :: forall a w. (BitPack a, IsWord w) => [w] -> Either Text a
unpackWords :: forall a w. (BitPack a, IsWord w) => [w] -> Either Text a
unpackWords [w]
words_ = do
    (a, rest) <- BitGet w a -> [w] -> Either Text (a, [w])
forall w a. IsWord w => BitGet w a -> [w] -> Either Text (a, [w])
runBitGetRest BitGet w a
forall w. IsWord w => BitGet w a
forall a w. (BitPack a, IsWord w) => BitGet w a
unpack [w]
words_
    unless (null rest) $ Left [i|words left over after decoding: #{length rest}|]
    return a

-- * Monads

data GetStep w a
    = Done a
    | NeedWord (Maybe w -> GetStep w a)
    | Fail Text

data GetState w = GetState
    { forall w. GetState w -> w
gsWord :: w
    , forall w. GetState w -> Int
gsUsed :: Int
    }

newtype BitGet w a = BitGet (GetState w -> GetStep w (a, GetState w))

instance Functor (BitGet w) where
    fmap :: forall a b. (a -> b) -> BitGet w a -> BitGet w b
fmap = (a -> b) -> BitGet w a -> BitGet w b
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM

instance Applicative (BitGet w) where
    pure :: forall a. a -> BitGet w a
pure a
a = (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
forall w a. (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
BitGet ((GetState w -> GetStep w (a, GetState w)) -> BitGet w a)
-> (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
forall a b. (a -> b) -> a -> b
$ \GetState w
st -> (a, GetState w) -> GetStep w (a, GetState w)
forall w a. a -> GetStep w a
Done (a
a, GetState w
st)
    <*> :: forall a b. BitGet w (a -> b) -> BitGet w a -> BitGet w b
(<*>) = BitGet w (a -> b) -> BitGet w a -> BitGet w b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad (BitGet w) where
    BitGet GetState w -> GetStep w (a, GetState w)
g >>= :: forall a b. BitGet w a -> (a -> BitGet w b) -> BitGet w b
>>= a -> BitGet w b
f = (GetState w -> GetStep w (b, GetState w)) -> BitGet w b
forall w a. (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
BitGet ((GetState w -> GetStep w (b, GetState w)) -> BitGet w b)
-> (GetState w -> GetStep w (b, GetState w)) -> BitGet w b
forall a b. (a -> b) -> a -> b
$ \GetState w
st -> GetStep w (a, GetState w) -> GetStep w (b, GetState w)
continue (GetState w -> GetStep w (a, GetState w)
g GetState w
st)
        where
            continue :: GetStep w (a, GetState w) -> GetStep w (b, GetState w)
continue (Done (a
a, GetState w
st')) = let BitGet GetState w -> GetStep w (b, GetState w)
h = a -> BitGet w b
f a
a in GetState w -> GetStep w (b, GetState w)
h GetState w
st'
            continue (NeedWord Maybe w -> GetStep w (a, GetState w)
k) = (Maybe w -> GetStep w (b, GetState w)) -> GetStep w (b, GetState w)
forall w a. (Maybe w -> GetStep w a) -> GetStep w a
NeedWord (GetStep w (a, GetState w) -> GetStep w (b, GetState w)
continue (GetStep w (a, GetState w) -> GetStep w (b, GetState w))
-> (Maybe w -> GetStep w (a, GetState w))
-> Maybe w
-> GetStep w (b, GetState w)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe w -> GetStep w (a, GetState w)
k)
            continue (Fail Text
err) = Text -> GetStep w (b, GetState w)
forall w a. Text -> GetStep w a
Fail Text
err

data PutState w = PutState
    { forall w. PutState w -> [w]
psWritten :: [w]
    , forall w. PutState w -> w
psWord :: w
    , forall w. PutState w -> Int
psUsed :: Int
    }

newtype BitPut w a = BitPut (StateT (PutState w) (Either Text) a)
    deriving (Functor (BitPut w)
Functor (BitPut w) =>
(forall a. a -> BitPut w a)
-> (forall a b. BitPut w (a -> b) -> BitPut w a -> BitPut w b)
-> (forall a b c.
    (a -> b -> c) -> BitPut w a -> BitPut w b -> BitPut w c)
-> (forall a b. BitPut w a -> BitPut w b -> BitPut w b)
-> (forall a b. BitPut w a -> BitPut w b -> BitPut w a)
-> Applicative (BitPut w)
forall w. Functor (BitPut w)
forall a. a -> BitPut w a
forall w a. a -> BitPut w a
forall a b. BitPut w a -> BitPut w b -> BitPut w a
forall a b. BitPut w a -> BitPut w b -> BitPut w b
forall a b. BitPut w (a -> b) -> BitPut w a -> BitPut w b
forall w a b. BitPut w a -> BitPut w b -> BitPut w a
forall w a b. BitPut w a -> BitPut w b -> BitPut w b
forall w a b. BitPut w (a -> b) -> BitPut w a -> BitPut w b
forall a b c.
(a -> b -> c) -> BitPut w a -> BitPut w b -> BitPut w c
forall w a b c.
(a -> b -> c) -> BitPut w a -> BitPut w b -> BitPut w c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall w a. a -> BitPut w a
pure :: forall a. a -> BitPut w a
$c<*> :: forall w a b. BitPut w (a -> b) -> BitPut w a -> BitPut w b
<*> :: forall a b. BitPut w (a -> b) -> BitPut w a -> BitPut w b
$cliftA2 :: forall w a b c.
(a -> b -> c) -> BitPut w a -> BitPut w b -> BitPut w c
liftA2 :: forall a b c.
(a -> b -> c) -> BitPut w a -> BitPut w b -> BitPut w c
$c*> :: forall w a b. BitPut w a -> BitPut w b -> BitPut w b
*> :: forall a b. BitPut w a -> BitPut w b -> BitPut w b
$c<* :: forall w a b. BitPut w a -> BitPut w b -> BitPut w a
<* :: forall a b. BitPut w a -> BitPut w b -> BitPut w a
Applicative, (forall a b. (a -> b) -> BitPut w a -> BitPut w b)
-> (forall a b. a -> BitPut w b -> BitPut w a)
-> Functor (BitPut w)
forall a b. a -> BitPut w b -> BitPut w a
forall a b. (a -> b) -> BitPut w a -> BitPut w b
forall w a b. a -> BitPut w b -> BitPut w a
forall w a b. (a -> b) -> BitPut w a -> BitPut w b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall w a b. (a -> b) -> BitPut w a -> BitPut w b
fmap :: forall a b. (a -> b) -> BitPut w a -> BitPut w b
$c<$ :: forall w a b. a -> BitPut w b -> BitPut w a
<$ :: forall a b. a -> BitPut w b -> BitPut w a
Functor, Applicative (BitPut w)
Applicative (BitPut w) =>
(forall a b. BitPut w a -> (a -> BitPut w b) -> BitPut w b)
-> (forall a b. BitPut w a -> BitPut w b -> BitPut w b)
-> (forall a. a -> BitPut w a)
-> Monad (BitPut w)
forall w. Applicative (BitPut w)
forall a. a -> BitPut w a
forall w a. a -> BitPut w a
forall a b. BitPut w a -> BitPut w b -> BitPut w b
forall a b. BitPut w a -> (a -> BitPut w b) -> BitPut w b
forall w a b. BitPut w a -> BitPut w b -> BitPut w b
forall w a b. BitPut w a -> (a -> BitPut w b) -> BitPut w b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall w a b. BitPut w a -> (a -> BitPut w b) -> BitPut w b
>>= :: forall a b. BitPut w a -> (a -> BitPut w b) -> BitPut w b
$c>> :: forall w a b. BitPut w a -> BitPut w b -> BitPut w b
>> :: forall a b. BitPut w a -> BitPut w b -> BitPut w b
$creturn :: forall w a. a -> BitPut w a
return :: forall a. a -> BitPut w a
Monad)

-- * Fields

-- | The next @n@ bits, unsigned.
getField :: forall w. (IsWord w) => Int -> BitGet w Int
getField :: forall w. IsWord w => Int -> BitGet w Int
getField Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxFieldBits = (GetState w -> GetStep w (Int, GetState w)) -> BitGet w Int
forall w a. (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
BitGet ((GetState w -> GetStep w (Int, GetState w)) -> BitGet w Int)
-> (GetState w -> GetStep w (Int, GetState w)) -> BitGet w Int
forall a b. (a -> b) -> a -> b
$ GetStep w (Int, GetState w)
-> GetState w -> GetStep w (Int, GetState w)
forall a b. a -> b -> a
const (GetStep w (Int, GetState w)
 -> GetState w -> GetStep w (Int, GetState w))
-> GetStep w (Int, GetState w)
-> GetState w
-> GetStep w (Int, GetState w)
forall a b. (a -> b) -> a -> b
$ Text -> GetStep w (Int, GetState w)
forall w a. Text -> GetStep w a
Fail [i|field width #{n} is outside 0..#{maxFieldBits}|]
    | Bool
otherwise = (GetState w -> GetStep w (Int, GetState w)) -> BitGet w Int
forall w a. (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
BitGet ((GetState w -> GetStep w (Int, GetState w)) -> BitGet w Int)
-> (GetState w -> GetStep w (Int, GetState w)) -> BitGet w Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> GetState w -> GetStep w (Int, GetState w)
go Int
n Int
0
    where
        go :: Int -> Int -> GetState w -> GetStep w (Int, GetState w)
go Int
0 Int
acc GetState w
st = (Int, GetState w) -> GetStep w (Int, GetState w)
forall w a. a -> GetStep w a
Done (Int
acc, GetState w
st)
        go Int
k Int
acc GetState{w
gsWord :: forall w. GetState w -> w
gsWord :: w
gsWord, Int
gsUsed :: forall w. GetState w -> Int
gsUsed :: Int
gsUsed}
            | Int
gsUsed Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== forall w. IsWord w => Int
wordBits @w = (Maybe w -> GetStep w (Int, GetState w))
-> GetStep w (Int, GetState w)
forall w a. (Maybe w -> GetStep w a) -> GetStep w a
NeedWord ((Maybe w -> GetStep w (Int, GetState w))
 -> GetStep w (Int, GetState w))
-> (Maybe w -> GetStep w (Int, GetState w))
-> GetStep w (Int, GetState w)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Maybe w -> GetStep w (Int, GetState w)
resume Int
k Int
acc
            | Bool
otherwise =
                let taken :: Int
taken = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
k (forall w. IsWord w => Int
wordBits @w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
gsUsed)
                    bits :: Int
bits = (w -> Int
forall w. IsWord w => w -> Int
unsigned w
gsWord Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` (forall w. IsWord w => Int
wordBits @w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
gsUsed Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
taken)) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
mask Int
taken
                 in Int -> Int -> GetState w -> GetStep w (Int, GetState w)
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
taken) ((Int
acc Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
taken) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
bits) GetState{w
gsWord :: w
gsWord :: w
gsWord, gsUsed :: Int
gsUsed = Int
gsUsed Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
taken}
        resume :: Int -> Int -> Maybe w -> GetStep w (Int, GetState w)
resume Int
_ Int
_ Maybe w
Nothing = Text -> GetStep w (Int, GetState w)
forall w a. Text -> GetStep w a
Fail Text
"out of words"
        resume Int
k Int
acc (Just w
w) = Int -> Int -> GetState w -> GetStep w (Int, GetState w)
go Int
k Int
acc GetState{gsWord :: w
gsWord = w
w, gsUsed :: Int
gsUsed = Int
0}

-- | Write @value@ into the next @n@ bits.
putField :: forall w. (IsWord w) => Int -> Int -> BitPut w ()
putField :: forall w. IsWord w => Int -> Int -> BitPut w ()
putField Int
n Int
value
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxFieldBits = StateT (PutState w) (Either Text) () -> BitPut w ()
forall w a. StateT (PutState w) (Either Text) a -> BitPut w a
BitPut (StateT (PutState w) (Either Text) () -> BitPut w ())
-> StateT (PutState w) (Either Text) () -> BitPut w ()
forall a b. (a -> b) -> a -> b
$ Text -> StateT (PutState w) (Either Text) ()
forall a. Text -> StateT (PutState w) (Either Text) a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError [i|field width #{n} is outside 0..#{maxFieldBits}|]
    | Int
value Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
value Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Int
mask Int
n = StateT (PutState w) (Either Text) () -> BitPut w ()
forall w a. StateT (PutState w) (Either Text) a -> BitPut w a
BitPut (StateT (PutState w) (Either Text) () -> BitPut w ())
-> StateT (PutState w) (Either Text) () -> BitPut w ()
forall a b. (a -> b) -> a -> b
$ Text -> StateT (PutState w) (Either Text) ()
forall a. Text -> StateT (PutState w) (Either Text) a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError [i|#{value} doesn't fit in #{n} bits|]
    | Bool
otherwise = StateT (PutState w) (Either Text) () -> BitPut w ()
forall w a. StateT (PutState w) (Either Text) a -> BitPut w a
BitPut (StateT (PutState w) (Either Text) () -> BitPut w ())
-> StateT (PutState w) (Either Text) () -> BitPut w ()
forall a b. (a -> b) -> a -> b
$ Int -> StateT (PutState w) (Either Text) ()
go Int
n
    where
        go :: Int -> StateT (PutState w) (Either Text) ()
go Int
0 = () -> StateT (PutState w) (Either Text) ()
forall a. a -> StateT (PutState w) (Either Text) a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        go Int
k = do
            PutState{psWritten, psWord, psUsed} <- StateT (PutState w) (Either Text) (PutState w)
forall s (m :: * -> *). MonadState s m => m s
get
            let taken = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
k (forall w. IsWord w => Int
wordBits @w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
psUsed)
                bits = (Int
value Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
taken)) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
mask Int
taken
                word' = w
psWord w -> w -> w
forall a. Bits a => a -> a -> a
.|. (Int -> w
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bits w -> Int -> w
forall a. Bits a => a -> Int -> a
`shiftL` (forall w. IsWord w => Int
wordBits @w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
psUsed Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
taken))
                used' = Int
psUsed Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
taken
            put $
                if used' == wordBits @w
                    then PutState{psWritten = word' : psWritten, psWord = 0, psUsed = 0}
                    else PutState{psWritten, psWord = word', psUsed = used'}
            go (k - taken)

-- | The next @n@ bits as a constructor of @e@; fails on a value no
-- constructor has.
getEnum :: forall e w. (Bounded e, Enum e, IsWord w) => Int -> BitGet w e
getEnum :: forall e w. (Bounded e, Enum e, IsWord w) => Int -> BitGet w e
getEnum Int
n = do
    value <- Int -> BitGet w Int
forall w. IsWord w => Int -> BitGet w Int
getField Int
n
    maybe (BitGet $ const $ Fail [i|unknown tag #{value}|]) return (safeToEnum value)

-- | Write a constructor of @e@ into the next @n@ bits.
putEnum :: (Enum e, IsWord w) => Int -> e -> BitPut w ()
putEnum :: forall e w. (Enum e, IsWord w) => Int -> e -> BitPut w ()
putEnum Int
n = Int -> Int -> BitPut w ()
forall w. IsWord w => Int -> Int -> BitPut w ()
putField Int
n (Int -> BitPut w ()) -> (e -> Int) -> e -> BitPut w ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Int
forall a. Enum a => a -> Int
fromEnum

-- | Skip @n@ bits, any number of them.
getPad :: (IsWord w) => Int -> BitGet w ()
getPad :: forall w. IsWord w => Int -> BitGet w ()
getPad = (Int -> BitGet w Int) -> [Int] -> BitGet w ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Int -> BitGet w Int
forall w. IsWord w => Int -> BitGet w Int
getField ([Int] -> BitGet w ()) -> (Int -> [Int]) -> Int -> BitGet w ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int]
padChunks

-- | @n@ zero bits, any number of them.
putPad :: (IsWord w) => Int -> BitPut w ()
putPad :: forall w. IsWord w => Int -> BitPut w ()
putPad = (Int -> BitPut w ()) -> [Int] -> BitPut w ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Int -> Int -> BitPut w ()
forall w. IsWord w => Int -> Int -> BitPut w ()
`putField` Int
0) ([Int] -> BitPut w ()) -> (Int -> [Int]) -> Int -> BitPut w ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int]
padChunks

-- * Structure

-- | Combinators that work in both directions.
class Layout f where
    -- | Run a layout that must cover exactly one word, starting on a word
    -- boundary.
    word :: f a -> f a

    -- | Prefix a failure with where it happened; nested labels read
    -- outermost first.
    label :: Text -> f a -> f a

instance (IsWord w) => Layout (BitGet w) where
    word :: forall a. BitGet w a -> BitGet w a
word (BitGet GetState w -> GetStep w (a, GetState w)
inner) = (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
forall w a. (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
BitGet ((GetState w -> GetStep w (a, GetState w)) -> BitGet w a)
-> (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
forall a b. (a -> b) -> a -> b
$ \GetState w
st ->
        if GetState w -> Int
forall w. GetState w -> Int
gsUsed GetState w
st Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= forall w. IsWord w => Int
wordBits @w
            then Text -> GetStep w (a, GetState w)
forall w a. Text -> GetStep w a
Fail Text
"word layout doesn't start on a word boundary"
            else Int -> GetStep w (a, GetState w) -> GetStep w (a, GetState w)
go (Int
0 :: Int) (GetState w -> GetStep w (a, GetState w)
inner GetState w
st)
        where
            go :: Int -> GetStep w (a, GetState w) -> GetStep w (a, GetState w)
go Int
n (NeedWord Maybe w -> GetStep w (a, GetState w)
k) = (Maybe w -> GetStep w (a, GetState w)) -> GetStep w (a, GetState w)
forall w a. (Maybe w -> GetStep w a) -> GetStep w a
NeedWord (Int -> GetStep w (a, GetState w) -> GetStep w (a, GetState w)
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (GetStep w (a, GetState w) -> GetStep w (a, GetState w))
-> (Maybe w -> GetStep w (a, GetState w))
-> Maybe w
-> GetStep w (a, GetState w)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe w -> GetStep w (a, GetState w)
k)
            go Int
n (Done (a
a, GetState w
st'))
                | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 Bool -> Bool -> Bool
&& GetState w -> Int
forall w. GetState w -> Int
gsUsed GetState w
st' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== forall w. IsWord w => Int
wordBits @w = (a, GetState w) -> GetStep w (a, GetState w)
forall w a. a -> GetStep w a
Done (a
a, GetState w
st')
                | Bool
otherwise = Text -> GetStep w (a, GetState w)
forall w a. Text -> GetStep w a
Fail [i|word layout doesn't fill exactly one #{wordBits @w}-bit word|]
            go Int
_ (Fail Text
err) = Text -> GetStep w (a, GetState w)
forall w a. Text -> GetStep w a
Fail Text
err
    label :: forall a. Text -> BitGet w a -> BitGet w a
label Text
name (BitGet GetState w -> GetStep w (a, GetState w)
g) = (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
forall w a. (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
BitGet ((GetState w -> GetStep w (a, GetState w)) -> BitGet w a)
-> (GetState w -> GetStep w (a, GetState w)) -> BitGet w a
forall a b. (a -> b) -> a -> b
$ GetStep w (a, GetState w) -> GetStep w (a, GetState w)
relabel (GetStep w (a, GetState w) -> GetStep w (a, GetState w))
-> (GetState w -> GetStep w (a, GetState w))
-> GetState w
-> GetStep w (a, GetState w)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GetState w -> GetStep w (a, GetState w)
g
        where
            relabel :: GetStep w (a, GetState w) -> GetStep w (a, GetState w)
relabel (Fail Text
err) = Text -> GetStep w (a, GetState w)
forall w a. Text -> GetStep w a
Fail (Text -> GetStep w (a, GetState w))
-> Text -> GetStep w (a, GetState w)
forall a b. (a -> b) -> a -> b
$ Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
err
            relabel (NeedWord Maybe w -> GetStep w (a, GetState w)
k) = (Maybe w -> GetStep w (a, GetState w)) -> GetStep w (a, GetState w)
forall w a. (Maybe w -> GetStep w a) -> GetStep w a
NeedWord (GetStep w (a, GetState w) -> GetStep w (a, GetState w)
relabel (GetStep w (a, GetState w) -> GetStep w (a, GetState w))
-> (Maybe w -> GetStep w (a, GetState w))
-> Maybe w
-> GetStep w (a, GetState w)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe w -> GetStep w (a, GetState w)
k)
            relabel GetStep w (a, GetState w)
done = GetStep w (a, GetState w)
done

instance (IsWord w) => Layout (BitPut w) where
    word :: forall a. BitPut w a -> BitPut w a
word (BitPut StateT (PutState w) (Either Text) a
inner) = StateT (PutState w) (Either Text) a -> BitPut w a
forall w a. StateT (PutState w) (Either Text) a -> BitPut w a
BitPut (StateT (PutState w) (Either Text) a -> BitPut w a)
-> StateT (PutState w) (Either Text) a -> BitPut w a
forall a b. (a -> b) -> a -> b
$ do
        before <- StateT (PutState w) (Either Text) (PutState w)
forall s (m :: * -> *). MonadState s m => m s
get
        unless (psUsed before == 0) $
            throwError "word layout doesn't start on a word boundary"
        a <- inner
        after <- get
        unless (psUsed after == 0 && length (psWritten after) == length (psWritten before) + 1) $
            throwError [i|word layout doesn't fill exactly one #{wordBits @w}-bit word|]
        return a
    label :: forall a. Text -> BitPut w a -> BitPut w a
label Text
name (BitPut StateT (PutState w) (Either Text) a
body) = StateT (PutState w) (Either Text) a -> BitPut w a
forall w a. StateT (PutState w) (Either Text) a -> BitPut w a
BitPut (StateT (PutState w) (Either Text) a -> BitPut w a)
-> StateT (PutState w) (Either Text) a -> BitPut w a
forall a b. (a -> b) -> a -> b
$ (PutState w -> Either Text (a, PutState w))
-> StateT (PutState w) (Either Text) a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((PutState w -> Either Text (a, PutState w))
 -> StateT (PutState w) (Either Text) a)
-> (PutState w -> Either Text (a, PutState w))
-> StateT (PutState w) (Either Text) a
forall a b. (a -> b) -> a -> b
$ (Text -> Text)
-> Either Text (a, PutState w) -> Either Text (a, PutState w)
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": ") Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) (Either Text (a, PutState w) -> Either Text (a, PutState w))
-> (PutState w -> Either Text (a, PutState w))
-> PutState w
-> Either Text (a, PutState w)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StateT (PutState w) (Either Text) a
-> PutState w -> Either Text (a, PutState w)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT (PutState w) (Either Text) a
body

infixl 1 <?>

-- | @layout <?> name@ is 'label' with the name last, after the field it names.
(<?>) :: (Layout f) => f a -> Text -> f a
<?> :: forall (f :: * -> *) a. Layout f => f a -> Text -> f a
(<?>) = (Text -> f a -> f a) -> f a -> Text -> f a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Text -> f a -> f a
forall a. Text -> f a -> f a
forall (f :: * -> *) a. Layout f => Text -> f a -> f a
label

-- * Running

-- | Decode, fetching word @k@ of the stream with @fetch k@ when the layout
-- reaches it.
runBitGetWith :: forall w m a. (IsWord w, Monad m) => (Int -> m w) -> BitGet w a -> m (Either Text a)
runBitGetWith :: forall w (m :: * -> *) a.
(IsWord w, Monad m) =>
(Int -> m w) -> BitGet w a -> m (Either Text a)
runBitGetWith Int -> m w
fetch (BitGet GetState w -> GetStep w (a, GetState w)
body) = Int -> GetStep w (a, GetState w) -> m (Either Text a)
go Int
0 (GetStep w (a, GetState w) -> m (Either Text a))
-> GetStep w (a, GetState w) -> m (Either Text a)
forall a b. (a -> b) -> a -> b
$ GetState w -> GetStep w (a, GetState w)
body GetState w
forall w. IsWord w => GetState w
emptyGetState
    where
        go :: Int -> GetStep w (a, GetState w) -> m (Either Text a)
go Int
_ (Done (a
a, GetState w
_)) = Either Text a -> m (Either Text a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Text a -> m (Either Text a))
-> Either Text a -> m (Either Text a)
forall a b. (a -> b) -> a -> b
$ a -> Either Text a
forall a b. b -> Either a b
Right a
a
        go Int
_ (Fail Text
err) = Either Text a -> m (Either Text a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Text a -> m (Either Text a))
-> Either Text a -> m (Either Text a)
forall a b. (a -> b) -> a -> b
$ Text -> Either Text a
forall a b. a -> Either a b
Left Text
err
        go Int
k (NeedWord Maybe w -> GetStep w (a, GetState w)
next) = Int -> m w
fetch Int
k m w -> (w -> m (Either Text a)) -> m (Either Text a)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> GetStep w (a, GetState w) -> m (Either Text a)
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (GetStep w (a, GetState w) -> m (Either Text a))
-> (w -> GetStep w (a, GetState w)) -> w -> m (Either Text a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe w -> GetStep w (a, GetState w)
next (Maybe w -> GetStep w (a, GetState w))
-> (w -> Maybe w) -> w -> GetStep w (a, GetState w)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. w -> Maybe w
forall a. a -> Maybe a
Just

-- | Decode a list of words; fails if the layout runs past them.
runBitGet :: forall w a. (IsWord w) => BitGet w a -> [w] -> Either Text a
runBitGet :: forall w a. IsWord w => BitGet w a -> [w] -> Either Text a
runBitGet BitGet w a
decoder = ((a, [w]) -> a) -> Either Text (a, [w]) -> Either Text a
forall a b. (a -> b) -> Either Text a -> Either Text b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, [w]) -> a
forall a b. (a, b) -> a
fst (Either Text (a, [w]) -> Either Text a)
-> ([w] -> Either Text (a, [w])) -> [w] -> Either Text a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BitGet w a -> [w] -> Either Text (a, [w])
forall w a. IsWord w => BitGet w a -> [w] -> Either Text (a, [w])
runBitGetRest BitGet w a
decoder

runBitGetRest :: forall w a. (IsWord w) => BitGet w a -> [w] -> Either Text (a, [w])
runBitGetRest :: forall w a. IsWord w => BitGet w a -> [w] -> Either Text (a, [w])
runBitGetRest (BitGet GetState w -> GetStep w (a, GetState w)
body) [w]
words_ = [w] -> GetStep w (a, GetState w) -> Either Text (a, [w])
forall {w} {a} {b}. [w] -> GetStep w (a, b) -> Either Text (a, [w])
go [w]
words_ (GetStep w (a, GetState w) -> Either Text (a, [w]))
-> GetStep w (a, GetState w) -> Either Text (a, [w])
forall a b. (a -> b) -> a -> b
$ GetState w -> GetStep w (a, GetState w)
body GetState w
forall w. IsWord w => GetState w
emptyGetState
    where
        go :: [w] -> GetStep w (a, b) -> Either Text (a, [w])
go [w]
rest (Done (a
a, b
_)) = (a, [w]) -> Either Text (a, [w])
forall a b. b -> Either a b
Right (a
a, [w]
rest)
        go [w]
_ (Fail Text
err) = Text -> Either Text (a, [w])
forall a b. a -> Either a b
Left Text
err
        go [w]
rest (NeedWord Maybe w -> GetStep w (a, b)
next) = case [w]
rest of
            [] -> [w] -> GetStep w (a, b) -> Either Text (a, [w])
go [] (GetStep w (a, b) -> Either Text (a, [w]))
-> GetStep w (a, b) -> Either Text (a, [w])
forall a b. (a -> b) -> a -> b
$ Maybe w -> GetStep w (a, b)
next Maybe w
forall a. Maybe a
Nothing
            w
w : [w]
rest' -> [w] -> GetStep w (a, b) -> Either Text (a, [w])
go [w]
rest' (GetStep w (a, b) -> Either Text (a, [w]))
-> GetStep w (a, b) -> Either Text (a, [w])
forall a b. (a -> b) -> a -> b
$ Maybe w -> GetStep w (a, b)
next (Maybe w -> GetStep w (a, b)) -> Maybe w -> GetStep w (a, b)
forall a b. (a -> b) -> a -> b
$ w -> Maybe w
forall a. a -> Maybe a
Just w
w

emptyGetState :: forall w. (IsWord w) => GetState w
emptyGetState :: forall w. IsWord w => GetState w
emptyGetState = GetState{gsWord :: w
gsWord = w
0, gsUsed :: Int
gsUsed = forall w. IsWord w => Int
wordBits @w}

emptyPutState :: (IsWord w) => PutState w
emptyPutState :: forall w. IsWord w => PutState w
emptyPutState = PutState{psWritten :: [w]
psWritten = [], psWord :: w
psWord = w
0, psUsed :: Int
psUsed = Int
0}

-- | Pack the fields into words; fails if a value doesn't fit its field or the
-- fields don't fill whole words.
runBitPut :: forall w. (IsWord w) => BitPut w () -> Either Text [w]
runBitPut :: forall w. IsWord w => BitPut w () -> Either Text [w]
runBitPut (BitPut StateT (PutState w) (Either Text) ()
body) = do
    PutState{psWritten, psUsed} <- PutState w
-> StateT (PutState w) (Either Text) () -> Either Text (PutState w)
forall (f :: * -> *) s a. Functor f => s -> StateT s f a -> f s
executingStateT PutState w
forall w. IsWord w => PutState w
emptyPutState StateT (PutState w) (Either Text) ()
body
    unless (psUsed == 0) $
        Left [i|#{length psWritten * wordBits @w + psUsed} bits don't fill whole #{wordBits @w}-bit words|]
    return $ reverse psWritten

-- * Helpers

wordBits :: forall w. (IsWord w) => Int
wordBits :: forall w. IsWord w => Int
wordBits = w -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (w
0 :: w)

unsigned :: (IsWord w) => w -> Int
unsigned :: forall w. IsWord w => w -> Int
unsigned = Unsign w -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Unsign w -> Int) -> (w -> Unsign w) -> w -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. w -> Unsign w
forall w. FromSign w => w -> Unsign w
fromSign

-- | The widest field whose value fits a non-negative 'Int'.
maxFieldBits :: Int
maxFieldBits :: Int
maxFieldBits = Int -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Int
0 :: Int) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

padChunks :: Int -> [Int]
padChunks :: Int -> [Int]
padChunks Int
n = Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
maxFieldBits) Int
maxFieldBits [Int] -> [Int] -> [Int]
forall a. Semigroup a => a -> a -> a
<> [Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
maxFieldBits | Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
maxFieldBits Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0]

mask :: Int -> Int
mask :: Int -> Int
mask Int
n = (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
n) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1