{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE TypeOperators #-}

-- | An approach to computing the abstract size of data using 'TypeRep'.
module Data.AbstractSize
  ( HasTypeReps,
    typeReps,
    abstractSize,
    AccountingMap,
    Size,
  )
where

import Cardano.Crypto.DSIGN.Class (SignedDSIGN (SignedDSIGN), VerKeyDSIGN)
import Cardano.Crypto.DSIGN.Mock (MockDSIGN, SigDSIGN (SigMockDSIGN))
import Cardano.Crypto.Hash (Hash, hashToBytes)
import Cardano.Crypto.Hash.Short (ShortHash)
import qualified Crypto.Hash as Crypto
import qualified Data.ByteString as BS
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Sequence (Seq, empty, (<|), (><))
import qualified Data.Sequence as Seq
import Data.Set (Set)
import Data.Typeable (TypeRep, Typeable, typeOf)
import Data.Word (Word16, Word32, Word64, Word8)
import GHC.Generics
  ( Generic,
    K1 (K1),
    M1 (M1),
    Rep,
    U1 (U1),
    from,
    (:*:) ((:*:)),
    (:+:) (L1, R1),
  )
import GHC.Natural (Natural)

-- | @abstractSize m a@ computes the abstract size of @a@, using the accounting
-- map @m@. The map @m@ determines the abstract size of each 'TypeRep'
-- contained in @a@, and this function simply adds all the individual abstract
-- sizes. To be able to extract the type representations ('TypeRep's) inside
-- @a@, we require it to be an instance of 'HasTypeReps'.
--
-- Examples:
--
-- >>> :set -XOverloadedLists
-- >>> import Data.Typeable (typeOf)
-- >>> abstractSize [(typeOf (undefined:: Char), 10)] 'a'
-- 10
--
-- >>> abstractSize [(typeOf 'x', 10)] "hello"
-- 50
--
-- >>> abstractSize [(typeOf 'x', 10), (typeOf True, 100)] ("hello", False)
-- 150
--
-- >>> abstractSize [(typeOf (undefined :: [Int]), 6), (typeOf (1 :: Int), 1)] ([0, 1, 2, 3] :: [Int])
-- 10
--
-- >>> abstractSize [(typeOf (undefined :: [Int]), 3), (typeOf (1 :: Int), -1)] ([0, 1, 2] :: [Int])
-- 0
abstractSize :: HasTypeReps a => AccountingMap -> a -> Size
abstractSize :: AccountingMap -> a -> Size
abstractSize AccountingMap
m a
a = Seq Size -> Size
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Seq Size -> Size) -> Seq Size -> Size
forall a b. (a -> b) -> a -> b
$ (TypeRep -> Size) -> Seq TypeRep -> Seq Size
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeRep -> Size
cost Seq TypeRep
trs
  where
    trs :: Seq TypeRep
trs = a -> Seq TypeRep
forall a. HasTypeReps a => a -> Seq TypeRep
typeReps a
a
    cost :: TypeRep -> Size
cost TypeRep
t = Size -> TypeRep -> AccountingMap -> Size
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Size
0 TypeRep
t AccountingMap
m

type Size = Int

type AccountingMap = Map TypeRep Size

--------------------------------------------------------------------------------
-- HasTypeReps class
--------------------------------------------------------------------------------

-- | The 'typeReps' function retrieves all the type representations found while
-- traversing the data given as parameter.
--
-- CAUTION: for newtypes, do not use 'deriving newtype (HasTypeReps)' to derive
-- instances, rather use 'deriving anyclass (HasTypeReps)'.
-- This is because we use these instances in 'abstractSize', and for that
-- we prefer to have the newtype wrapper type available for "costing".
-- The difference between 'newtype' and 'anyclass' instances is as follows:
--
--  newtype Hash = Hash { unHash :: Int }
--      deriving newtype (..., HasTypeReps)
--  > typeReps someHash = Seq.fromList [Int]
--  vs
--  newtype Hash = Hash { unHash :: Int }
--      deriving stock (...,Generics); deriving anyclass (HasTypeReps)
--  > typeReps someHash = Seq.fromList [Hash, Int]
--
-- Examples:
--
-- >>> typeReps "a"
-- fromList [[Char],Char]
--
-- >>> typeReps "ab"
-- fromList [[Char],Char,Char]
--
-- >>> typeReps ([] :: [Int])
-- fromList [[Int]]
--
-- >>> :set -XDeriveGeneric
-- >>> import GHC.Generics (Generic)
-- >>> data Foo = Foo [Int] (Char, Char) deriving (Generic)
-- >>> instance HasTypeReps Foo
-- >>> typeReps $ Foo [1, 2] ('a', 'b')
-- fromList [Foo,[Int],Int,Int,(Char,Char),Char,Char]
class HasTypeReps a where
  typeReps :: a -> Seq TypeRep
  default typeReps ::
    ( Generic a,
      GHasTypeReps (Rep a),
      Typeable a
    ) =>
    a ->
    Seq TypeRep
  typeReps a
