{-# LANGUAGE OverloadedStrings #-}

module PlutusCore.Size
    ( Size (..)
    , kindSize
    , typeSize
    , tyVarDeclSize
    , termSize
    , varDeclSize
    , programSize
    , serialisedSize
    ) where

import PlutusPrelude

import PlutusCore.Core

import Control.Lens
import Data.ByteString qualified as BS
import Data.Monoid
import Flat hiding (to)

newtype Size = Size
    { Size -> Integer
unSize :: Integer
    } deriving stock (Int -> Size -> ShowS
[Size] -> ShowS
Size -> String
(Int -> Size -> ShowS)
-> (Size -> String) -> ([Size] -> ShowS) -> Show Size
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Size -> ShowS
showsPrec :: Int -> Size -> ShowS
$cshow :: Size -> String
show :: Size -> String
$cshowList :: [Size] -> ShowS
showList :: [Size] -> ShowS
Show)
      deriving newtype ((forall ann. Size -> Doc ann)
-> (forall ann. [Size] -> Doc ann) -> Pretty Size
forall ann. [Size] -> Doc ann
forall ann. Size -> Doc ann
forall a.
(forall ann. a -> Doc ann)
-> (forall ann. [a] -> Doc ann) -> Pretty a
$cpretty :: forall ann. Size -> Doc ann
pretty :: forall ann. Size -> Doc ann
$cprettyList :: forall ann. [Size] -> Doc ann
prettyList :: forall ann. [Size] -> Doc ann
Pretty, Size -> Size -> Bool
(Size -> Size -> Bool) -> (Size -> Size -> Bool) -> Eq Size
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Size -> Size -> Bool
== :: Size -> Size -> Bool
$c/= :: Size -> Size -> Bool
/= :: Size -> Size -> Bool
Eq, Eq Size
Eq Size =>
(Size -> Size -> Ordering)
-> (Size -> Size -> Bool)
-> (Size -> Size -> Bool)
-> (Size -> Size -> Bool)
-> (Size -> Size -> Bool)
-> (Size -> Size -> Size)
-> (Size -> Size -> Size)
-> Ord Size
Size -> Size -> Bool
Size -> Size -> Ordering
Size -> Size -> Size
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Size -> Size -> Ordering
compare :: Size -> Size -> Ordering
$c< :: Size -> Size -> Bool
< :: Size -> Size -> Bool
$c<= :: Size -> Size -> Bool
<= :: Size -> Size -> Bool
$c> :: Size -> Size -> Bool
> :: Size -> Size -> Bool
$c>= :: Size -> Size -> Bool
>= :: Size -> Size -> Bool
$cmax :: Size -> Size -> Size
max :: Size -> Size -> Size
$cmin :: Size -> Size -> Size
min :: Size -> Size -> Size
Ord, Integer -> Size
Size -> Size
Size -> Size -> Size
(Size -> Size -> Size)
-> (Size -> Size -> Size)
-> (Size -> Size -> Size)
-> (Size -> Size)
-> (Size -> Size)
-> (Size -> Size)
-> (Integer -> Size)
-> Num Size
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
$c+ :: Size -> Size -> Size
+ :: Size -> Size -> Size
$c- :: Size -> Size -> Size
- :: Size -> Size -> Size
$c* :: Size -> Size -> Size
* :: Size -> Size -> Size
$cnegate :: Size -> Size
negate :: Size -> Size
$cabs :: Size -> Size
abs :: Size -> Size
$csignum :: Size -> Size
signum :: Size -> Size
$cfromInteger :: Integer -> Size
fromInteger :: Integer -> Size
Num)
      deriving (NonEmpty Size -> Size
Size -> Size -> Size
(Size -> Size -> Size)
-> (NonEmpty Size -> Size)
-> (forall b. Integral b => b -> Size -> Size)
-> Semigroup Size
forall b. Integral b => b -> Size -> Size
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
$c<> :: Size -> Size -> Size
<> :: Size -> Size -> Size
$csconcat :: NonEmpty Size -> Size
sconcat :: NonEmpty Size -> Size
$cstimes :: forall b. Integral b => b -> Size -> Size
stimes :: forall b. Integral b => b -> Size -> Size
Semigroup, Semigroup Size
Size
Semigroup Size =>
Size -> (Size -> Size -> Size) -> ([Size] -> Size) -> Monoid Size
[Size] -> Size
Size -> Size -> Size
forall a.
Semigroup a =>
a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
$cmempty :: Size
mempty :: Size
$cmappend :: Size -> Size -> Size
mappend :: Size -> Size -> Size
$cmconcat :: [Size] -> Size
mconcat :: [Size] -> Size
Monoid) via Sum Integer

-- | Count the number of AST nodes in a kind.
--
-- >>> kindSize $ Type ()
-- Size {unSize = 1}
-- >>> kindSize $ KindArrow () (KindArrow () (Type ()) (Type ())) (Type ())
-- Size {unSize = 5}
kindSize :: Kind a -> Size
kindSize :: forall a. Kind a -> Size
kindSize Kind a
kind = [Size] -> Size
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
    [ Integer -> Size
Size Integer
1
    , Kind a
kind Kind a -> Getting Size (Kind a) Size -> Size
forall s a. s -> Getting a s a -> a
^. (Kind a -> Const Size (Kind a)) -> Kind a -> Const Size (Kind a)
forall ann (f :: * -> *).
Applicative f =>
(Kind ann -> f (Kind ann)) -> Kind ann -> f (Kind ann)
kindSubkinds ((Kind a -> Const Size (Kind a)) -> Kind a -> Const Size (Kind a))
-> Getting Size (Kind a) Size -> Getting Size (Kind a) Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Kind a -> Size) -> Getting Size (Kind a) Size
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Kind a -> Size
forall a. Kind a -> Size
kindSize
    ]

