Builtin Value: efficient retrieval of total size and max inner map size (#7319)

Author: zliu41

plutus-core/plutus-core/src/PlutusCore/Evaluation/Machine/ExMemoryUsage.hs

@@ -12,6 +12,8 @@ module PlutusCore.Evaluation.Machine.ExMemoryUsage
     , flattenCostRose
     , NumBytesCostedAsNumWords(..)
     , IntegerCostedLiterally(..)
+    , ValueTotalSize(..)
+    , ValueOuterOrMaxInner(..)
     ) where
 
 import PlutusCore.Crypto.BLS12_381.G1 as BLS12_381.G1
@@ -20,7 +22,8 @@ import PlutusCore.Crypto.BLS12_381.Pairing as BLS12_381.Pairing
 import PlutusCore.Data
 import PlutusCore.Evaluation.Machine.CostStream
 import PlutusCore.Evaluation.Machine.ExMemory
-import PlutusCore.Value
+import PlutusCore.Value (Value)
+import PlutusCore.Value qualified as Value
 
 import Data.ByteString qualified as BS
 import Data.Functor
@@ -371,11 +374,22 @@ instance ExMemoryUsage Data where
             B b        -> memoryUsage b
 
 instance ExMemoryUsage Value where
-    memoryUsage (Value v) = case Map.toList v of
-        []     -> CostRose 0 []
-        x : xs -> CostRose (f x) (flip CostRose [] . f <$> xs)
+    memoryUsage = singletonRose . fromIntegral . Value.totalSize
+
+-- | Measure the size of a `Value` by its `Value.totalSize`.
+newtype ValueTotalSize = ValueTotalSize { unValueTotalSize :: Value }
+
+instance ExMemoryUsage ValueTotalSize where
+    memoryUsage = singletonRose . fromIntegral . Value.totalSize . unValueTotalSize
+
+-- | Measure the size of a `Value` by taking the max of
+-- (size of the outer map, size of the largest inner map).
+newtype ValueOuterOrMaxInner = ValueOuterOrMaxInner { unValueOuterOrMaxInner :: Value }
+
+instance ExMemoryUsage ValueOuterOrMaxInner where
+    memoryUsage (ValueOuterOrMaxInner v) = singletonRose (fromIntegral size)
       where
