{-# 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
showList :: [Size] -> ShowS
$cshowList :: [Size] -> ShowS
show :: Size -> String
$cshow :: Size -> String
showsPrec :: Int -> Size -> ShowS
$cshowsPrec :: Int -> Size -> ShowS
Show)
      deriving newtype ([Size] -> Doc ann
Size -> Doc ann
(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
prettyList :: [Size] -> Doc ann
$cprettyList :: forall ann. [Size] -> Doc ann
pretty :: Size -> Doc ann
$cpretty :: 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
/= :: Size -> Size -> Bool
$c/= :: Size -> Size -> Bool
== :: Size -> Size -> Bool
$c== :: 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
min :: Size -> Size -> Size
$cmin :: Size -> Size -> Size
max :: Size -> Size -> Size
$cmax :: Size -> Size -> Size
>= :: Size -> Size -> Bool
$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
compare :: Size -> Size -> Ordering
$ccompare :: Size -> Size -> Ordering
$cp1Ord :: Eq 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
fromInteger :: Integer -> Size
$cfromInteger :: Integer -> Size
signum :: Size -> Size
$csignum :: Size -> Size
abs :: Size -> Size
$cabs :: Size -> Size
negate :: Size -> Size
$cnegate :: Size -> Size
* :: Size -> Size -> Size
$c* :: Size -> Size -> Size
- :: Size -> Size -> Size
$c- :: Size -> Size -> Size
+ :: Size -> Size -> Size
$c+ :: Size -> Size -> Size
Num)
      deriving (b -> Size -> Size
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
stimes :: b -> Size -> Size
$cstimes :: forall b. Integral b => b -> Size -> Size
sconcat :: NonEmpty Size -> Size
$csconcat :: NonEmpty Size -> Size
<> :: Size -> Size -> Size
$c<> :: Size -> 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
mconcat :: [Size] -> Size
$cmconcat :: [Size] -> Size
mappend :: Size -> Size -> Size
$cmappend :: Size -> Size -> Size
mempty :: Size
$cmempty :: Size
$cp1Monoid :: Semigroup 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 :: Kind a -> Size
kindSize Kind a
kind = [Size] -> Size
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. Traversal' (Kind ann) (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 :: Type tyname uni ann -> Size
typeSize Type tyname uni ann
ty = [Size] -> Size
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.
Traversal' (Type tyname uni ann) (Kind 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.
Traversal' (Type tyname uni ann) (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 :: TyVarDecl tyname ann -> Size
tyVarDeclSize TyVarDecl tyname ann
tyVarDecl = [Size] -> Size
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. Traversal' (TyVarDecl tyname a) (Kind 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 :: Term tyname name uni fun ann -> Size
termSize Term tyname name uni fun ann
term = [Size] -> Size
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.
Traversal' (Term tyname name uni fun ann) (Kind 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.
Traversal' (Term tyname name uni fun ann) (Type tyname uni 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.
Traversal'
  (Term tyname name uni fun ann) (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 fun ann -> Size
varDeclSize :: VarDecl tyname name uni fun ann -> Size
varDeclSize VarDecl tyname name uni fun ann
varDecl = [Size] -> Size
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
    [ Integer -> Size
Size Integer
1
    , VarDecl tyname name uni fun ann
varDecl VarDecl tyname name uni fun ann
-> Getting Size (VarDecl 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))
-> VarDecl tyname name uni fun ann
-> Const Size (VarDecl tyname name uni fun ann)
forall tyname name (uni :: * -> *) fun a.
Traversal' (VarDecl tyname name uni fun a) (Type tyname uni a)
varDeclSubtypes ((Type tyname uni ann -> Const Size (Type tyname uni ann))
 -> VarDecl tyname name uni fun ann
 -> Const Size (VarDecl tyname name uni fun ann))
-> ((Size -> Const Size Size)
    -> Type tyname uni ann -> Const Size (Type tyname uni ann))
-> Getting Size (VarDecl 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
    ]

-- | Count the number of AST nodes in a program.
programSize :: Program tyname name uni fun ann -> Size
programSize :: Program tyname name uni fun ann -> Size
programSize (Program ann
_ Version ann
_ 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 :: 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