a = a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf a
a TypeRep -> Seq TypeRep -> Seq TypeRep
forall a. a -> Seq a -> Seq a
<| Rep a Any -> Seq TypeRep
forall (f :: * -> *) a. GHasTypeReps f => f a -> Seq TypeRep
gTypeReps (a -> Rep a Any
forall a x. Generic a => a -> Rep a x
from a
a)

class GHasTypeReps f where
  gTypeReps :: f a -> Seq TypeRep

--------------------------------------------------------------------------------
-- GHasTypeReps instances
--------------------------------------------------------------------------------

-- | No types to report for a constructor without arguments.
instance GHasTypeReps U1 where
  gTypeReps :: U1 a -> Seq TypeRep
gTypeReps U1 a
U1 = Seq TypeRep
forall a. Seq a
empty

-- | The types in a product is the concatenation of the types found in the
-- values of the product terms.
instance (GHasTypeReps a, GHasTypeReps b) => GHasTypeReps (a :*: b) where
  gTypeReps :: (:*:) a b a -> Seq TypeRep
gTypeReps (a a
a :*: b a
b) = a a -> Seq TypeRep
forall (f :: * -> *) a. GHasTypeReps f => f a -> Seq TypeRep
gTypeReps a a
a Seq TypeRep -> Seq TypeRep -> Seq TypeRep
forall a. Seq a -> Seq a -> Seq a
>< b a -> Seq TypeRep
forall (f :: * -> *) a. GHasTypeReps f => f a -> Seq TypeRep
gTypeReps b a
b

instance (GHasTypeReps a, GHasTypeReps b) => GHasTypeReps (a :+: b) where
  gTypeReps :: (:+:) a b a -> Seq TypeRep
gTypeReps (L1 a a
a) = a a -> Seq TypeRep
forall (f :: * -> *) a. GHasTypeReps f => f a -> Seq TypeRep
gTypeReps a a
a
  gTypeReps (R1 b a
b) = b a -> Seq TypeRep
forall (f :: * -> *) a. GHasTypeReps f => f a -> Seq TypeRep
gTypeReps b a
b

-- | We do need to do anything for the metadata.
instance (GHasTypeReps a) => GHasTypeReps (M1 i c a) where
  gTypeReps :: M1 i c a a -> Seq TypeRep
gTypeReps (M1 a a
x) = a a -> Seq TypeRep
forall (f :: * -> *) a. GHasTypeReps f => f a -> Seq TypeRep
gTypeReps a a
x

-- | And the only interesting case, get the type of a type constructor
instance (HasTypeReps a) => GHasTypeReps (K1 i a) where
  gTypeReps :: K1 i a a -> Seq TypeRep
gTypeReps (K1 a
x) = a -> Seq TypeRep
forall a. HasTypeReps a => a -> Seq TypeRep
typeReps a
x

--------------------------------------------------------------------------------
-- HasTypeReps instances
--------------------------------------------------------------------------------

instance (Typeable a, HasTypeReps a) => HasTypeReps (Maybe a) where
  typeReps :: Maybe a -> Seq TypeRep
typeReps Maybe a
x = Maybe a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Maybe a
x TypeRep -> Seq TypeRep -> Seq TypeRep
forall a. a -> Seq a -> Seq a
<| Seq TypeRep -> (a -> Seq TypeRep) -> Maybe a -> Seq TypeRep
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] a -> Seq TypeRep
forall a. HasTypeReps a => a -> Seq TypeRep
typeReps Maybe a
x

instance (Typeable a, HasTypeReps a) => HasTypeReps [a] where
  typeReps :: [a] -> Seq TypeRep
typeReps [a]
xs = [a] -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf [a]
xs TypeRep -> Seq TypeRep -> Seq TypeRep
forall a. a -> Seq a -> Seq a
<| (a -> Seq TypeRep) -> [a] -> Seq TypeRep
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> Seq TypeRep
forall a. HasTypeReps a => a -> Seq TypeRep
typeReps [a]
xs

instance (Typeable a, HasTypeReps a) => HasTypeReps (Set a) where
  typeReps :: Set a -> Seq TypeRep