-        f = fromIntegral . Map.size . snd
+        size = Map.size (Value.unpack v) `max` Value.maxInnerSize v
 
 {- Note [Costing constant-size types]
 The memory usage of each of the BLS12-381 types is constant, so we may be able

plutus-core/plutus-core/src/PlutusCore/Pretty/PrettyConst.hs

@@ -15,7 +15,7 @@ module PlutusCore.Pretty.PrettyConst where
 
 import PlutusCore.Data
 import PlutusCore.Pretty.Readable
-import PlutusCore.Value (Value (..))
+import PlutusCore.Value (Value)
 import PlutusCore.Value qualified as Value
 
 import Control.Lens hiding (List)
@@ -29,7 +29,7 @@ import Data.Typeable
 import Data.Vector.Strict (Vector)
 import Data.Word (Word8)
 import Numeric (showHex)
-import Prettyprinter as Prettyprinter
+import Prettyprinter
 import Prettyprinter.Internal (Doc (Text))
 import Text.PrettyBy
 import Text.PrettyBy.Internal (DefaultPrettyBy (..))

plutus-core/plutus-core/src/PlutusCore/Value.hs

@@ -1,11 +1,17 @@
+{-# LANGUAGE DeriveAnyClass    #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE ViewPatterns      #-}
 
 module PlutusCore.Value (
-  Value (..),
+  Value, -- Do not expose data constructor
+  NestedMap,
+  unpack,
+  pack,
   empty,
   fromList,
   toList,
   totalSize,
+  maxInnerSize,
 ) where
 
 import Codec.Serialise (Serialise)
@@ -14,33 +20,77 @@ import Data.Bifunctor
 import Data.ByteString (ByteString)
 import Data.ByteString.Base64 qualified as Base64
 import Data.Hashable (Hashable)
+import Data.IntMap.Strict (IntMap)
+import Data.IntMap.Strict qualified as IntMap
 import Data.Map.Strict (Map)
 import Data.Map.Strict qualified as Map
 import Data.Text.Encoding qualified as Text
+import GHC.Generics
 
 import PlutusPrelude (Pretty (..))
 
-newtype Value = Value {unValue :: Map ByteString (Map ByteString Integer)}
-  deriving newtype (Eq, Show, Hashable, Serialise, NFData)
+type NestedMap = Map ByteString (Map ByteString Integer)
+
+-- | The underlying type of the UPLC built-in type @Value@.
+data Value
+  = Value
+      !NestedMap
+      {- ^ Map from (currency symbol, token name) to amount.
+
+      Invariants: no empty inner map, and no zero amount.
+      -}
+      !(IntMap Int)
+      {- ^ Map from size to the number of inner maps that have that size,
+      useful for efficient retrieval of the size of the largest inner map.
+
+      Invariant: all values are positive.
+      -}
+      {-# UNPACK #-} !Int
+      -- ^ Total size, i.e., sum total of inner map sizes
+  deriving stock (Eq, Show, Generic)
+  deriving anyclass (Hashable, Serialise, NFData)
+
+{-| Unpack a `Value` into a map from (currency symbol, token name) to amount.
+
+The map is guaranteed to not contain empty inner map or zero amount.
+-}
+unpack :: Value -> NestedMap
+unpack (Value v _ _) = v
+
+{-| Pack a map from (currency symbol, token name) to amount into a `Value`.
+
+The map will be filtered so that it does not contain empty inner map or zero amount.
+-}
+pack :: NestedMap -> Value
+pack (normalize -> v) = Value v sizes size
+ where
+  sizes = Map.foldr' (IntMap.alter (maybe (Just 1) (Just . (+ 1))) . Map.size) mempty v
+  size = Map.foldr' ((+) . Map.size) 0 v
+
+{-| Total size, i.e., the number of distinct `(currency symbol, token name)` pairs
+contained in the `Value`.
+-}
+totalSize :: Value -> Int
+totalSize (Value _ _ size) = size
+
+-- | Size of the largest inner map.
+maxInnerSize :: Value -> Int
+maxInnerSize (Value _ sizes _) = maybe 0 fst (IntMap.lookupMax sizes)
 
 empty :: Value
-empty = Value mempty
+empty = Value mempty mempty 0
 
 toList :: Value -> [(ByteString, [(ByteString, Integer)])]
-toList = Map.toList . Map.map Map.toList . unValue
+toList = Map.toList . Map.map Map.toList . unpack
 
 fromList :: [(ByteString, [(ByteString, Integer)])] -> Value
 fromList =
-  normalize
-    . Value
+  pack
     . Map.fromListWith (Map.unionWith (+))
     . fmap (second (Map.fromListWith (+)))
 
-normalize :: Value -> Value
-normalize = Value . Map.filter (not . Map.null) . Map.map (Map.filter (/= 0)) . unValue
-
-totalSize :: Value -> Int
-totalSize = Map.foldl' (\n -> (n +) .  Map.size) 0 . unValue
+normalize :: NestedMap -> NestedMap
+normalize = Map.filter (not . Map.null) . Map.map (Map.filter (/= 0))
 
 instance Pretty Value where
   pretty = pretty . fmap (bimap toText (fmap (first toText))) . toList

plutus-core/testlib/PlutusCore/Generators/QuickCheck/Builtin.hs

@@ -20,12 +20,12 @@ import PlutusCore.Crypto.BLS12_381.Pairing qualified as BLS12_381.Pairing
 import PlutusCore.Data
 import PlutusCore.Generators.QuickCheck.GenerateKinds ()
 import PlutusCore.Generators.QuickCheck.Split (multiSplit0, multiSplit1, multiSplit1In)
-import PlutusCore.Value
+import PlutusCore.Value (Value)
+import PlutusCore.Value qualified as Value
 
 import Data.ByteString (ByteString, empty)
 import Data.Int
 import Data.Kind qualified as GHC
-import Data.Map (Map)
 import Data.Maybe
 import Data.Proxy
 import Data.Text (Text)
@@ -244,7 +244,9 @@ instance Arbitrary Data where
     arbitrary = arbitraryBuiltin
     shrink = shrinkBuiltin
 
-deriving via Map ByteString (Map ByteString Integer) instance Arbitrary Value
+instance Arbitrary Value where
+    arbitrary = Value.pack <$> arbitrary
+    shrink = fmap Value.pack . shrink . Value.unpack
 
 instance ArbitraryBuiltin Value
 
@@ -422,6 +424,7 @@ instance KnownKind k => Arbitrary (MaybeSomeTypeOf k) where
                , JustSomeType DefaultUniBLS12_381_G1_Element
                , JustSomeType DefaultUniBLS12_381_G2_Element
                , JustSomeType DefaultUniBLS12_381_MlResult
+               , JustSomeType DefaultUniValue
                ]
            SingType `SingKindArrow` SingType ->
                 [ genDefaultUniApply | size > 10 ]