module UntypedPlutusCore.Size
    ( Size (..)
    , termSize
    , programSize
    , serialisedSize
    ) where

import PlutusCore.Size (Size (..))
import UntypedPlutusCore.Core

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

-- | Count the number of AST nodes in a term.
termSize :: Term name uni fun ann -> Size
termSize :: forall name (uni :: * -> *) fun ann. Term name uni fun ann -> Size
termSize Term 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 name uni fun ann
term Term name uni fun ann
-> Getting Size (Term name uni fun ann) Size -> Size
forall s a. s -> Getting a s a -> a
^. (Term name uni fun ann -> Const Size (Term name uni fun ann))
-> Term name uni fun ann -> Const Size (Term name uni fun ann)
forall name (uni :: * -> *) fun ann (f :: * -> *).
Applicative f =>
(Term name uni fun ann -> f (Term name uni fun ann))
-> Term name uni fun ann -> f (Term name uni fun ann)
termSubterms ((Term name uni fun ann -> Const Size (Term name uni fun ann))
 -> Term name uni fun ann -> Const Size (Term name uni fun ann))
-> Getting Size (Term name uni fun ann) Size
-> Getting Size (Term name uni fun ann) Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term name uni fun ann -> Size)
-> Getting Size (Term name uni fun ann) Size
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to Term name uni fun ann -> Size
forall name (uni :: * -> *) fun ann. Term name uni fun ann -> Size
termSize
    ]

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

-- | Compute the size of the serialized 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