-- | Count the number of AST nodes in a type.
typeSize :: Type tyname uni ann -> Size
typeSize :: forall tyname (uni :: * -> *) ann. Type tyname uni ann -> Size
typeSize Type tyname uni ann
ty = [Size] -> Size
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
    [ Integer -> Size
Size Integer
1
    , Type tyname uni ann
ty Type tyname uni ann
-> Getting Size (Type tyname uni ann) Size -> Size
forall s a. s -> Getting a s a -> a
^. (Kind ann -> Const Size (Kind ann))
-> Type tyname uni ann -> Const Size (Type tyname uni ann)
forall tyname (uni :: * -> *) ann (f :: * -> *).
Applicative f =>
(Kind ann -> f (Kind ann))
-> Type tyname uni ann -> f (Type tyname uni ann)
typeSubkinds ((Kind ann -> Const Size (Kind ann))
 -> Type tyname uni ann -> Const Size (Type tyname uni ann))
-> ((Size -> Const Size Size) -> Kind ann -> Const Size (Kind ann))
-> Getting Size (Type tyname uni ann) Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Kind ann -> Size)
-> (Size -> Const Size Size) -> Kind ann -> Const Size (Kind ann)
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Kind ann -> Size
forall a. Kind a -> Size
kindSize
    , Type tyname uni ann
ty Type tyname uni ann
-> Getting Size (Type tyname uni ann) Size -> Size
forall s a. s -> Getting a s a -> a
^. (Type tyname uni ann -> Const Size (Type tyname uni ann))
-> Type tyname uni ann -> Const Size (Type tyname uni ann)
forall tyname (uni :: * -> *) ann (f :: * -> *).
Applicative f =>
(Type tyname uni ann -> f (Type tyname uni ann))
-> Type tyname uni ann -> f (Type tyname uni ann)
typeSubtypes ((Type tyname uni ann -> Const Size (Type tyname uni ann))
 -> Type tyname uni ann -> Const Size (Type tyname uni ann))
