{-# LANGUAGE DeriveAnyClass #-}

module PlutusCore.AstSize (
  AstSize (..),
  kindAstSize,
  typeAstSize,
  tyVarDeclAstSize,
  termAstSize,
  varDeclAstSize,
  programAstSize,
  serialisedAstSize,
) where

import PlutusPrelude

import PlutusCore.Core

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

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

{-| Count the number of AST nodes in a kind.

>>> kindAstSize $ Type ()
AstSize {unAstSize = 1}
>>> kindAstSize $ KindArrow () (KindArrow () (Type ()) (Type ())) (Type ())
AstSize {unAstSize = 5}
-}
kindAstSize :: Kind a -> AstSize
kindAstSize :: forall a. Kind a -> AstSize
kindAstSize Kind a
kind =
  [AstSize] -> AstSize
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
    [ Integer -> AstSize
AstSize Integer
1
    , Kind a
kind Kind a -> Getting AstSize (Kind a) AstSize -> AstSize
forall s a. s -> Getting a s a -> a
^. (Kind a -> Const AstSize (Kind a))
-> Kind a -> Const AstSize (Kind a)
forall ann (f :: * -> *).
Applicative f =>
(Kind ann -> f (Kind ann)) -> Kind ann -> f (Kind ann)
kindSubkinds ((Kind a -> Const AstSize (Kind a))
 -> Kind a -> Const AstSize (Kind a))
-> Getting AstSize (Kind a) AstSize
-> Getting AstSize (Kind a) AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Kind a -> AstSize) -> Getting AstSize (Kind a) AstSize
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Kind a -> AstSize
forall a. Kind a -> AstSize
kindAstSize
    ]

-- | Count the number of AST nodes in a type.
typeAstSize :: Type tyname uni ann -> AstSize
typeAstSize :: forall tyname (uni :: * -> *) ann. Type tyname uni ann -> AstSize
typeAstSize Type tyname uni ann
ty =
  [AstSize] -> AstSize
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
    [ Integer -> AstSize
AstSize Integer
1
    , Type tyname uni ann
ty Type tyname uni ann
-> Getting AstSize (Type tyname uni ann) AstSize -> AstSize
forall s a. s -> Getting a s a -> a
^. (Kind ann -> Const AstSize (Kind ann))
-> Type tyname uni ann -> Const AstSize (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 AstSize (Kind ann))
 -> Type tyname uni ann -> Const AstSize (Type tyname uni ann))
-> ((AstSize -> Const AstSize AstSize)
    -> Kind ann -> Const AstSize (Kind ann))
-> Getting AstSize (Type tyname uni ann) AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Kind ann -> AstSize)
-> (AstSize -> Const AstSize AstSize)
-> Kind ann
-> Const AstSize (Kind ann)
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Kind ann -> AstSize
forall a. Kind a -> AstSize
kindAstSize
    , Type tyname uni ann
ty Type tyname uni ann
-> Getting AstSize (Type tyname uni ann) AstSize -> AstSize
forall s a. s -> Getting a s a -> a
^. (Type tyname uni ann -> Const AstSize (Type tyname uni ann))
-> Type tyname uni ann -> Const AstSize (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 AstSize (Type tyname uni ann))
 -> Type tyname uni ann -> Const AstSize (Type tyname uni ann))
-> Getting AstSize (Type tyname uni ann) AstSize
-> Getting AstSize (Type tyname uni ann) AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type tyname uni ann -> AstSize)
-> Getting AstSize (Type tyname uni ann) AstSize
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Type tyname uni ann -> AstSize
forall tyname (uni :: * -> *) ann. Type tyname uni ann -> AstSize
typeAstSize
    ]

tyVarDeclAstSize :: TyVarDecl tyname ann -> AstSize
tyVarDeclAstSize :: forall tyname ann. TyVarDecl tyname ann -> AstSize
tyVarDeclAstSize TyVarDecl tyname ann
tyVarDecl =
  [AstSize] -> AstSize
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
    [ Integer -> AstSize
AstSize Integer
1
    , TyVarDecl tyname ann
tyVarDecl TyVarDecl tyname ann
-> Getting AstSize (TyVarDecl tyname ann) AstSize -> AstSize
forall s a. s -> Getting a s a -> a
^. (Kind ann -> Const AstSize (Kind ann))
-> TyVarDecl tyname ann -> Const AstSize (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 AstSize (Kind ann))
 -> TyVarDecl tyname ann -> Const AstSize (TyVarDecl tyname ann))
-> ((AstSize -> Const AstSize AstSize)
    -> Kind ann -> Const AstSize (Kind ann))
-> Getting AstSize (TyVarDecl tyname ann) AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Kind ann -> AstSize)
-> (AstSize -> Const AstSize AstSize)
-> Kind ann
-> Const AstSize (Kind ann)
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Kind ann -> AstSize
forall a. Kind a -> AstSize
kindAstSize
    ]

