{-# LANGUAGE LambdaCase #-}

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

import UntypedPlutusCore.Core

import Data.ByteString qualified as BS
import Flat

-- | Count the number of AST nodes in a term.
termSize :: Term name uni fun ann -> Integer
termSize :: Term name uni fun ann -> Integer
termSize = \case
    Var{}        -> Integer
1
    Delay ann
_ Term name uni fun ann
t    -> Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Term name uni fun ann -> Integer
forall name (uni :: * -> *) fun ann.
Term name uni fun ann -> Integer
termSize Term name uni fun ann
t
    Apply ann
_ Term name uni fun ann
t Term name uni fun ann
t' -> Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Term name uni fun ann -> Integer
forall name (uni :: * -> *) fun ann.
Term name uni fun ann -> Integer
termSize Term name uni fun ann
t Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Term name uni fun ann -> Integer
forall name (uni :: * -> *) fun ann.
Term name uni fun ann -> Integer
termSize Term name uni fun ann
t'
    LamAbs ann
_ name
_ Term name uni fun ann
t -> Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Term name uni fun ann -> Integer
forall name (uni :: * -> *) fun ann.
Term name uni fun ann -> Integer
termSize Term name uni fun ann
t
    Constant{}   -> Integer
1
    Builtin{}    -> Integer
1
    Force ann
_ Term name uni fun ann
t    -> Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Term name uni fun ann -> Integer
forall name (uni :: * -> *) fun ann.
Term name uni fun ann -> Integer
termSize Term name uni fun ann
t
    Error ann
_      -> Integer
1

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

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