-> Getting Size (Type tyname uni ann) Size
-> Getting Size (Type tyname uni ann) Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type tyname uni ann -> Size)
-> Getting Size (Type tyname uni ann) Size
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Type tyname uni ann -> Size
forall tyname (uni :: * -> *) ann. Type tyname uni ann -> Size
typeSize
    ]

tyVarDeclSize :: TyVarDecl tyname ann -> Size
tyVarDeclSize :: forall tyname ann. TyVarDecl tyname ann -> Size
tyVarDeclSize TyVarDecl tyname ann
tyVarDecl = [Size] -> Size
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
    [ Integer -> Size
Size Integer
1
    , TyVarDecl tyname ann
tyVarDecl TyVarDecl tyname ann
-> Getting Size (TyVarDecl tyname ann) Size -> Size
forall s a. s -> Getting a s a -> a
^. (Kind ann -> Const Size (Kind ann))
-> TyVarDecl tyname ann -> Const Size (TyVarDecl tyname ann)
forall tyname a (f :: * -> *).
Applicative f =>
(Kind a -> f (Kind a))
-> TyVarDecl tyname a -> f (TyVarDecl tyname a)
tyVarDeclSubkinds ((Kind ann -> Const Size (Kind ann))
 -> TyVarDecl tyname ann -> Const Size (TyVarDecl tyname ann))
-> ((Size -> Const Size Size) -> Kind ann -> Const Size (Kind ann))
-> Getting Size (TyVarDecl tyname ann) Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Kind ann -> Size)
-> (Size -> Const Size Size) -> Kind ann -> Const Size (Kind ann)
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Kind ann -> Size
forall a. Kind a -> Size
kindSize
    ]

-- | Count the number of AST nodes in a term.
termSize :: Term tyname name uni fun ann -> Size
termSize :: forall tyname name (uni :: * -> *) fun ann.
Term tyname name uni fun ann -> Size
termSize Term tyname name uni fun ann
term = [Size] -> Size
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
    [ Integer -> Size
Size Integer
1
    , Term tyname name uni fun ann
term Term tyname name uni fun ann
-> Getting Size (Term tyname name uni fun ann) Size -> Size
forall s a. s -> Getting a s a -> a
^. (Kind ann -> Const Size (Kind ann))
-> Term tyname name uni fun ann
-> Const Size (Term tyname name uni fun ann)
forall tyname name (uni :: * -> *) fun ann (f :: * -> *).
Applicative f =>
(Kind ann -> f (Kind ann))
-> Term tyname name uni fun ann -> f (Term tyname name uni fun ann)
termSubkinds ((Kind ann -> Const Size (Kind ann))
 -> Term tyname name uni fun ann
 -> Const Size (Term tyname name uni fun ann))
-> ((Size -> Const Size Size) -> Kind ann -> Const Size (Kind ann))
-> Getting Size (Term tyname name uni fun ann) Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Kind ann -> Size)
-> (Size -> Const Size Size) -> Kind ann -> Const Size (Kind ann)
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Kind ann -> Size
forall a. Kind a -> Size
kindSize
    , Term tyname name uni fun ann
term Term tyname name uni fun ann
-> Getting Size (Term tyname name uni fun ann) Size -> Size
forall s a. s -> Getting a s a -> a
^. (Type tyname uni ann -> Const Size (Type tyname uni ann))
-> Term tyname name uni fun ann
-> Const Size (Term tyname name uni fun ann)
forall tyname name (uni :: * -> *) fun ann (f :: * -> *).
Applicative f =>
(Type tyname uni ann -> f (Type tyname uni ann))
-> Term tyname name uni fun ann -> f (Term tyname name uni fun ann)
termSubtypes ((Type tyname uni ann -> Const Size (Type tyname uni ann))
 -> Term tyname name uni fun ann
 -> Const Size (Term tyname name uni fun ann))
-> ((Size -> Const Size Size)
    -> Type tyname uni ann -> Const Size (Type tyname uni ann))
