{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Safe #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- Module      :  Text.Megaparsec.Byte.Binary
-- Copyright   :  © 2021–present Megaparsec contributors
-- License     :  FreeBSD
--
-- Maintainer  :  Mark Karpov <markkarpov92@gmail.com>
-- Stability   :  experimental
-- Portability :  portable
--
-- Binary-format number parsers.
--
-- @since 9.2.0
module Text.Megaparsec.Byte.Binary
  ( -- * Generic parsers
    BinaryChunk (..),
    anyLE,
    anyBE,

    -- * Parsing unsigned values
    word8,
    word16le,
    word16be,
    word32le,
    word32be,
    word64le,
    word64be,

    -- * Parsing signed values
    int8,
    int16le,
    int16be,
    int32le,
    int32be,
    int64le,
    int64be,
  )
where

import Data.Bits
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.Int
import Data.Word
import Text.Megaparsec

-- | Data types that can be converted to little- or big- endian numbers.
class BinaryChunk chunk where
  convertChunkBE :: (Bits a, Num a) => chunk -> a
  convertChunkLE :: (Bits a, Num a) => chunk -> a

instance BinaryChunk B.ByteString where
  convertChunkBE :: ByteString -> a
convertChunkBE = (a -> Word8 -> a) -> a -> ByteString -> a
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
B.foldl' a -> Word8 -> a
forall a a. (Bits a, Integral a, Num a) => a -> a -> a
go a
0
    where
      go :: a -> a -> a
go a
acc a
byte = (a
acc a -> Int -> a
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
8) a -> a -> a
forall a. Bits a => a -> a -> a
.|. a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
byte
  convertChunkLE :: ByteString -> a
convertChunkLE = (a -> Word8 -> a) -> a -> ByteString -> a
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
B.foldl' a -> Word8 -> a
forall a a. (Bits a, Integral a, Num a) => a -> a -> a
go a
0
    where
      go :: a -> a -> a
go a
acc a
byte = (a
acc a -> a -> a
forall a. Bits a => a -> a -> a
.|. a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
byte) a -> Int -> a
forall a. Bits a => a -> Int -> a
`rotateR` Int
8

instance BinaryChunk BL.ByteString where
  convertChunkBE :: ByteString -> a
convertChunkBE = (a -> Word8 -> a) -> a -> ByteString -> a
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
BL.foldl' a -> Word8 -> a
forall a a. (Bits a, Integral a, Num a) => a -> a -> a
go a
0
    where
      go :: a -> a -> a
go a
acc a
byte = (a
acc a -> Int -> a
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
8) a -> a -> a
forall a. Bits a => a -> a -> a
.|. a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
byte
  convertChunkLE :: ByteString -> a
convertChunkLE = (a -> Word8 -> a) -> a -> ByteString -> a
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
BL.foldl' a -> Word8 -> a
forall a a. (Bits a, Integral a, Num a) => a -> a -> a
go a
0
    where
      go :: a -> a -> a
go a
acc a
byte = (a
acc a -> a -> a
forall a. Bits a => a -> a -> a
.|. a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
byte) a -> Int -> a
forall a. Bits a => a -> Int -> a
`rotateR` Int
8

----------------------------------------------------------------------------
-- Generic parsers

-- | Parse a little-endian number.
--
-- You may wish to call this with a visible type application:
--
-- > number <- anyLE (Just "little-endian 32 bit word") @Word32
anyLE ::
  forall a e s m.
  (MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
  -- | Label, if any
  Maybe String ->
  m a
anyLE :: Maybe String -> m a
anyLE Maybe String
mlabel = Tokens s -> a
forall chunk a. (BinaryChunk chunk, Bits a, Num a) => chunk -> a
convertChunkLE (Tokens s -> a) -> m (Tokens s) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe String -> Int -> m (Tokens s)
forall e s (m :: * -> *).
MonadParsec e s m =>
Maybe String -> Int -> m (Tokens s)
takeP Maybe String
mlabel (FiniteBits a => Int
forall a. FiniteBits a => Int
finiteByteSize @a)
{-# INLINE anyLE #-}

-- | Parse a big-endian number.
--
-- You may wish to call this with a visible type application:
--
-- > number <- anyBE (Just "big-endian 32 bit word") @Word32
anyBE ::
  forall a e s m.
  (MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
  -- | Label, if any
  Maybe String ->
  m a
anyBE :: Maybe String -> m a
anyBE Maybe String
mlabel = Tokens s -> a
forall chunk a. (BinaryChunk chunk, Bits a, Num a) => chunk -> a
convertChunkBE (Tokens s -> a) -> m (Tokens s) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe String -> Int -> m (Tokens s)
forall e s (m :: * -> *).
MonadParsec e s m =>
Maybe String -> Int -> m (Tokens s)
takeP Maybe String
mlabel (FiniteBits a => Int
forall a. FiniteBits a => Int
finiteByteSize @a)
{-# INLINE anyBE #-}

--------------------------------------------------------------------------------
-- Parsing unsigned values

-- | Parse a 'Word8'.
word8 :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Word8
word8 :: m Word8
word8 = Maybe String -> m Word8
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyBE (String -> Maybe String
forall a. a -> Maybe a
Just String
"8 bit word")
{-# INLINE word8 #-}

-- | Parse a little-endian 'Word16'.
word16le :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Word16
word16le :: m Word16
word16le = Maybe String -> m Word16
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyLE (String -> Maybe String
forall a. a -> Maybe a
Just String
"little-endian 16 bit word")
{-# INLINE word16le #-}

-- | Parse a big-endian 'Word16'.
word16be :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Word16
word16be :: m Word16
word16be = Maybe String -> m Word16
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyBE (String -> Maybe String
forall a. a -> Maybe a
Just String
"big-endian 16 bit word")
{-# INLINE word16be #-}

-- | Parse a little-endian 'Word32'.
word32le :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Word32
word32le :: m Word32
word32le = Maybe String -> m Word32
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyLE (String -> Maybe String
forall a. a -> Maybe a
Just String
"little-endian 32 bit word")
{-# INLINE word32le #-}

-- | Parse a big-endian 'Word32'.
word32be :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Word32
word32be :: m Word32
word32be = Maybe String -> m Word32
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyBE (String -> Maybe String
forall a. a -> Maybe a
Just String
"big-endian 32 bit word")
{-# INLINE word32be #-}

-- | Parse a little-endian 'Word64'.
word64le :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Word64
word64le :: m Word64
word64le = Maybe String -> m Word64
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyLE (String -> Maybe String
forall a. a -> Maybe a
Just String
"little-endian 64 word")
{-# INLINE word64le #-}

-- | Parse a big-endian 'Word64'.
word64be :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Word64
word64be :: m Word64
word64be = Maybe String -> m Word64
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyBE (String -> Maybe String
forall a. a -> Maybe a
Just String
"big-endian 64 word")
{-# INLINE word64be #-}

----------------------------------------------------------------------------
-- Parsing signed values

-- | Parse a 'Int8'.
int8 :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Int8
int8 :: m Int8
int8 = Maybe String -> m Int8
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyBE (String -> Maybe String
forall a. a -> Maybe a
Just String
"8 bit int")
{-# INLINE int8 #-}

-- | Parse a little-endian 'Int16'.
int16le :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Int16
int16le :: m Int16
int16le = Maybe String -> m Int16
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyLE (String -> Maybe String
forall a. a -> Maybe a
Just String
"little-endian 16 bit int")
{-# INLINE int16le #-}

-- | Parse a big-endian 'Int16'.
int16be :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Int16
int16be :: m Int16
int16be = Maybe String -> m Int16
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyBE (String -> Maybe String
forall a. a -> Maybe a
Just String
"big-endian 16 bit int")
{-# INLINE int16be #-}

-- | Parse a little-endian 'Int32'.
int32le :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Int32
int32le :: m Int32
int32le = Maybe String -> m Int32
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyLE (String -> Maybe String
forall a. a -> Maybe a
Just String
"little-endian 32 bit int")
{-# INLINE int32le #-}

-- | Parse a big-endian 'Int32'.
int32be :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Int32
int32be :: m Int32
int32be = Maybe String -> m Int32
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyBE (String -> Maybe String
forall a. a -> Maybe a
Just String
"big-endian 32 bit int")
{-# INLINE int32be #-}

-- | Parse a little-endian 'Int64'.
int64le :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Int64
int64le :: m Int64
int64le = Maybe String -> m Int64
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyLE (String -> Maybe String
forall a. a -> Maybe a
Just String
"little-endian 64 int")
{-# INLINE int64le #-}

-- | Parse a big-endian 'Int64'.
int64be :: (MonadParsec e s m, BinaryChunk (Tokens s)) => m Int64
int64be :: m Int64
int64be = Maybe String -> m Int64
forall a e s (m :: * -> *).
(MonadParsec e s m, FiniteBits a, Num a, BinaryChunk (Tokens s)) =>
Maybe String -> m a
anyBE (String -> Maybe String
forall a. a -> Maybe a
Just String
"big-endian 64 int")
{-# INLINE int64be #-}

--------------------------------------------------------------------------------
-- Helpers

-- | Return the number of bytes in the argument.
--
-- Performs ceiling division, so byte-unaligned types (bitsize not a
-- multiple of 8) should work, but further usage is not tested.
finiteByteSize :: forall a. FiniteBits a => Int
finiteByteSize :: Int
finiteByteSize = a -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize @a a
forall a. HasCallStack => a
undefined Int -> Int -> Int
forall a. Integral a => a -> a -> a
`ceilDiv` Int
8
  where
    ceilDiv :: a -> a -> a
ceilDiv a
x a
y = (a
x a -> a -> a
forall a. Num a => a -> a -> a
+ a
y a -> a -> a
forall a. Num a => a -> a -> a
- a
1) a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
y
{-# INLINE finiteByteSize #-}