{-# LANGUAGE AllowAmbiguousTypes, UndecidableInstances #-}

{- | Efficient type-level bytestring parsing via chunking.

See 'Bytezap.Struct.TypeLits.Bytes' for an explanation on the chunking design.

On mismatch, the index of the failing byte and its value are returned. (This is
over-engineered to be extremely efficient.)

Type classes take a 'Natural' for tracking the current index in the type-level
bytestring. We do this on the type level for performance. Use @\@0@ when
calling.

The parsers take an error wrapper function to enable wrapping the error into any
parser with confidence that it won't do extra allocations/wrapping.

The parsers here either return the unit '()' or a pretty error. No 'Fail#'.

TODO check generated Core, assembly
-}

module Bytezap.Parser.Struct.TypeLits.Bytes
  ( ParseReifyBytesW64(parseReifyBytesW64)
  , ParseReifyBytesW32(parseReifyBytesW32)
  , ParseReifyBytesW16(parseReifyBytesW16)
  , ParseReifyBytesW8(parseReifyBytesW8)
  ) where

import Bytezap.Parser.Struct
import Data.Word ( Word8 )
import GHC.TypeNats ( Natural, type (+), KnownNat )
import Data.Type.Byte ( ReifyW8, reifyW64, reifyW32, reifyW16, reifyW8 )
import GHC.Exts ( (+#), Int(I#), Int#, Addr# )
import Util.TypeNats ( natValInt )
import Raehik.Compat.Data.Primitive.Types ( indexWord8OffAddrAs# )
import Data.Bits

-- | Parse a type-level bytestring, largest grouping 'Word64'.
class ParseReifyBytesW64 (idx :: Natural) (bs :: [Natural]) where
    parseReifyBytesW64 :: (Int -> Word8 -> e) -> ParserT st e ()

-- | Enough bytes to make a 'Word64'.
instance {-# OVERLAPPING #-}
  ( ReifyW8 b0
  , ReifyW8 b1
  , ReifyW8 b2
  , ReifyW8 b3
  , ReifyW8 b4
  , ReifyW8 b5
  , ReifyW8 b6
  , ReifyW8 b7
  , KnownNat idx
  , ParseReifyBytesW64 (idx+8) bs
  ) => ParseReifyBytesW64 idx (b0 ': b1 ': b2 ': b3 ': b4 ': b5 ': b6 ': b7 ': bs) where
    {-# INLINE parseReifyBytesW64 #-}
    parseReifyBytesW64 :: forall e (st :: ZeroBitType).
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW64 = forall (idx :: Natural) (len :: Natural) a e (st :: ZeroBitType).
(KnownNat idx, KnownNat len, Integral a, FiniteBits a) =>
a
-> (Addr# -> Int# -> a)
-> ((Int -> Word8 -> e) -> ParserT st e ())
-> (Int -> Word8 -> e)
-> ParserT st e ()
parseReifyBytesHelper @idx @8
        Word64
wExpect Addr# -> Int# -> Word64
forall a. Prim' a => Addr# -> Int# -> a
indexWord8OffAddrAs# (forall (idx :: Natural) (bs :: [Natural]) e (st :: ZeroBitType).
ParseReifyBytesW64 idx bs =>
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW64 @(idx+8) @bs)
      where
        wExpect :: Word64
wExpect = forall (n1 :: Natural) (n2 :: Natural) (n3 :: Natural)
       (n4 :: Natural) (n5 :: Natural) (n6 :: Natural) (n7 :: Natural)
       (n8 :: Natural).
(ReifyW8 n1, ReifyW8 n2, ReifyW8 n3, ReifyW8 n4, ReifyW8 n5,
 ReifyW8 n6, ReifyW8 n7, ReifyW8 n8) =>
Word64
reifyW64 @b0 @b1 @b2 @b3 @b4 @b5 @b6 @b7

-- | Try to group 'Word32's next.
instance ParseReifyBytesW32 idx bs => ParseReifyBytesW64 idx bs where
    {-# INLINE parseReifyBytesW64 #-}
    parseReifyBytesW64 :: forall e (st :: ZeroBitType).
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW64 = forall (idx :: Natural) (bs :: [Natural]) e (st :: ZeroBitType).
ParseReifyBytesW32 idx bs =>
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW32 @idx @bs

-- | Parse a type-level bytestring, largest grouping 'Word32'.
class ParseReifyBytesW32 (idx :: Natural) (bs :: [Natural]) where
    parseReifyBytesW32 :: (Int -> Word8 -> e) -> ParserT st e ()

-- | Enough bytes to make a 'Word32'.
instance {-# OVERLAPPING #-}
  ( ReifyW8 b0
  , ReifyW8 b1
  , ReifyW8 b2
  , ReifyW8 b3
  , KnownNat idx
  , ParseReifyBytesW32 (idx+4) bs
  ) => ParseReifyBytesW32 idx (b0 ': b1 ': b2 ': b3 ': bs) where
    {-# INLINE parseReifyBytesW32 #-}
    parseReifyBytesW32 :: forall e (st :: ZeroBitType).
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW32 = forall (idx :: Natural) (len :: Natural) a e (st :: ZeroBitType).
(KnownNat idx, KnownNat len, Integral a, FiniteBits a) =>
a
-> (Addr# -> Int# -> a)
-> ((Int -> Word8 -> e) -> ParserT st e ())
-> (Int -> Word8 -> e)
-> ParserT st e ()
parseReifyBytesHelper @idx @4
        Word32
wExpect Addr# -> Int# -> Word32
forall a. Prim' a => Addr# -> Int# -> a
indexWord8OffAddrAs# (forall (idx :: Natural) (bs :: [Natural]) e (st :: ZeroBitType).
ParseReifyBytesW32 idx bs =>
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW32 @(idx+4) @bs)
      where
        wExpect :: Word32
wExpect = forall (n1 :: Natural) (n2 :: Natural) (n3 :: Natural)
       (n4 :: Natural).
(ReifyW8 n1, ReifyW8 n2, ReifyW8 n3, ReifyW8 n4) =>
Word32
reifyW32 @b0 @b1 @b2 @b3

-- | Try to group 'Word16's next.
instance ParseReifyBytesW16 idx bs => ParseReifyBytesW32 idx bs where
    {-# INLINE parseReifyBytesW32 #-}
    parseReifyBytesW32 :: forall e (st :: ZeroBitType).
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW32 = forall (idx :: Natural) (bs :: [Natural]) e (st :: ZeroBitType).
ParseReifyBytesW16 idx bs =>
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW16 @idx @bs

-- | Parse a type-level bytestring, largest grouping 'Word16'.
class ParseReifyBytesW16 (idx :: Natural) (bs :: [Natural]) where
    parseReifyBytesW16 :: (Int -> Word8 -> e) -> ParserT st e ()

-- | Enough bytes to make a 'Word16'.
instance {-# OVERLAPPING #-}
  ( ReifyW8 b0
  , ReifyW8 b1
  , KnownNat idx
  , ParseReifyBytesW16 (idx+2) bs
  ) => ParseReifyBytesW16 idx (b0 ': b1 ': bs) where
    {-# INLINE parseReifyBytesW16 #-}
    parseReifyBytesW16 :: forall e (st :: ZeroBitType).
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW16 = forall (idx :: Natural) (len :: Natural) a e (st :: ZeroBitType).
(KnownNat idx, KnownNat len, Integral a, FiniteBits a) =>
a
-> (Addr# -> Int# -> a)
-> ((Int -> Word8 -> e) -> ParserT st e ())
-> (Int -> Word8 -> e)
-> ParserT st e ()
parseReifyBytesHelper @idx @2
        Word16
wExpect Addr# -> Int# -> Word16
forall a. Prim' a => Addr# -> Int# -> a
indexWord8OffAddrAs# (forall (idx :: Natural) (bs :: [Natural]) e (st :: ZeroBitType).
ParseReifyBytesW16 idx bs =>
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW16 @(idx+2) @bs)
      where
        wExpect :: Word16
wExpect = forall (n1 :: Natural) (n2 :: Natural).
(ReifyW8 n1, ReifyW8 n2) =>
Word16
reifyW16 @b0 @b1

-- | Parse byte-by-byte next.
instance ParseReifyBytesW8 idx bs => ParseReifyBytesW16 idx bs where
    {-# INLINE parseReifyBytesW16 #-}
    parseReifyBytesW16 :: forall e (st :: ZeroBitType).
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW16 = forall (idx :: Natural) (bs :: [Natural]) e (st :: ZeroBitType).
ParseReifyBytesW8 idx bs =>
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW8 @idx @bs

-- | Serialize a type-level bytestring, byte-by-byte.
class ParseReifyBytesW8 (idx :: Natural) (bs :: [Natural]) where
    parseReifyBytesW8 :: (Int -> Word8 -> e) -> ParserT st e ()

-- | Parse the next byte.
instance
  ( ReifyW8 b0
  , KnownNat idx
  , ParseReifyBytesW8 (idx+1) bs
  ) => ParseReifyBytesW8 idx (b0 ': bs) where
    {-# INLINE parseReifyBytesW8 #-}
    parseReifyBytesW8 :: forall e (st :: ZeroBitType).
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW8 Int -> Word8 -> e
f = ParserT# st e () -> ParserT st e ()
forall (st :: ZeroBitType) e a. ParserT# st e a -> ParserT st e a
ParserT (ParserT# st e () -> ParserT st e ())
-> ParserT# st e () -> ParserT st e ()
forall a b. (a -> b) -> a -> b
$ \ForeignPtrContents
fpc Addr#
base# Int#
os# st
st ->
        let bExpect :: Word8
bExpect = forall (n :: Natural). ReifyW8 n => Word8
reifyW8 @b0
            bActual :: Word8
bActual = Addr# -> Int# -> Word8
forall a. Prim' a => Addr# -> Int# -> a
indexWord8OffAddrAs# Addr#
base# Int#
os#
            idx :: Int
idx     = forall (n :: Natural). KnownNat n => Int
natValInt @idx
        in  if   Word8
bExpect Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
bActual
            then ParserT st e () -> ParserT# st e ()
forall (st :: ZeroBitType) e a. ParserT st e a -> ParserT# st e a
runParserT# (forall (idx :: Natural) (bs :: [Natural]) e (st :: ZeroBitType).
ParseReifyBytesW8 idx bs =>
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW8 @(idx+1) @bs Int -> Word8 -> e
f) ForeignPtrContents
fpc Addr#
base# (Int#
os# Int# -> Int# -> Int#
+# Int#
1#) st
st
            else st -> e -> Res# st e ()
forall (st :: ZeroBitType) e a. st -> e -> Res# st e a
Err# st
st (Int -> Word8 -> e
f Int
idx Word8
bActual)

-- | End of the line.
instance ParseReifyBytesW8 idx '[] where
    {-# INLINE parseReifyBytesW8 #-}
    parseReifyBytesW8 :: forall e (st :: ZeroBitType).
(Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesW8 Int -> Word8 -> e
_f = ParserT# st e () -> ParserT st e ()
forall (st :: ZeroBitType) e a. ParserT# st e a -> ParserT st e a
ParserT (ParserT# st e () -> ParserT st e ())
-> ParserT# st e () -> ParserT st e ()
forall a b. (a -> b) -> a -> b
$ \ForeignPtrContents
_fpc Addr#
_base# Int#
_os# st
st -> st -> () -> Res# st e ()
forall (st :: ZeroBitType) a e. st -> a -> Res# st e a
OK# st
st ()

parseReifyBytesHelper
    :: forall (idx :: Natural) (len :: Natural) a e st
    .  (KnownNat idx, KnownNat len, Integral a, FiniteBits a)
    => a -> (Addr# -> Int# -> a)
    -> ((Int -> Word8 -> e) -> ParserT st e ())
    -> (Int -> Word8 -> e) -> ParserT st e ()
parseReifyBytesHelper :: forall (idx :: Natural) (len :: Natural) a e (st :: ZeroBitType).
(KnownNat idx, KnownNat len, Integral a, FiniteBits a) =>
a
-> (Addr# -> Int# -> a)
-> ((Int -> Word8 -> e) -> ParserT st e ())
-> (Int -> Word8 -> e)
-> ParserT st e ()
parseReifyBytesHelper a
a Addr# -> Int# -> a
indexWord8OffAddrAsA# (Int -> Word8 -> e) -> ParserT st e ()
pCont Int -> Word8 -> e
f = (Int -> a -> e)
-> Int#
-> a
-> (Addr# -> Int# -> a)
-> ParserT st e ()
-> ParserT st e ()
forall a e (st :: ZeroBitType) r.
(Num a, FiniteBits a) =>
(Int -> a -> e)
-> Int#
-> a
-> (Addr# -> Int# -> a)
-> ParserT st e r
-> ParserT st e r
withLitErr
    (\Int
idx a
b -> Int -> Word8 -> e
f (forall (n :: Natural). KnownNat n => Int
natValInt @idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
idx) (a -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
b))
    Int#
len# a
a Addr# -> Int# -> a
indexWord8OffAddrAsA# ((Int -> Word8 -> e) -> ParserT st e ()
pCont Int -> Word8 -> e
f)
  where
    !(I# Int#
len#) = forall (n :: Natural). KnownNat n => Int
natValInt @len