{-# LANGUAGE UndecidableInstances #-} -- for 'CBLenly', 'TypeError'
{-# LANGUAGE AllowAmbiguousTypes #-} -- for 'cblen', 'natValInt'

{- | Byte length as a simple pure function, no bells or whistles.

Non-reallocating serializers like store, bytezap or ptr-poker request the
expected total byte length when serializing. Thus, they need some way to measure
byte length *before* serializing. This is that.

It should be very efficient to calculate serialized byte length for most
binrep-compatible Haskell types. If it isn't, consider whether the
representation is appropriate for binrep.
-}

module Binrep.BLen.Simple where

import Binrep.CBLen
import GHC.TypeNats
import Util.TypeNats ( natValInt )

import Binrep.Util.Class
import GHC.TypeLits ( TypeError )

import Data.Void
import Data.ByteString qualified as B
import Data.Word
import Data.Int
import Bytezap ( Write(..) )

import Data.Monoid ( Sum(..) )
import GHC.Generics
import Generic.Data.Function.FoldMap
import Generic.Data.Rep.Assert
import Generic.Data.Function.Common

class BLen a where blen :: a -> Int

-- newtype sum monoid for generic foldMap
newtype BLen' a = BLen' { forall a. BLen' a -> a
getBLen' :: a }
    deriving (NonEmpty (BLen' a) -> BLen' a
BLen' a -> BLen' a -> BLen' a
(BLen' a -> BLen' a -> BLen' a)
-> (NonEmpty (BLen' a) -> BLen' a)
-> (forall b. Integral b => b -> BLen' a -> BLen' a)
-> Semigroup (BLen' a)
forall b. Integral b => b -> BLen' a -> BLen' a
forall a. Num a => NonEmpty (BLen' a) -> BLen' a
forall a. Num a => BLen' a -> BLen' a -> BLen' a
forall a b. (Num a, Integral b) => b -> BLen' a -> BLen' a
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
$c<> :: forall a. Num a => BLen' a -> BLen' a -> BLen' a
<> :: BLen' a -> BLen' a -> BLen' a
$csconcat :: forall a. Num a => NonEmpty (BLen' a) -> BLen' a
sconcat :: NonEmpty (BLen' a) -> BLen' a
$cstimes :: forall a b. (Num a, Integral b) => b -> BLen' a -> BLen' a
stimes :: forall b. Integral b => b -> BLen' a -> BLen' a
Semigroup, Semigroup (BLen' a)
BLen' a
Semigroup (BLen' a)
-> BLen' a
-> (BLen' a -> BLen' a -> BLen' a)
-> ([BLen' a] -> BLen' a)
-> Monoid (BLen' a)
[BLen' a] -> BLen' a
BLen' a -> BLen' a -> BLen' a
forall a. Num a => Semigroup (BLen' a)
forall a. Num a => BLen' a
forall a. Num a => [BLen' a] -> BLen' a
forall a. Num a => BLen' a -> BLen' a -> BLen' a
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
$cmempty :: forall a. Num a => BLen' a
mempty :: BLen' a
$cmappend :: forall a. Num a => BLen' a -> BLen' a -> BLen' a
mappend :: BLen' a -> BLen' a -> BLen' a
$cmconcat :: forall a. Num a => [BLen' a] -> BLen' a
mconcat :: [BLen' a] -> BLen' a
Monoid) via Sum a

instance GenericFoldMap (BLen' Int) where
    type GenericFoldMapC (BLen' Int) a = BLen a
    genericFoldMapF :: forall a. GenericFoldMapC (BLen' Int) a => a -> BLen' Int
genericFoldMapF = Int -> BLen' Int
forall a. a -> BLen' a
BLen' (Int -> BLen' Int) -> (a -> Int) -> a -> BLen' Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int
forall a. BLen a => a -> Int
blen

-- | Measure the byte length of a term of the non-sum type @a@ via its 'Generic'
--   instance.
blenGenericNonSum
    :: forall {cd} {f} {asserts} a
    .  ( Generic a, Rep a ~ D1 cd f, GFoldMapNonSum (BLen' Int) f
       , asserts ~ '[ 'NoEmpty, 'NoSum], ApplyGCAsserts asserts f)
    => a -> Int
blenGenericNonSum :: forall {cd :: Meta} {f :: Type -> Type} {asserts :: [GCAssert]} a.
(Generic a, Rep a ~ D1 cd f, GFoldMapNonSum (BLen' Int) f,
 asserts ~ '[ 'NoEmpty, 'NoSum], ApplyGCAsserts asserts f) =>
a -> Int
blenGenericNonSum = BLen' Int -> Int
forall a. BLen' a -> a
getBLen' (BLen' Int -> Int) -> (a -> BLen' Int) -> a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (asserts :: [GCAssert]) m a.
(Generic a, Rep a ~ D1 cd f, GFoldMapNonSum m f,
 ApplyGCAsserts asserts f) =>
a -> m
forall {cd :: Meta} {f :: Type -> Type} (asserts :: [GCAssert]) m
       a.
(Generic a, Rep a ~ D1 cd f, GFoldMapNonSum m f,
 ApplyGCAsserts asserts f) =>
a -> m
genericFoldMapNonSum @asserts

-- | Measure the byte length of a term of the sum type @a@ via its 'Generic'
--   instance.
--
-- You must provide a function to obtain the byte length for the prefix tag, via
-- inspecting the reified constructor names. This is regrettably inefficient.
-- Alas. Do write your own instance if you want better performance!
blenGenericSum
    :: forall {cd} {f} {asserts} a
    .  (Generic a, Rep a ~ D1 cd f, GFoldMapSum 'SumOnly (BLen' Int) f
       , asserts ~ '[ 'NoEmpty, 'NeedSum], ApplyGCAsserts asserts f)
    => (String -> Int) -> a -> Int
blenGenericSum :: forall {cd :: Meta} {f :: Type -> Type} {asserts :: [GCAssert]} a.
(Generic a, Rep a ~ D1 cd f, GFoldMapSum 'SumOnly (BLen' Int) f,
 asserts ~ '[ 'NoEmpty, 'NeedSum], ApplyGCAsserts asserts f) =>
(String -> Int) -> a -> Int
blenGenericSum String -> Int
f = BLen' Int -> Int
forall a. BLen' a -> a
getBLen' (BLen' Int -> Int) -> (a -> BLen' Int) -> a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {cd :: Meta} {f :: Type -> Type} (opts :: SumOpts)
       (asserts :: [GCAssert]) m a.
(Generic a, Rep a ~ D1 cd f, GFoldMapSum opts m f,
 ApplyGCAsserts asserts f) =>
(String -> m) -> a -> m
forall (opts :: SumOpts) (asserts :: [GCAssert]) m a.
(Generic a, Rep a ~ D1 cd f, GFoldMapSum opts m f,
 ApplyGCAsserts asserts f) =>
(String -> m) -> a -> m
genericFoldMapSum @'SumOnly @asserts (Int -> BLen' Int
forall a. a -> BLen' a
BLen' (Int -> BLen' Int) -> (String -> Int) -> String -> BLen' Int
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Int
f)

instance TypeError ENoEmpty => BLen Void where blen :: Void -> Int
blen = Void -> Int
forall a. HasCallStack => a
undefined
instance TypeError ENoSum => BLen (Either a b) where blen :: Either a b -> Int
blen = Either a b -> Int
forall a. HasCallStack => a
undefined

instance BLen Write where
    {-# INLINE blen #-}
    blen :: Write -> Int
blen = Write -> Int
writeSize

-- | Unit type has length 0.
instance BLen () where
    {-# INLINE blen #-}
    blen :: () -> Int
blen () = Int
0

-- | Sum tuples.
instance (BLen l, BLen r) => BLen (l, r) where
    {-# INLINE blen #-}
    blen :: (l, r) -> Int
blen (l
l, r
r) = l -> Int
forall a. BLen a => a -> Int
blen l
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ r -> Int
forall a. BLen a => a -> Int
blen r
r

-- | _O(n)_ Sum the length of each element of a list.
instance BLen a => BLen [a] where
    {-# INLINE blen #-}
    blen :: [a] -> Int
blen = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> ([a] -> [Int]) -> [a] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Int) -> [a] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map a -> Int
forall a. BLen a => a -> Int
blen

-- | Length of a bytestring is fairly obvious.
instance BLen B.ByteString where
    {-# INLINE blen #-}
    blen :: ByteString -> Int
blen = ByteString -> Int
B.length

-- Machine integers have a constant byte length.
deriving via CBLenly Word8  instance BLen Word8
deriving via CBLenly  Int8  instance BLen  Int8
deriving via CBLenly Word16 instance BLen Word16
deriving via CBLenly  Int16 instance BLen  Int16
deriving via CBLenly Word32 instance BLen Word32
deriving via CBLenly  Int32 instance BLen  Int32
deriving via CBLenly Word64 instance BLen Word64
deriving via CBLenly  Int64 instance BLen  Int64

--------------------------------------------------------------------------------

-- | Deriving via wrapper for types which may derive a 'BLen' instance through
--   an existing 'IsCBLen' instance.
--
-- Examples of such types include machine integers, and explicitly-sized types
-- (e.g. "Binrep.Type.Sized").
newtype CBLenly a = CBLenly { forall a. CBLenly a -> a
unCBLenly :: a }
instance KnownNat (CBLen a) => BLen (CBLenly a) where
    {-# INLINE blen #-}
    blen :: CBLenly a -> Int
blen CBLenly a
_ = forall {k} (a :: k) (n :: Natural).
(n ~ CBLen a, KnownNat n) =>
Int
forall a (n :: Natural). (n ~ CBLen a, KnownNat n) => Int
cblen @a

-- | Reify a type's constant byte length to the term level.
cblen :: forall a n. (n ~ CBLen a, KnownNat n) => Int
cblen :: forall {k} (a :: k) (n :: Natural).
(n ~ CBLen a, KnownNat n) =>
Int
cblen = forall (n :: Natural). KnownNat n => Int
natValInt @n
{-# INLINE cblen #-}