{-# 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 :: forall a. HasTypeReps a => AccountingMap -> a -> Size
abstractSize AccountingMap
m a
a = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeRep -> Size
cost Seq TypeRep
trs
  where
    trs :: Seq TypeRep
trs = forall a. HasTypeReps a => a -> Seq TypeRep
typeReps a
a
    cost :: TypeRep -> Size
cost TypeRep
t = 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 = forall a. Typeable a => a -> TypeRep
typeOf a
a forall a. a -> Seq a -> Seq a
<| forall (f :: * -> *) a. GHasTypeReps f => f a -> Seq TypeRep
gTypeReps (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 :: forall a. U1 a -> Seq TypeRep
gTypeReps U1 a
U1 = 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 :: forall a. (:*:) a b a -> Seq TypeRep
gTypeReps (a a
a :*: b a
b) = forall (f :: * -> *) a. GHasTypeReps f => f a -> Seq TypeRep
gTypeReps a a
a forall a. Seq a -> Seq a -> Seq a
>< forall (f :: * -> *) a. GHasTypeReps f => f a -> Seq TypeRep
gTypeReps b a
b

instance (GHasTypeReps a, GHasTypeReps b) => GHasTypeReps (a :+: b) where
  gTypeReps :: forall a. (:+:) a b a -> Seq TypeRep
gTypeReps (L1 a a
a) = forall (f :: * -> *) a. GHasTypeReps f => f a -> Seq TypeRep
gTypeReps a a
a
  gTypeReps (R1 b a
b) = 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 :: forall a. M1 i c a a -> Seq TypeRep
gTypeReps (M1 a a
x) = 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 :: forall a. K1 i a a -> Seq TypeRep
gTypeReps (K1 a
x) = 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 = forall a. Typeable a => a -> TypeRep
typeOf Maybe a
x forall a. a -> Seq a -> Seq a
<| forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] 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 = forall a. Typeable a => a -> TypeRep
typeOf [a]
xs forall a. a -> Seq a -> Seq a
<| forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap 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 = forall a. Typeable a => a -> TypeRep
typeOf Set a
xs forall a. a -> Seq a -> Seq a
<| forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap 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) = forall a. Typeable a => a -> TypeRep
typeOf (a, b)
t forall a. a -> Seq a -> Seq a
<| (forall a. HasTypeReps a => a -> Seq TypeRep
typeReps a
a forall a. Seq a -> Seq a -> Seq a
>< forall a. HasTypeReps a => a -> Seq TypeRep
typeReps b
b)

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

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

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

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

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

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

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

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

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

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

instance HasTypeReps Word64 where
  typeReps :: Word64 -> Seq TypeRep
typeReps Word64
x = [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 = [forall a. Typeable a => a -> TypeRep
typeOf Digest SHA256
x]

instance HasTypeReps ShortHash where
  typeReps :: ShortHash -> Seq TypeRep
typeReps ShortHash
x = [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 = [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 Hash ShortHash ()
h Word64
i)) =
    forall a. Typeable a => a -> TypeRep
typeOf Word64
i forall a. a -> Seq a -> Seq a
<| forall a. Size -> a -> Seq a
Seq.replicate (ByteString -> Size
BS.length (forall h a. Hash h a -> ByteString
hashToBytes Hash ShortHash ()
h)) (forall a. Typeable a => a -> TypeRep
typeOf (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
_ = [forall a. Typeable a => a -> TypeRep
typeOf (forall a. HasCallStack => a
undefined :: Int)]