typeReps Set a
xs = Set a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Set a
xs TypeRep -> Seq TypeRep -> Seq TypeRep
forall a. a -> Seq a -> Seq a
<| (a -> Seq TypeRep) -> Set a -> Seq TypeRep
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> Seq TypeRep
forall a. HasTypeReps a => a -> Seq TypeRep
typeReps Set a
xs

instance
  ( Typeable a,
    Typeable b,
    HasTypeReps a,
    HasTypeReps b
  ) =>
  HasTypeReps (a, b)
  where
  typeReps :: (a, b) -> Seq TypeRep
typeReps t :: (a, b)
t@(a
a, b
b) = (a, b) -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf (a, b)
t TypeRep -> Seq TypeRep -> Seq TypeRep
forall a. a -> Seq a -> Seq a
<| (a -> Seq TypeRep
forall a. HasTypeReps a => a -> Seq TypeRep
typeReps a
a Seq TypeRep -> Seq TypeRep -> Seq TypeRep
forall a. Seq a -> Seq a -> Seq a
>< b -> Seq TypeRep
forall a. HasTypeReps a => a -> Seq TypeRep
typeReps b
b)

instance HasTypeReps Bool where
  typeReps :: Bool -> Seq TypeRep
typeReps Bool
x = [Bool -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Bool
x]

instance HasTypeReps Char where
  typeReps :: Char -> Seq TypeRep
typeReps Char
x = [Char -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Char
x]

instance HasTypeReps Int where
  typeReps :: Size -> Seq TypeRep
typeReps Size
x = [Size -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Size
x]

instance HasTypeReps Integer where
  typeReps :: Integer -> Seq TypeRep
typeReps Integer
x = [Integer -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Integer
x]

instance HasTypeReps Double where
  typeReps :: Double -> Seq TypeRep
typeReps Double
x = [Double -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Double
x]

instance HasTypeReps Natural where
  typeReps :: Natural -> Seq TypeRep
typeReps Natural
x = [Natural -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Natural
x]

instance HasTypeReps Word where
  typeReps :: Word -> Seq TypeRep
typeReps Word
x = [Word -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Word
x]

instance HasTypeReps Word8 where
  typeReps :: Word8 -> Seq TypeRep
typeReps Word8
x = [Word8 -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Word8
x]

instance HasTypeReps Word16 where
  typeReps :: Word16 -> Seq TypeRep
typeReps Word16
x = [Word16 -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Word16
x]

instance HasTypeReps Word32 where
  typeReps :: Word32 -> Seq TypeRep
typeReps Word32
x = [Word32 -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Word32
x]

instance HasTypeReps Word64 where
  typeReps :: Word64 -> Seq TypeRep
typeReps Word64
x = [Word64 -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Word64
x]

instance HasTypeReps (Crypto.Digest Crypto.SHA256) where
  typeReps :: Digest SHA256 -> Seq TypeRep
typeReps Digest SHA256
x = [Digest SHA256 -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Digest SHA256
x]

instance HasTypeReps ShortHash where
  typeReps :: ShortHash -> Seq TypeRep
typeReps ShortHash
x = [ShortHash -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf ShortHash
x]

instance Typeable a => HasTypeReps (Hash ShortHash a) where
  typeReps :: Hash ShortHash a -> Seq TypeRep
typeReps Hash ShortHash a
x = [Hash ShortHash a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Hash ShortHash a
x]

instance HasTypeReps (SignedDSIGN MockDSIGN a) where
  -- A mock signature consists of a 'ByteString' (which is in turn a short hash)
  -- and a 'Word64'. For the 'ByteString' representation we return one character
  -- per byte.
  typeReps :: SignedDSIGN MockDSIGN a -> Seq TypeRep
typeReps (SignedDSIGN (SigMockDSIGN h i)) =
    Word64 -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf Word64
i TypeRep -> Seq TypeRep -> Seq TypeRep
forall a. a -> Seq a -> Seq a
<| Size -> TypeRep -> Seq TypeRep
forall a. Size -> a -> Seq a
Seq.replicate (ByteString -> Size
BS.length (Hash ShortHash () -> ByteString
forall h a. Hash h a -> ByteString
hashToBytes Hash ShortHash ()
h)) (Char -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf (Char
forall a. HasCallStack => a
undefined :: Char))

instance HasTypeReps (VerKeyDSIGN MockDSIGN) where
  -- A mock verification key is just an 'Int'.
  typeReps :: VerKeyDSIGN MockDSIGN -> Seq TypeRep
typeReps VerKeyDSIGN MockDSIGN
_ = [Size -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf (Size
forall a. HasCallStack => a
undefined :: Int)]