-> Getting Size (Term tyname name uni fun ann) Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type tyname uni ann -> Size)
-> (Size -> Const Size Size)
-> Type tyname uni ann
-> Const Size (Type tyname uni ann)
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Type tyname uni ann -> Size
forall tyname (uni :: * -> *) ann. Type tyname uni ann -> Size
typeSize
    , Term tyname name uni fun ann
term Term tyname name uni fun ann
-> Getting Size (Term tyname name uni fun ann) Size -> Size
forall s a. s -> Getting a s a -> a
^. (Term tyname name uni fun ann
 -> Const Size (Term tyname name uni fun ann))
-> Term tyname name uni fun ann
-> Const Size (Term tyname name uni fun ann)
forall tyname name (uni :: * -> *) fun ann (f :: * -> *).
Applicative f =>
(Term tyname name uni fun ann -> f (Term tyname name uni fun ann))
-> Term tyname name uni fun ann -> f (Term tyname name uni fun ann)
termSubterms ((Term tyname name uni fun ann
  -> Const Size (Term tyname name uni fun ann))
 -> Term tyname name uni fun ann
 -> Const Size (Term tyname name uni fun ann))
-> Getting Size (Term tyname name uni fun ann) Size
-> Getting Size (Term tyname name uni fun ann) Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term tyname name uni fun ann -> Size)
-> Getting Size (Term tyname name uni fun ann) Size
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Term tyname name uni fun ann -> Size
forall tyname name (uni :: * -> *) fun ann.
Term tyname name uni fun ann -> Size
termSize
    ]

varDeclSize :: VarDecl tyname name uni ann -> Size
varDeclSize :: forall tyname name (uni :: * -> *) ann.
VarDecl tyname name uni ann -> Size
varDeclSize VarDecl tyname name uni ann
varDecl = [Size] -> Size
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
    [ Integer -> Size
Size Integer
1
    , VarDecl tyname name uni ann
varDecl VarDecl tyname name uni ann
-> Getting Size (VarDecl tyname name uni ann) Size -> Size
forall s a. s -> Getting a s a -> a
^. (Type tyname uni ann -> Const Size (Type tyname uni ann))
-> VarDecl tyname name uni ann
-> Const Size (VarDecl tyname name uni ann)
forall tyname name (uni :: * -> *) a (f :: * -> *).
Applicative f =>
(Type tyname uni a -> f (Type tyname uni a))
-> VarDecl tyname name uni a -> f (VarDecl tyname name uni a)
varDeclSubtypes ((Type tyname uni ann -> Const Size (Type tyname uni ann))
 -> VarDecl tyname name uni ann
 -> Const Size (VarDecl tyname name uni ann))
-> ((Size -> Const Size Size)
    -> Type tyname uni ann -> Const Size (Type tyname uni ann))
-> Getting Size (VarDecl tyname name uni ann) Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type tyname uni ann -> Size)
-> (Size -> Const Size Size)
-> Type tyname uni ann
-> Const Size (Type tyname uni ann)
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Type tyname uni ann -> Size
forall tyname (uni :: * -> *) ann. Type tyname uni ann -> Size
typeSize
    ]

-- | Count the number of AST nodes in a program.
programSize :: Program tyname name uni fun ann -> Size
programSize :: forall tyname name (uni :: * -> *) fun ann.
Program tyname name uni fun ann -> Size
programSize (Program ann
_ Version
_ Term tyname name uni fun ann
t) = Term tyname name uni fun ann -> Size
forall tyname name (uni :: * -> *) fun ann.
Term tyname name uni fun ann -> Size
termSize Term tyname name uni fun ann
t

-- | Compute the size of the serializabled form of a value.
serialisedSize :: Flat a => a -> Integer
serialisedSize :: forall a. Flat a => a -> Integer
serialisedSize = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> (a -> Int) -> a -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Int
BS.length (ByteString -> Int) -> (a -> ByteString) -> a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ByteString
forall a. Flat a => a -> ByteString
flat