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

module Binrep.Type.Prefix.Size where

import Binrep.Util.Prefix
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 Refined hiding ( Weaken(..), strengthen )
import Refined.Unsafe

import Data.Typeable ( Typeable, typeRep )
import Data.Kind

data SizePrefix (pfx :: Type)
type SizePrefixed pfx = Refined (SizePrefix pfx)

instance (KnownNat (Max pfx), BLen a, Typeable pfx)
  => Predicate (SizePrefix pfx) a where
    validate :: Proxy (SizePrefix pfx) -> a -> Maybe RefineException
validate Proxy (SizePrefix pfx)
p a
a
      | 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 @(Max pfx) = Maybe RefineException
forall a. Maybe a
Nothing
      | Bool
otherwise = TypeRep -> Text -> Maybe RefineException
throwRefineOtherException (Proxy (SizePrefix pfx) -> TypeRep
forall {k} (proxy :: k -> Type) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep Proxy (SizePrefix pfx)
p) (Text -> Maybe RefineException) -> Text -> Maybe RefineException
forall a b. (a -> b) -> a -> b
$
          Text
"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 (Prefix 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. Prefix a => Int -> a
lenToPfx @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) x. Refined p x -> x
unrefine SizePrefixed pfx a
ra

instance (Prefix 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. Prefix a => Int -> a
lenToPfx @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) x. Refined p x -> x
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 E a -> ParserT PureMode E 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 E a -> ParserT PureMode E 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 (Prefix 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. Prefix a => a -> Int
pfxToLen pfx
pfx)
        SizePrefixed pfx a -> Getter (SizePrefixed pfx a)
forall a. a -> ParserT PureMode E 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} x (p :: k). x -> Refined p x
reallyUnsafeRefine a
a