Add insertCoin and unionValue implementations and tests (#7322)

Author: zliu41

plutus-core/changelog.d/20250909_084412_unsafeFixIO_insertCoin_unionValue.md

@@ -0,0 +1,4 @@
+
+### Added
+
+- Built-in type `Value`, and the implementations of `insertCoin` and `unionValue`.

plutus-core/plutus-core.cabal

@@ -378,6 +378,7 @@ test-suite plutus-core-test
     Parser.Spec
     Pretty.Readable
     TypeSynthesis.Spec
+    Value.Spec
 
   default-language: Haskell2010
   build-depends:
@@ -396,6 +397,8 @@ test-suite plutus-core-test
     , plutus-core                      ^>=1.53
     , plutus-core:plutus-core-testlib
     , prettyprinter
+    , QuickCheck
+    , safe
     , serialise
     , tasty
     , tasty-golden
@@ -798,6 +801,7 @@ library plutus-core-testlib
     PlutusCore.Generators.QuickCheck.Substitutions
     PlutusCore.Generators.QuickCheck.Unification
     PlutusCore.Generators.QuickCheck.Utils
+    PlutusCore.Generators.QuickCheck.Value
     PlutusCore.Test
     PlutusIR.Generators.AST
     PlutusIR.Generators.QuickCheck
@@ -813,6 +817,7 @@ library plutus-core-testlib
 
   build-depends:
     , base                        >=4.9     && <5
+    , base16-bytestring
     , bifunctors
     , bytestring
     , containers
@@ -836,6 +841,7 @@ library plutus-core-testlib
     , QuickCheck
     , quickcheck-instances
     , quickcheck-transformer
+    , random
     , size-based
     , Stream
     , tagged

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

@@ -1,5 +1,6 @@
 {-# LANGUAGE DeriveAnyClass    #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE LambdaCase        #-}
 {-# LANGUAGE ViewPatterns      #-}
 
 module PlutusCore.Value (
@@ -10,8 +11,12 @@ module PlutusCore.Value (
   empty,
   fromList,
   toList,
+  toFlatList,
   totalSize,
   maxInnerSize,
+  insertCoin,
+  deleteCoin,
+  unionValue,
 ) where
 
 import Codec.Serialise (Serialise)
@@ -22,8 +27,10 @@ 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.Merge.Strict qualified as M
 import Data.Map.Strict (Map)
 import Data.Map.Strict qualified as Map
+import Data.Maybe
 import Data.Text.Encoding qualified as Text
 import GHC.Generics
 
@@ -56,32 +63,50 @@ The map is guaranteed to not contain empty inner map or zero amount.
 -}
 unpack :: Value -> NestedMap
 unpack (Value v _ _) = v
+{-# INLINE unpack #-}
 
 {-| 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
+pack = pack' . normalize
+{-# INLINE pack #-}
+
+-- | Like `pack` but does not normalize.
+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
+  (sizes, size) = Map.foldl' alg (mempty, 0) v
+  alg (ss, s) inner =
+    ( IntMap.alter (maybe (Just 1) (Just . (+ 1))) (Map.size inner) ss
+    , s + Map.size inner
+    )
+{-# INLINEABLE pack' #-}
 
 {-| Total size, i.e., the number of distinct `(currency symbol, token name)` pairs
 contained in the `Value`.
 -}
 totalSize :: Value -> Int
 totalSize (Value _ _ size) = size
+{-# INLINE totalSize #-}
 
 -- | Size of the largest inner map.
 maxInnerSize :: Value -> Int
 maxInnerSize (Value _ sizes _) = maybe 0 fst (IntMap.lookupMax sizes)
+{-# INLINE maxInnerSize #-}
 
 empty :: Value
 empty = Value mempty mempty 0
+{-# INLINE empty #-}
 
 toList :: Value -> [(ByteString, [(ByteString, Integer)])]
 toList = Map.toList . Map.map Map.toList . unpack
+{-# INLINEABLE toList #-}
+
+toFlatList :: Value -> [(ByteString, ByteString, Integer)]
+toFlatList (toList -> xs) = [(c, t, a) | (c, ys) <- xs, (t, a) <- ys]
+{-# INLINEABLE toFlatList #-}
 
 fromList :: [(ByteString, [(ByteString, Integer)])] -> Value
 fromList =
@@ -91,8 +116,85 @@ fromList =
 
 normalize :: NestedMap -> NestedMap
 normalize = Map.filter (not . Map.null) . Map.map (Map.filter (/= 0))
+{-# INLINEABLE normalize #-}
 
 instance Pretty Value where
   pretty = pretty . fmap (bimap toText (fmap (first toText))) . toList
    where
     toText = Text.decodeLatin1 . Base64.encode
+
+{-| \(O(\log \max(m, k))\), where \(m\) is the size of the outer map, and \(k\) is
+the size of the largest inner map.
+-}
+insertCoin :: ByteString -> ByteString -> Integer -> Value -> Value
+insertCoin currency token amt v@(Value outer sizes size)
+  | amt == 0 = deleteCoin currency token v
+  | otherwise =
+      let (mold, outer') = Map.alterF f currency outer
+          (sizes', size') = case mold of
+            Just old -> (updateSizes old (old + 1) sizes, size + 1)
+            Nothing  -> (sizes, size)
+       in Value outer' sizes' size'
+ where
+  f
+    :: Maybe (Map ByteString Integer)
+    -> ( Maybe Int -- Just (old size of inner map) if the total size grows by 1, otherwise Nothing
+       , Maybe (Map ByteString Integer)
+       )
+  f = \case
+    Nothing -> (Just 0, Just (Map.singleton token amt))
+    Just inner ->
+      let (isJust -> exists, inner') = Map.insertLookupWithKey (\_ _ _ -> amt) token amt inner
+       in (if exists then Nothing else Just (Map.size inner), Just inner')
+{-# INLINEABLE insertCoin #-}
+
+-- TODO: implement properly
+deleteCoin :: ByteString -> ByteString -> Value -> Value
+deleteCoin currency token (Value outer _ _) =
+  pack $ case Map.lookup currency outer of
+    Nothing    -> outer
+    Just inner -> Map.insert currency (Map.delete token inner) outer
+
+{-| The precise complexity is complicated, but an upper bound
+is \(O(n_{1} \log n_{2}) + O(m)\), where \(n_{1}\) is the total size of the smaller
+value, \(n_{2}\) is the total size of the bigger value, and \(m\) is the
+combined size of the outer maps.
+-}
+unionValue :: Value -> Value -> Value
+unionValue (unpack -> vA) (unpack -> vB) =
+  pack' $
+    M.merge
+      M.preserveMissing
+      M.preserveMissing
+      ( M.zipWithMaybeMatched $ \_ innerA innerB ->
+          let inner =
+                M.merge
+                  M.preserveMissing
+                  M.preserveMissing
+                  ( M.zipWithMaybeMatched $ \_ x y ->
+                      let z = x + y in if z == 0 then Nothing else Just z
+                  )
+                  innerA
+                  innerB
+           in if Map.null inner
+                then Nothing
+                else
+                  Just inner
+      )
+      vA
+      vB
+{-# INLINEABLE unionValue #-}
+
+-- | Decrement bucket @old@, and increment bucket @new@.
+updateSizes :: Int -> Int -> IntMap Int -> IntMap Int
+updateSizes old new = dec . inc
+ where
+  inc =
+    if new == 0
+      then id
+      else IntMap.alter (maybe (Just 1) (Just . (+ 1))) new
+  dec =
+    if old == 0
+      then id
+      else IntMap.update (\n -> if n <= 1 then Nothing else Just (n - 1)) old
+{-# INLINEABLE updateSizes #-}

plutus-core/plutus-core/test/Spec.hs

@@ -23,6 +23,7 @@ import Normalization.Type
 import Parser.Spec qualified as Parser
 import Pretty.Readable
 import TypeSynthesis.Spec (test_typecheck)
+import Value.Spec qualified as Value
 
 import PlutusCore
 import PlutusCore.Check.Uniques qualified as Uniques
@@ -257,5 +258,6 @@ allTests plcFiles rwFiles typeFiles typeErrorFiles =
     , Check.tests
     , NEAT.tests
     , Parser.tests
+    , Value.tests
     , test_utils
     ]

plutus-core/plutus-core/test/Value/Spec.hs

@@ -0,0 +1,107 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ViewPatterns      #-}
+{-# OPTIONS_GHC -fno-warn-orphans #-}
+
+module Value.Spec (tests) where
+
+import Data.Foldable qualified as F
+import Data.Map.Strict qualified as Map
+import Data.Maybe
+import PlutusCore.Generators.QuickCheck.Builtin ()
+import PlutusCore.Generators.QuickCheck.Value (genShortHex)
+import PlutusCore.Value (Value)
+import PlutusCore.Value qualified as V
+import Safe.Foldable (maximumMay)
+import Test.QuickCheck
+import Test.Tasty
+import Test.Tasty.QuickCheck
+
+prop_packUnpackRoundtrip :: Property
+prop_packUnpackRoundtrip = forAll arbitrary $ \v ->
+  v === V.pack (V.unpack v)
+
+-- | Verifies that @pack@ correctly updates the sizes
+prop_packBookkeeping :: Property
+prop_packBookkeeping = forAll arbitrary $ \nm ->
+  checkSizes (V.pack nm)
+
+{-| Verifies that @pack@ preserves @Value@ invariants, i.e.,
+no empty inner map or zero amount.
+-}
+prop_packPreservesInvariants :: Property
+prop_packPreservesInvariants = forAll arbitrary $ \nm ->
+  checkInvariants (V.pack nm)
+
+-- | Verifies that @insertCoin@ correctly updates the sizes
+prop_insertCoinBookkeeping :: Property
+prop_insertCoinBookkeeping = forAll arbitrary $ \(v, amt) ->
+  forAll (genShortHex (V.totalSize v)) $ \currency ->
+    forAll (genShortHex (V.totalSize v)) $ \token ->
+      let v' = V.insertCoin currency token amt v
+       in checkSizes v'
+
+-- | Verifies that @insertCoin@ preserves @Value@ invariants
+prop_insertCoinPreservesInvariants :: Property
+prop_insertCoinPreservesInvariants = forAll arbitrary $ \(v, amt) ->
+  forAll (genShortHex (V.totalSize v)) $ \currency ->
+    forAll (genShortHex (V.totalSize v)) $ \token ->
+      let v' = V.insertCoin currency token amt v
+       in checkInvariants v'
+
+prop_unionCommutative :: Property
+prop_unionCommutative = forAll arbitrary $ \(v, v') ->
+  V.unionValue v v' === V.unionValue v' v
+
+prop_unionAssociative :: Property
+prop_unionAssociative = forAll arbitrary $ \(v1, v2, v3) ->
+  V.unionValue v1 (V.unionValue v2 v3) === V.unionValue (V.unionValue v1 v2) v3
+
+prop_insertCoinIdempotent :: Property
+prop_insertCoinIdempotent = forAll arbitrary $ \v ->
+  let fm = V.toFlatList v
+   in v === F.foldl' (\acc (c, t, a) -> V.insertCoin c t a acc) v fm
+
+checkSizes :: Value -> Property
+checkSizes v =
+  (expectedMaxInnerSize === actualMaxInnerSize)
+    .&&. (expectedSize === actualSize)
+ where
+  expectedMaxInnerSize = fromMaybe 0 . maximumMay $ Map.map Map.size (V.unpack v)
+  actualMaxInnerSize = V.maxInnerSize v
+  expectedSize = sum $ Map.map Map.size (V.unpack v)
+  actualSize = V.totalSize v
+
+checkInvariants :: Value -> Property
+checkInvariants (V.unpack -> v) =
+  property ((not . any Map.null) v)
+    .&&. property ((not . any (elem 0)) v)
+
+tests :: TestTree
+tests =
+  testGroup
+    "Value"
+    [ testProperty
+        "packUnpackRoundtrip"
+        prop_packUnpackRoundtrip
+    , testProperty
+        "packBookkeeping"
+        prop_packBookkeeping
+    , testProperty
+        "packPreservesInvariants"
+        prop_packPreservesInvariants
+    , testProperty
+        "insertCoinBookkeeping"
+        prop_insertCoinBookkeeping
+    , testProperty
+        "insertCoinPreservesInvariants"
+        prop_insertCoinPreservesInvariants
+    , testProperty
+        "unionCommutative"
+        prop_unionCommutative
+    , testProperty
+        "unionAssociative"
+        prop_unionAssociative
+    , testProperty
+        "insertCoinIdempotent"
+        prop_insertCoinIdempotent
+    ]

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

@@ -21,8 +21,8 @@ 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.Generators.QuickCheck.Value ()
 import PlutusCore.Value (Value)
-import PlutusCore.Value qualified as Value
 
 import Data.ByteString (ByteString, empty)
 import Data.Int
@@ -245,10 +245,6 @@ instance Arbitrary Data where
     arbitrary = arbitraryBuiltin
     shrink = shrinkBuiltin
 
-instance Arbitrary Value where
-    arbitrary = Value.pack <$> arbitrary
-    shrink = fmap Value.pack . shrink . Value.unpack
-
 instance ArbitraryBuiltin Value
 
 instance ArbitraryBuiltin BLS12_381.G1.Element where