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

module Binrep.Type.Prefix.Count where

import Binrep.Type.Prefix.Internal
import Binrep
import Control.Monad.Combinators qualified as Monad

import GHC.TypeNats
import Util.TypeNats ( natValInt )

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

import Data.Kind ( type Type )

import Data.Foldable qualified as Foldable

-- TODO put monofoldable in here, instead of that useless @(f a)@ stuff

data CountPrefix (pfx :: Type)

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

instance
  ( KnownPredicateName (CountPrefix pfx), KnownNat (LenNatMax pfx), Foldable f
  ) => Refine1 (CountPrefix pfx) f where
    validate1 :: forall a. Proxy# (CountPrefix pfx) -> f a -> Maybe RefineFailure
validate1 Proxy# (CountPrefix pfx)
p f a
fa =
        Proxy# (CountPrefix pfx) -> Bool -> Builder -> Maybe RefineFailure
forall {k} (p :: k).
(Predicate p, KnownPredicateName p) =>
Proxy# p -> Bool -> Builder -> Maybe RefineFailure
validateBool Proxy# (CountPrefix pfx)
p (f a -> Int
forall a. f a -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
Foldable.length f a
fa 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
"TODO too large for count prefix"

instance
  ( KnownPredicateName (CountPrefix pfx), KnownNat (LenNatMax pfx), Foldable f
  ) => Refine (CountPrefix pfx) (f a) where
    validate :: Proxy# (CountPrefix pfx) -> f a -> Maybe RefineFailure
validate = Proxy# (CountPrefix pfx) -> f a -> Maybe RefineFailure
forall a. Proxy# (CountPrefix pfx) -> f a -> Maybe RefineFailure
forall {k} {k1} (p :: k) (f :: k1 -> Type) (a :: k1).
Refine1 p f =>
Proxy# p -> f a -> Maybe RefineFailure
validate1

type CountPrefixed pfx = Refined1 (CountPrefix pfx)

-- | We can know byte length at compile time /if/ we know it for the prefix and
--   the list-like.
--
-- This is extremely unlikely, because then what counting are we even
-- performing for the list-like? But it's a valid instance.
instance IsCBLen (CountPrefixed pfx f a) where
    type CBLen (CountPrefixed pfx f a) = CBLen pfx + CBLen (f a)

-- | The byte length of a count-prefixed type is the length of the prefix type
--   (holding the length of the type) plus the length of the type.
--
-- Bit confusing. How to explain this? TODO
instance (LenNat pfx, Foldable f, BLen pfx, BLen (f a))
  => BLen (CountPrefixed pfx f a) where
    blen :: CountPrefixed pfx f a -> Int
blen CountPrefixed pfx f a
rfa = pfx -> Int
forall a. BLen a => a -> Int
blen (forall a. LenNat a => Int -> a
lenToNat @pfx (f a -> Int
forall a. f a -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
Foldable.length f a
fa)) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ f a -> Int
forall a. BLen a => a -> Int
blen f a
fa
      where fa :: f a
fa = CountPrefixed pfx f a -> f a
forall {k1} {k2} (p :: k1) (f :: k2 -> Type) (a :: k2).
Refined1 p f a -> f a
unrefine1 CountPrefixed pfx f a
rfa

instance (LenNat pfx, Foldable f, Put pfx, Put (f a))
  => Put (CountPrefixed pfx f a) where
    put :: CountPrefixed pfx f a -> Putter
put CountPrefixed pfx f a
rfa = pfx -> Putter
forall a. Put a => a -> Putter
put (forall a. LenNat a => Int -> a
lenToNat @pfx (f a -> Int
forall a. f a -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
Foldable.length f a
fa)) Putter -> Putter -> Putter
forall a. Semigroup a => a -> a -> a
<> f a -> Putter
forall a. Put a => a -> Putter
put f a
fa
      where fa :: f a
fa = CountPrefixed pfx f a -> f a
forall {k1} {k2} (p :: k1) (f :: k2 -> Type) (a :: k2).
Refined1 p f a -> f a
unrefine1 CountPrefixed pfx f a
rfa

class GetCount f where getCount :: Get a => Int -> Getter (f a)
instance GetCount [] where getCount :: forall a. Get a => Int -> Getter [a]
getCount Int
n = Int
-> ParserT PureMode (ParseError Pos Builder) a
-> ParserT PureMode (ParseError Pos Builder) [a]
forall (m :: Type -> Type) a. Monad m => Int -> m a -> m [a]
Monad.count Int
n ParserT PureMode (ParseError Pos Builder) a
forall a. Get a => Getter a
get

instance (LenNat pfx, GetCount f, Get pfx, Get a)
  => Get (CountPrefixed pfx f a) where
    get :: Getter (CountPrefixed pfx f a)
get = do
        pfx
pfx <- forall a. Get a => Getter a
get @pfx
        f a
fa <- Int -> Getter (f a)
forall a. Get a => Int -> Getter (f a)
forall (f :: Type -> Type) a.
(GetCount f, Get a) =>
Int -> Getter (f a)
getCount (pfx -> Int
forall a. LenNat a => a -> Int
natToLen pfx
pfx)
        CountPrefixed pfx f a -> Getter (CountPrefixed pfx f a)
forall a. a -> ParserT PureMode (ParseError Pos Builder) a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (CountPrefixed pfx f a -> Getter (CountPrefixed pfx f a))
-> CountPrefixed pfx f a -> Getter (CountPrefixed pfx f a)
forall a b. (a -> b) -> a -> b
$ f a -> CountPrefixed pfx f a
forall {k1} {k2} (f :: k1 -> Type) (a :: k1) (p :: k2).
f a -> Refined1 p f a
unsafeRefine1 f a
fa