{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}


-- |Utilities to represent and display bit sequences
module Flat.Bits (
    Bits,
    toBools,
    fromBools,
    bits,
    paddedBits,
    asBytes,
    asBits,
    takeBits,
    takeAllBits,
) where
-- TODO: AsBits Class?

import           Data.Bits                      (FiniteBits (finiteBitSize),
                                                 testBit)
import qualified Data.ByteString                as B
import qualified Data.Vector.Unboxed            as V
import           Data.Word                      (Word8)
import           Flat.Class                     (Flat)
import           Flat.Decoder                   (Decoded)
import           Flat.Filler                    (PostAligned (PostAligned),
                                                 fillerLength)
import           Flat.Run                       (flat, unflatRaw)
import           Text.PrettyPrint.HughesPJClass (Doc, Pretty (pPrint), hsep,
                                                 text)

-- |A sequence of bits
type Bits = V.Vector Bool

toBools :: Bits -> [Bool]
toBools :: Bits -> [Bool]
toBools = forall a. Unbox a => Vector a -> [a]
V.toList

fromBools :: [Bool] -> Bits
fromBools :: [Bool] -> Bits
fromBools = forall a. Unbox a => [a] -> Vector a
V.fromList

{- $setup
>>> import Data.Word
>>> import Flat.Instances.Base
>>> import Flat.Instances.Test(tst,prettyShow)
-}

{- |The sequence of bits corresponding to the serialization of the passed value (without any final byte padding)

>>> bits True
[True]
-}
bits :: forall a. Flat a => a -> Bits
bits :: forall a. Flat a => a -> Bits
bits a
v =
    let lbs :: ByteString
lbs = forall a. Flat a => a -> ByteString
flat a
v
    in case forall a b. (Flat a, AsByteString b) => b -> Decoded a
unflatRaw ByteString
lbs :: Decoded (PostAligned a) of 
            Right (PostAligned a
_ Filler
f) -> Int -> ByteString -> Bits
takeBits (Int
8 forall a. Num a => a -> a -> a
* ByteString -> Int
B.length ByteString
lbs forall a. Num a => a -> a -> a
- forall a. Num a => Filler -> a
fillerLength Filler
f) ByteString
lbs
            Left DecodeException
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"incorrect coding or decoding, please inform the maintainer of this package"

{- |The sequence of bits corresponding to the byte-padded serialization of the passed value

>>> paddedBits True
[True,False,False,False,False,False,False,True]
-}
paddedBits :: forall a. Flat a => a -> Bits
paddedBits :: forall a. Flat a => a -> Bits
paddedBits a
v = let lbs :: ByteString
lbs = forall a. Flat a => a -> ByteString
flat a
v in ByteString -> Bits
takeAllBits ByteString
lbs

takeAllBits :: B.ByteString -> Bits
takeAllBits :: ByteString -> Bits
takeAllBits ByteString
lbs= Int -> ByteString -> Bits
takeBits (Int
8 forall a. Num a => a -> a -> a
* ByteString -> Int
B.length ByteString
lbs) ByteString
lbs

takeBits :: Int -> B.ByteString -> Bits
takeBits :: Int -> ByteString -> Bits
takeBits Int
numBits ByteString
lbs =
    forall a. Unbox a => Int -> (Int -> a) -> Vector a
V.generate
        (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
numBits)
        ( \Int
n ->
            let (Int
bb, Int
b) = Int
n forall a. Integral a => a -> a -> (a, a)
`divMod` Int
8
             in forall a. Bits a => a -> Int -> Bool
testBit (HasCallStack => ByteString -> Int -> Word8
B.index ByteString
lbs (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bb)) (Int
7 forall a. Num a => a -> a -> a
- Int
b)
        )

{- |Convert an integral value to its equivalent bit representation

>>> asBits (5::Word8)
[False,False,False,False,False,True,False,True]
-}
asBits :: FiniteBits a => a -> Bits
asBits :: forall a. FiniteBits a => a -> Bits
asBits a
w = let s :: Int
s = forall b. FiniteBits b => b -> Int
finiteBitSize a
w in forall a. Unbox a => Int -> (Int -> a) -> Vector a
V.generate Int
s (forall a. Bits a => a -> Int -> Bool
testBit a
w forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
s forall a. Num a => a -> a -> a
- Int
1 forall a. Num a => a -> a -> a
-))

{- |Convert a sequence of bits to the corresponding list of bytes

>>> asBytes $ asBits (256+3::Word16)
[1,3]
-}
asBytes :: Bits -> [Word8]
asBytes :: Bits -> [Word8]
asBytes = forall a b. (a -> b) -> [a] -> [b]
map [Bool] -> Word8
byteVal forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t. [t] -> [[t]]
bytes forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Unbox a => Vector a -> [a]
V.toList

-- |Convert to the corresponding value (most significant bit first)
byteVal :: [Bool] -> Word8
byteVal :: [Bool] -> Word8
byteVal = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\ Word8
e Bool
b -> (if Bool
b then Word8
e else Word8
0)) ([Word8
2 forall a b. (Num a, Integral b) => a -> b -> a
^ Int
n | Int
n <- [Int
7 :: Int, Int
6 .. Int
0]])

-- |Split a list in groups of 8 elements or less
bytes :: [t] -> [[t]]
bytes :: forall t. [t] -> [[t]]
bytes [] = []
bytes [t]
l  = let ([t]
w, [t]
r) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
8 [t]
l in [t]
w forall a. a -> [a] -> [a]
: forall t. [t] -> [[t]]
bytes [t]
r

{- |
>>> prettyShow $ asBits (256+3::Word16)
"00000001 00000011"
-}
instance Pretty Bits where
    pPrint :: Bits -> Doc
pPrint = [Doc] -> Doc
hsep forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *). Foldable t => t Bool -> Doc
prettyBits forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t. [t] -> [[t]]
bytes forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Unbox a => Vector a -> [a]
V.toList

prettyBits :: Foldable t => t Bool -> Doc
prettyBits :: forall (t :: * -> *). Foldable t => t Bool -> Doc
prettyBits t Bool
l =
    [Char] -> Doc
text forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> [a] -> [a]
take (forall (t :: * -> *) a. Foldable t => t a -> Int
length t Bool
l) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\Bool
b -> if Bool
b then [Char]
"1" else [Char]
"0") forall a b. (a -> b) -> a -> b
$ t Bool
l