-- | Count the number of AST nodes in a term.
termAstSize :: Term tyname name uni fun ann -> AstSize
termAstSize :: forall tyname name (uni :: * -> *) fun ann.
Term tyname name uni fun ann -> AstSize
termAstSize Term tyname name uni fun ann
term =
  [AstSize] -> AstSize
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
    [ Integer -> AstSize
AstSize Integer
1
    , Term tyname name uni fun ann
term Term tyname name uni fun ann
-> Getting AstSize (Term tyname name uni fun ann) AstSize
-> AstSize
forall s a. s -> Getting a s a -> a
^. (Kind ann -> Const AstSize (Kind ann))
-> Term tyname name uni fun ann
-> Const AstSize (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 AstSize (Kind ann))
 -> Term tyname name uni fun ann
 -> Const AstSize (Term tyname name uni fun ann))
-> ((AstSize -> Const AstSize AstSize)
    -> Kind ann -> Const AstSize (Kind ann))
-> Getting AstSize (Term tyname name uni fun ann) AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Kind ann -> AstSize)
-> (AstSize -> Const AstSize AstSize)
-> Kind ann
-> Const AstSize (Kind ann)
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Kind ann -> AstSize
forall a. Kind a -> AstSize
kindAstSize
    , Term tyname name uni fun ann
term Term tyname name uni fun ann
-> Getting AstSize (Term tyname name uni fun ann) AstSize
-> AstSize
forall s a. s -> Getting a s a -> a
^. (Type tyname uni ann -> Const AstSize (Type tyname uni ann))
-> Term tyname name uni fun ann
-> Const AstSize (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 AstSize (Type tyname uni ann))
 -> Term tyname name uni fun ann
 -> Const AstSize (Term tyname name uni fun ann))
-> ((AstSize -> Const AstSize AstSize)
    -> Type tyname uni ann -> Const AstSize (Type tyname uni ann))
-> Getting AstSize (Term tyname name uni fun ann) AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type tyname uni ann -> AstSize)
-> (AstSize -> Const AstSize AstSize)
-> Type tyname uni ann
-> Const AstSize (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 -> AstSize
forall tyname (uni :: * -> *) ann. Type tyname uni ann -> AstSize
typeAstSize
    , Term tyname name uni fun ann
term Term tyname name uni fun ann
-> Getting AstSize (Term tyname name uni fun ann) AstSize
-> AstSize
forall s a. s -> Getting a s a -> a
^. (Term tyname name uni fun ann
 -> Const AstSize (Term tyname name uni fun ann))
-> Term tyname name uni fun ann
-> Const AstSize (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 AstSize (Term tyname name uni fun ann))
 -> Term tyname name uni fun ann
 -> Const AstSize (Term tyname name uni fun ann))
-> Getting AstSize (Term tyname name uni fun ann) AstSize
-> Getting AstSize (Term tyname name uni fun ann) AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term tyname name uni fun ann -> AstSize)
-> Getting AstSize (Term tyname name uni fun ann) AstSize
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Term tyname name uni fun ann -> AstSize
forall tyname name (uni :: * -> *) fun ann.
Term tyname name uni fun ann -> AstSize
termAstSize
    ]

varDeclAstSize :: VarDecl tyname name uni ann -> AstSize
varDeclAstSize :: forall tyname name (uni :: * -> *) ann.
VarDecl tyname name uni ann -> AstSize
varDeclAstSize VarDecl tyname name uni ann
varDecl =
  [AstSize] -> AstSize
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold
    [ Integer -> AstSize
AstSize Integer
1
    , VarDecl tyname name uni ann
varDecl VarDecl tyname name uni ann
-> Getting AstSize (VarDecl tyname name uni ann) AstSize -> AstSize
forall s a. s -> Getting a s a -> a
^. (Type tyname uni ann -> Const AstSize (Type tyname uni ann))
-> VarDecl tyname name uni ann
-> Const AstSize (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 AstSize (Type tyname uni ann))
 -> VarDecl tyname name uni ann
 -> Const AstSize (VarDecl tyname name uni ann))
-> ((AstSize -> Const AstSize AstSize)
    -> Type tyname uni ann -> Const AstSize (Type tyname uni ann))
-> Getting AstSize (VarDecl tyname name uni ann) AstSize
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type tyname uni ann -> AstSize)
-> (AstSize -> Const AstSize AstSize)
-> Type tyname uni ann
-> Const AstSize (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 -> AstSize
forall tyname (uni :: * -> *) ann. Type tyname uni ann -> AstSize
typeAstSize
    ]

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

-- | Compute the size of the serializabled form of a value.
serialisedAstSize :: (Flat a) => a -> Integer
serialisedAstSize :: forall a. Flat a => a -> Integer
serialisedAstSize = 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