module Rattletrap.BitGet where

import qualified Control.Exception as Exception
import qualified Control.Monad as Monad
import qualified Data.Bits as Bits
import qualified Data.ByteString as ByteString
import qualified Data.Functor.Identity as Identity
import qualified Rattletrap.BitString as BitString
import qualified Rattletrap.ByteGet as ByteGet
import qualified Rattletrap.Exception.NotEnoughInput as NotEnoughInput
import qualified Rattletrap.Get as Get

type BitGet = Get.Get BitString.BitString Identity.Identity

toByteGet :: BitGet a -> ByteGet.ByteGet a
toByteGet :: forall a. BitGet a -> ByteGet a
toByteGet BitGet a
g = do
  ByteString
s1 <- forall (m :: * -> *) s. Applicative m => Get s m s
Get.get
  case forall a. Identity a -> a
Identity.runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (m :: * -> *) a.
Get s m a -> s -> m (Either ([String], SomeException) (s, a))
Get.run BitGet a
g forall a b. (a -> b) -> a -> b
$ ByteString -> BitString
BitString.fromByteString ByteString
s1 of
    Left ([String]
ls, SomeException
e) -> forall (m :: * -> *) s a.
Functor m =>
[String] -> Get s m a -> Get s m a
Get.labels [String]
ls forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => e -> ByteGet a
ByteGet.throw SomeException
e
    Right (BitString
s2, a
x) -> do
      forall (m :: * -> *) s. Applicative m => s -> Get s m ()
Get.put forall a b. (a -> b) -> a -> b
$ BitString -> ByteString
BitString.byteString BitString
s2
      forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

fromByteGet :: ByteGet.ByteGet a -> Int -> BitGet a
fromByteGet :: forall a. ByteGet a -> Int -> BitGet a
fromByteGet ByteGet a
f Int
n = do
  ByteString
x <- Int -> BitGet ByteString
byteString Int
n
  forall (m :: * -> *) s a t. Monad m => Get s m a -> s -> Get t m a
Get.embed ByteGet a
f ByteString
x

bits :: (Bits.Bits a) => Int -> BitGet a
bits :: forall a. Bits a => Int -> BitGet a
bits Int
n = do
  let f :: (Bits.Bits a) => Bool -> a -> a
      f :: forall a. Bits a => Bool -> a -> a
f Bool
bit a
x = let y :: a
y = forall a. Bits a => a -> Int -> a
Bits.shiftL a
x Int
1 in if Bool
bit then forall a. Bits a => a -> Int -> a
Bits.setBit a
y Int
0 else a
y
  [Bool]
xs <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
Monad.replicateM Int
n BitGet Bool
bool
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall a. Bits a => Bool -> a -> a
f forall a. Bits a => a
Bits.zeroBits [Bool]
xs

bool :: BitGet Bool
bool :: BitGet Bool
bool = do
  BitString
s1 <- forall (m :: * -> *) s. Applicative m => Get s m s
Get.get
  case BitString -> Maybe (Bool, BitString)
BitString.pop BitString
s1 of
    Maybe (Bool, BitString)
Nothing -> forall e a. Exception e => e -> BitGet a
throw NotEnoughInput
NotEnoughInput.NotEnoughInput
    Just (Bool
x, BitString
s2) -> do
      forall (m :: * -> *) s. Applicative m => s -> Get s m ()
Get.put BitString
s2
      forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
x

byteString :: Int -> BitGet ByteString.ByteString
byteString :: Int -> BitGet ByteString
byteString Int
n = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Word8] -> ByteString
ByteString.pack forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
Monad.replicateM Int
n forall a b. (a -> b) -> a -> b
$ forall a. Bits a => Int -> BitGet a
bits Int
8

throw :: (Exception.Exception e) => e -> BitGet a
throw :: forall e a. Exception e => e -> BitGet a
throw = forall e (m :: * -> *) s a.
(Exception e, Applicative m) =>
e -> Get s m a
Get.throw

label :: String -> BitGet a -> BitGet a
label :: forall a. String -> BitGet a -> BitGet a
label = forall (m :: * -> *) s a.
Functor m =>
String -> Get s m a -> Get s m a
Get.label