{-# LANGUAGE UndecidableInstances #-} -- required for type-level stuff
{-# LANGUAGE OverloadedStrings #-} -- required for refined errors

module Binrep.Type.Prefix.Size where

import Binrep.Type.Prefix.Internal
import Binrep.Type.Thin
import Binrep
import FlatParse.Basic qualified as FP

import GHC.TypeNats
import Util.TypeNats ( natValInt )
import Data.ByteString qualified as B

import Rerefined.Predicate.Common
import Rerefined.Refine
import TypeLevelShow.Utils

import Data.Kind ( type Type )

data SizePrefix (pfx :: Type)

instance Predicate (SizePrefix pfx) where
    type PredicateName d (SizePrefix pfx) = ShowParen (d > 9)
        ("SizePrefix " ++ LenNatName pfx)

type SizePrefixed pfx = Refined (SizePrefix pfx)

instance
  ( KnownPredicateName (SizePrefix pfx), KnownNat (LenNatMax pfx), BLen a
  ) => Refine (SizePrefix pfx) a where
    validate :: Proxy# (SizePrefix pfx) -> a -> Maybe RefineFailure
validate Proxy# (SizePrefix pfx)
p a
a = Proxy# (SizePrefix pfx) -> Bool -> Builder -> Maybe RefineFailure
forall {k} (p :: k).
(Predicate p, KnownPredicateName p) =>
Proxy# p -> Bool -> Builder -> Maybe RefineFailure
validateBool Proxy# (SizePrefix pfx)
p (a -> Int
forall a. BLen a => a -> Int
blen a
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= forall (n :: Nat). KnownNat n => Int
natValInt @(LenNatMax pfx)) (Builder -> Maybe RefineFailure) -> Builder -> Maybe RefineFailure
forall a b. (a -> b) -> a -> b
$
        Builder
"thing too big for length prefix type"

-- TODO no idea if this is sensible
instance IsCBLen (SizePrefixed pfx a) where
    type CBLen (SizePrefixed pfx a) = CBLen pfx + CBLen a

instance (LenNat pfx, BLen a, BLen pfx)
  => BLen (SizePrefixed pfx a) where
    blen :: SizePrefixed pfx a -> Int
blen SizePrefixed pfx a
ra = pfx -> Int
forall a. BLen a => a -> Int
blen (forall a. LenNat a => Int -> a
lenToNat @pfx (a -> Int
forall a. BLen a => a -> Int
blen a
a)) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int
forall a. BLen a => a -> Int
blen a
a
      where a :: a
a = SizePrefixed pfx a -> a
forall {k} (p :: k) a. Refined p a -> a
unrefine SizePrefixed pfx a
ra

instance (LenNat pfx, BLen a, Put pfx, Put a)
  => Put (SizePrefixed pfx a) where
    put :: SizePrefixed pfx a -> Putter
put SizePrefixed pfx a
ra = pfx -> Putter
forall a. Put a => a -> Putter
put (forall a. LenNat a => Int -> a
lenToNat @pfx (a -> Int
forall a. BLen a => a -> Int
blen a
a)) Putter -> Putter -> Putter
forall a. Semigroup a => a -> a -> a
<> a -> Putter
forall a. Put a => a -> Putter
put a
a
      where a :: a
a = SizePrefixed pfx a -> a
forall {k} (p :: k) a. Refined p a -> a
unrefine SizePrefixed pfx a
ra

class GetSize a where getSize :: Int -> Getter a
instance GetSize       B.ByteString  where getSize :: Int -> Getter ByteString
getSize = (ByteString -> ByteString)
-> Getter ByteString -> Getter ByteString
forall a b.
(a -> b)
-> ParserT PureMode (ParseError Pos Builder) a
-> ParserT PureMode (ParseError Pos Builder) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> ByteString
B.copy (Getter ByteString -> Getter ByteString)
-> (Int -> Getter ByteString) -> Int -> Getter ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Getter ByteString
forall (st :: ZeroBitType) e. Int -> ParserT st e ByteString
FP.take
instance GetSize (Thin B.ByteString) where getSize :: Int -> Getter (Thin ByteString)
getSize = (ByteString -> Thin ByteString)
-> Getter ByteString -> Getter (Thin ByteString)
forall a b.
(a -> b)
-> ParserT PureMode (ParseError Pos Builder) a
-> ParserT PureMode (ParseError Pos Builder) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> Thin ByteString
forall a. a -> Thin a
Thin (Getter ByteString -> Getter (Thin ByteString))
-> (Int -> Getter ByteString) -> Int -> Getter (Thin ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Getter ByteString
forall (st :: ZeroBitType) e. Int -> ParserT st e ByteString
FP.take

instance (LenNat pfx, GetSize a, Get pfx)
  => Get (SizePrefixed pfx a) where
    get :: Getter (SizePrefixed pfx a)
get = do
        pfx
pfx <- forall a. Get a => Getter a
get @pfx
        a
a <- Int -> Getter a
forall a. GetSize a => Int -> Getter a
getSize (pfx -> Int
forall a. LenNat a => a -> Int
natToLen pfx
pfx)
        SizePrefixed pfx a -> Getter (SizePrefixed pfx a)
forall a. a -> ParserT PureMode (ParseError Pos Builder) a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (SizePrefixed pfx a -> Getter (SizePrefixed pfx a))
-> SizePrefixed pfx a -> Getter (SizePrefixed pfx a)
forall a b. (a -> b) -> a -> b
$ a -> SizePrefixed pfx a
forall {k} a (p :: k). a -> Refined p a
unsafeRefine a
a