From 5879f8e9fc68a444b6ab8645759916060090a593 Mon Sep 17 00:00:00 2001
From: Ana Pantilie <45069775+ana-pantilie@users.noreply.github.com>
Date: Tue, 19 Mar 2024 19:49:46 +0200
Subject: [PATCH] Document unsafe operations of `AssocMap` (#5838)

---
 .../Marlowe/Core/V1/Semantics.hs              |  6 +-
 .../src/PlutusLedgerApi/V1/Value.hs           | 19 ++++-
 plutus-ledger-api/src/PlutusLedgerApi/V2.hs   |  4 +-
 plutus-ledger-api/src/PlutusLedgerApi/V3.hs   |  2 +-
 plutus-ledger-api/test-plugin/Spec/Budget.hs  |  3 +-
 plutus-ledger-api/test-plugin/Spec/Value.hs   |  2 +-
 .../testlib/PlutusLedgerApi/Test/V1/Value.hs  |  2 +-
 ..._ana.pantilie95_plt_9511_audit_assocmap.md |  8 ++
 plutus-tx/src/PlutusTx/AssocMap.hs            | 76 ++++++++++++++-----
 9 files changed, 91 insertions(+), 31 deletions(-)
 create mode 100644 plutus-tx/changelog.d/20240315_151519_ana.pantilie95_plt_9511_audit_assocmap.md

diff --git a/plutus-benchmark/marlowe/src/PlutusBenchmark/Marlowe/Core/V1/Semantics.hs b/plutus-benchmark/marlowe/src/PlutusBenchmark/Marlowe/Core/V1/Semantics.hs
index 221603c2b35..918462e4a5d 100644
--- a/plutus-benchmark/marlowe/src/PlutusBenchmark/Marlowe/Core/V1/Semantics.hs
+++ b/plutus-benchmark/marlowe/src/PlutusBenchmark/Marlowe/Core/V1/Semantics.hs
@@ -328,12 +328,12 @@ refundOne accounts = case Map.toList accounts of
     -- Isabelle semantics in that it returns the least-recently
     -- added account-token combination rather than the first
     -- lexicographically ordered one. Also, the sequence
-    -- `Map.fromList . tail . Map.toList` preserves the
+    -- `Map.unsafeFromList . tail . Map.toList` preserves the
     -- invariants of order and non-duplication.
     ((accId, token), balance) : rest ->
         if balance > 0
-        then Just ((accId, token, balance), Map.fromList rest)
-        else refundOne (Map.fromList rest)
+        then Just ((accId, token, balance), Map.unsafeFromList rest)
+        else refundOne (Map.unsafeFromList rest)
 
 
 -- | Obtains the amount of money available an account.
diff --git a/plutus-ledger-api/src/PlutusLedgerApi/V1/Value.hs b/plutus-ledger-api/src/PlutusLedgerApi/V1/Value.hs
index 014004a78ad..99eace37ae5 100644
--- a/plutus-ledger-api/src/PlutusLedgerApi/V1/Value.hs
+++ b/plutus-ledger-api/src/PlutusLedgerApi/V1/Value.hs
@@ -262,6 +262,7 @@ instance MeetSemiLattice Value where
 
 {-# INLINABLE valueOf #-}
 -- | Get the quantity of the given currency in the 'Value'.
+-- Assumes that the underlying map doesn't contain duplicate keys.
 valueOf :: Value -> CurrencySymbol -> TokenName -> Integer
 valueOf (Value mp) cur tn =
     case Map.lookup cur mp of
@@ -271,6 +272,8 @@ valueOf (Value mp) cur tn =
             Just v  -> v
 
 {-# INLINABLE currencySymbolValueOf #-}
+-- | Get the total value of the currency symbol in the 'Value' map.
+-- Assumes that the underlying map doesn't contain duplicate keys.
 currencySymbolValueOf :: Value -> CurrencySymbol -> Integer
 currencySymbolValueOf (Value mp) cur = case Map.lookup cur mp of
     Nothing     -> 0
@@ -290,8 +293,8 @@ singleton :: CurrencySymbol -> TokenName -> Integer -> Value
 singleton c tn i = Value (Map.singleton c (Map.singleton tn i))
 
 {-# INLINABLE lovelaceValue #-}
-lovelaceValue :: Lovelace -> Value
 -- | A 'Value' containing the given quantity of Lovelace.
+lovelaceValue :: Lovelace -> Value
 lovelaceValue = singleton adaSymbol adaToken . getLovelace
 
 {-# INLINABLE lovelaceValueOf #-}
@@ -310,7 +313,7 @@ assetClassValueOf :: Value -> AssetClass -> Integer
 assetClassValueOf v (AssetClass (c, t)) = valueOf v c t
 
 {-# INLINABLE unionVal #-}
--- | Combine two 'Value' maps
+-- | Combine two 'Value' maps, assumes the well-definedness of the two maps.
 unionVal :: Value -> Value -> Map.Map CurrencySymbol (Map.Map TokenName (These Integer Integer))
 unionVal (Value l) (Value r) =
     let
@@ -322,6 +325,8 @@ unionVal (Value l) (Value r) =
     in unThese <$> combined
 
 {-# INLINABLE unionWith #-}
+-- | Combine two 'Value' maps with the argument function.
+-- Assumes the well-definedness of the two maps.
 unionWith :: (Integer -> Integer -> Integer) -> Value -> Value -> Value
 unionWith f ls rs =
     let
@@ -336,6 +341,7 @@ unionWith f ls rs =
 -- | Convert a 'Value' to a simple list, keeping only the non-zero amounts.
 -- Note that the result isn't sorted, meaning @v1 == v2@ doesn't generally imply
 -- @flattenValue v1 == flattenValue v2@.
+-- Also assumes that there are no duplicate keys in the 'Value' 'Map'.
 flattenValue :: Value -> [(CurrencySymbol, TokenName, Integer)]
 flattenValue v = goOuter [] (Map.toList $ getValue v)
   where
@@ -355,6 +361,8 @@ isZero :: Value -> Bool
 isZero (Value xs) = Map.all (Map.all (\i -> 0 == i)) xs
 
 {-# INLINABLE checkPred #-}
+-- | Checks whether a predicate holds for all the values in a 'Value'
+-- union. Assumes the well-definedness of the two underlying 'Map's.
 checkPred :: (These Integer Integer -> Bool) -> Value -> Value -> Bool
 checkPred f l r =
     let
@@ -417,8 +425,11 @@ split (Value mp) = (negate (Value neg), Value pos) where
     (l, r) = Map.mapThese (\i -> if i <= 0 then This i else That i) mp'
 
 {-# INLINABLE unordEqWith #-}
-{- | Check equality of two lists given a function checking whether a 'Value' is zero and a function
-checking equality of values.
+{- | Check equality of two lists of distinct key-value pairs, each value being uniquely
+identified by a key, given a function checking whether a 'Value' is zero and a function
+checking equality of values. Note that the caller must ensure that the two lists are
+well-defined in this sense. This is not checked or enforced in `unordEqWith`, and therefore
+it might yield undefined results for ill-defined input.
 
 This function recurses on both the lists in parallel and checks whether the key-value pairs are
 equal pointwise. If there is a mismatch, then it tries to find the left key-value pair in the right
diff --git a/plutus-ledger-api/src/PlutusLedgerApi/V2.hs b/plutus-ledger-api/src/PlutusLedgerApi/V2.hs
index 6cd8f5d998c..ab17125c809 100644
--- a/plutus-ledger-api/src/PlutusLedgerApi/V2.hs
+++ b/plutus-ledger-api/src/PlutusLedgerApi/V2.hs
@@ -100,7 +100,7 @@ module PlutusLedgerApi.V2 (
 
   -- *** Association maps
   Map,
-  fromList,
+  unsafeFromList,
 
   -- *** Newtypes and hash types
   ScriptHash (..),
@@ -138,7 +138,7 @@ import PlutusLedgerApi.V2.ParamName
 import PlutusLedgerApi.V2.Tx (OutputDatum (..))
 
 import PlutusCore.Data qualified as PLC
-import PlutusTx.AssocMap (Map, fromList)
+import PlutusTx.AssocMap (Map, unsafeFromList)
 
 import Control.Monad.Except (MonadError)
 
diff --git a/plutus-ledger-api/src/PlutusLedgerApi/V3.hs b/plutus-ledger-api/src/PlutusLedgerApi/V3.hs
index 936834c23c3..d6249a9e8d9 100644
--- a/plutus-ledger-api/src/PlutusLedgerApi/V3.hs
+++ b/plutus-ledger-api/src/PlutusLedgerApi/V3.hs
@@ -118,7 +118,7 @@ module PlutusLedgerApi.V3 (
 
   -- *** Association maps
   V2.Map,
-  V2.fromList,
+  V2.unsafeFromList,
 
   -- *** Newtypes and hash types
   V2.ScriptHash (..),
diff --git a/plutus-ledger-api/test-plugin/Spec/Budget.hs b/plutus-ledger-api/test-plugin/Spec/Budget.hs
index ec22c9b461c..f2393bf912f 100644
--- a/plutus-ledger-api/test-plugin/Spec/Budget.hs
+++ b/plutus-ledger-api/test-plugin/Spec/Budget.hs
@@ -50,7 +50,8 @@ compiledCurrencySymbolValueOf :: CompiledCode (Value -> CurrencySymbol -> Intege
 compiledCurrencySymbolValueOf = $$(compile [||currencySymbolValueOf||])
 
 mkValue :: [(Integer, [(Integer, Integer)])] -> Value
-mkValue = Value . Map.fromList . fmap (bimap toSymbol (Map.fromList . fmap (first toToken)))
+mkValue =
+    Value . Map.unsafeFromList . fmap (bimap toSymbol (Map.unsafeFromList . fmap (first toToken)))
 
 toSymbol :: Integer -> CurrencySymbol
 toSymbol = currencySymbol . fromString . show
diff --git a/plutus-ledger-api/test-plugin/Spec/Value.hs b/plutus-ledger-api/test-plugin/Spec/Value.hs
index b2b920793a9..87ad1c6ec24 100644
--- a/plutus-ledger-api/test-plugin/Spec/Value.hs
+++ b/plutus-ledger-api/test-plugin/Spec/Value.hs
@@ -147,7 +147,7 @@ currencyLongListOptions =
     ListTx.concatMap (maybe longCurrencyChunk pure) currencyListWithHooks
 
 listsToValue :: [(CurrencySymbol, [(TokenName, Integer)])] -> Value
-listsToValue = Value . AssocMap.fromList . ListTx.map (fmap AssocMap.fromList)
+listsToValue = Value . AssocMap.unsafeFromList . ListTx.map (fmap AssocMap.unsafeFromList)
 
 valueToLists :: Value -> [(CurrencySymbol, [(TokenName, Integer)])]
 valueToLists = ListTx.map (fmap AssocMap.toList) . AssocMap.toList . getValue
diff --git a/plutus-ledger-api/testlib/PlutusLedgerApi/Test/V1/Value.hs b/plutus-ledger-api/testlib/PlutusLedgerApi/Test/V1/Value.hs
index 3386e38947b..90afa8a1ee4 100644
--- a/plutus-ledger-api/testlib/PlutusLedgerApi/Test/V1/Value.hs
+++ b/plutus-ledger-api/testlib/PlutusLedgerApi/Test/V1/Value.hs
@@ -18,7 +18,7 @@ import Test.QuickCheck
 
 -- | Convert a list representation of a 'Value' to the 'Value'.
 listsToValue :: [(CurrencySymbol, [(TokenName, Integer)])] -> Value
-listsToValue = Value . AssocMap.fromList . ListTx.map (fmap AssocMap.fromList)
+listsToValue = Value . AssocMap.unsafeFromList . ListTx.map (fmap AssocMap.unsafeFromList)
 
 -- | Convert a 'Value' to its list representation.
 valueToLists :: Value -> [(CurrencySymbol, [(TokenName, Integer)])]
diff --git a/plutus-tx/changelog.d/20240315_151519_ana.pantilie95_plt_9511_audit_assocmap.md b/plutus-tx/changelog.d/20240315_151519_ana.pantilie95_plt_9511_audit_assocmap.md
new file mode 100644
index 00000000000..37a9d4d0a9f
--- /dev/null
+++ b/plutus-tx/changelog.d/20240315_151519_ana.pantilie95_plt_9511_audit_assocmap.md
@@ -0,0 +1,8 @@
+### Added
+
+- Documented functions which unsafely construct `PlutusTx.AssocMap.Map`s, or depend on the precondition that the input `Map`s do not contain duplicate entries.
+
+### Changed
+
+- Renamed `PlutusTx.AssocMap.Map.fromList` to `PlutusTx.AssocMap.Map.unsafeFromList`.
+- Renamed `PlutusTx.AssocMap.Map.fromListSafe` to `PlutusTx.AssocMap.Map.safeFromList`.
diff --git a/plutus-tx/src/PlutusTx/AssocMap.hs b/plutus-tx/src/PlutusTx/AssocMap.hs
index 2053038e206..12cfce5684c 100644
--- a/plutus-tx/src/PlutusTx/AssocMap.hs
+++ b/plutus-tx/src/PlutusTx/AssocMap.hs
@@ -18,8 +18,8 @@ module PlutusTx.AssocMap (
   singleton,
   empty,
   null,
-  fromList,
-  fromListSafe,
+  unsafeFromList,
+  safeFromList,
   toList,
   keys,
   elems,
@@ -53,15 +53,27 @@ import GHC.Generics (Generic)
 import Language.Haskell.TH.Syntax as TH (Lift)
 import Prettyprinter (Pretty (..))
 
-{- HLINT ignore "Use newtype instead of data" -}
-
 -- See Note [Optimising Value].
 -- | A 'Map' of key-value pairs.
+-- A 'Map' is considered well-defined if there are no key collisions, meaning that each value
+-- is uniquely identified by a key.
+--
+-- Use 'safeFromList' to create well-defined 'Map's from arbitrary lists of pairs.
+--
+-- If cost minimisation is required, then you can use 'unsafeFromList' but you must
+-- be certain that the list you are converting to a 'Map' abides by the well-definedness condition.
+--
+-- Most operations on 'Map's are definedness-preserving, meaning that for the resulting 'Map' to be
+-- well-defined then the input 'Map'(s) have to also be well-defined. This is not checked explicitly
+-- unless mentioned in the documentation.
+--
+-- Take care when using 'fromBuiltinData' and 'unsafeFromBuiltinData', as neither function performs
+-- deduplication of the input collection and may create invalid 'Map's!
 newtype Map k v = Map {unMap :: [(k, v)]}
   deriving stock (Generic, Haskell.Eq, Haskell.Show, Data, TH.Lift)
   deriving newtype (Eq, Ord, NFData)
 
--- Hand-written instances to use the underlying 'Map' type in 'Data', and
+-- | Hand-written instances to use the underlying 'Map' type in 'Data', and
 -- to be reasonably efficient.
 instance (ToData k, ToData v) => ToData (Map k v) where
   toBuiltinData (Map es) = BI.mkMap (mapToBuiltin es)
@@ -74,6 +86,11 @@ instance (ToData k, ToData v) => ToData (Map k v) where
           go []            = BI.mkNilPairData BI.unitval
           go ((k, v) : xs) = BI.mkCons (BI.mkPairData (toBuiltinData k) (toBuiltinData v)) (go xs)
 
+-- | A hand-written transformation from 'Data' to 'Map'. Compared to 'unsafeFromBuiltinData',
+-- it is safe to call when it is unknown if the 'Data' is built with 'Data's 'Map' constructor.
+-- Note that it is, however, unsafe in the sense that it assumes that any map
+-- encoded in the 'Data' is well-formed, i.e. 'fromBuiltinData' does not perform any
+-- deduplication of keys or of key-value pairs!
 instance (FromData k, FromData v) => FromData (Map k v) where
   fromBuiltinData d =
     P.matchData'
@@ -104,6 +121,12 @@ instance (FromData k, FromData v) => FromData (Map k v) where
               )
               ()
 
+-- | A hand-written transformation from 'Data' to 'Map'. It is unsafe because the
+-- caller must provide the guarantee that the 'Data' is constructed using the 'Data's
+-- 'Map' constructor.
+-- Note that it assumes, like the 'fromBuiltinData' transformation, that the map
+-- encoded in the 'Data' is well-formed, i.e. 'unsafeFromBuiltinData' does not perform
+-- any deduplication of keys or of key-value pairs!
 instance (UnsafeFromData k, UnsafeFromData v) => UnsafeFromData (Map k v) where
   -- The `~` here enables `BI.unsafeDataAsMap d` to be inlined, which reduces costs slightly.
   -- Without the `~`, the inliner would consider it not effect safe to inline.
@@ -151,23 +174,27 @@ instance (Eq k, Semigroup v) => Monoid (Map k v) where
 instance (Pretty k, Pretty v) => Pretty (Map k v) where
   pretty (Map mp) = pretty mp
 
-{-# INLINEABLE fromList #-}
-fromList :: [(k, v)] -> Map k v
-fromList = Map
+{-# INLINEABLE unsafeFromList #-}
+-- | Unsafely create a 'Map' from a list of pairs. This should _only_ be applied to lists which
+-- have been checked to not contain duplicate keys, otherwise the resulting 'Map' will contain
+-- conflicting entries (two entries sharing the same key).
+-- As usual, the "keys" are considered to be the first element of the pair.
+unsafeFromList :: [(k, v)] -> Map k v
+unsafeFromList = Map
 
-{-# INLINEABLE fromListSafe #-}
+{-# INLINEABLE safeFromList #-}
 -- | In case of duplicates, this function will keep only one entry (the one that precedes).
 -- In other words, this function de-duplicates the input list.
-fromListSafe :: Eq k => [(k, v)] -> Map k v
-fromListSafe = foldr (uncurry insert) empty
+safeFromList :: Eq k => [(k, v)] -> Map k v
+safeFromList = foldr (uncurry insert) empty
 
 {-# INLINEABLE toList #-}
 toList :: Map k v -> [(k, v)]
 toList (Map l) = l
 
 {-# INLINEABLE lookup #-}
-
--- | Find an entry in a 'Map'.
+-- | Find an entry in a 'Map'. If the 'Map' is not well-formed (it contains duplicate keys)
+-- then this will return the value of the left-most pair in the underlying list of pairs.
 lookup :: forall k v. (Eq k) => k -> Map k v -> Maybe v
 lookup c (Map xs) =
   let
@@ -191,6 +218,9 @@ insert k v (Map xs) = Map (go xs)
     go ((k', v') : rest) = if k == k' then (k, v) : rest else (k', v') : go rest
 
 {-# INLINEABLE delete #-}
+-- | Delete an entry from the 'Map'. Assumes that the 'Map' is well-formed, i.e. if the
+-- underlying list of pairs contains pairs with duplicate keys then only the left-most
+-- pair will be removed.
 delete :: forall k v. (Eq k) => k -> Map k v -> Map k v
 delete key (Map ls) = Map (go ls)
   where
@@ -200,11 +230,17 @@ delete key (Map ls) = Map (go ls)
       | otherwise = (k, v) : go rest
 
 {-# INLINEABLE keys #-}
--- | The keys of a 'Map'.
+-- | The keys of a 'Map'. Semantically, the resulting list is only a set if the 'Map'
+-- didn't contain duplicate keys.
 keys :: Map k v -> [k]
 keys (Map xs) = P.fmap (\(k, _ :: v) -> k) xs
 
--- | Combine two 'Map's.
+-- | Combine two 'Map's. Keeps both values on key collisions.
+-- Note that well-formedness is only preserved if the two input maps
+-- are also well-formed.
+-- Also, as an implementation detail, in the case that the right map contains
+-- duplicate keys, and there exists a collision between the two maps,
+-- then only the left-most value of the right map will be kept.
 union :: forall k v r. (Eq k) => Map k v -> Map k r -> Map k (These v r)
 union (Map ls) (Map rs) =
   let
@@ -216,6 +252,7 @@ union (Map ls) (Map rs) =
     ls' :: [(k, These v r)]
     ls' = P.fmap (\(c, i) -> (c, f i (lookup c (Map rs)))) ls
 
+    -- Keeps only those keys which don't appear in the left map.
     rs' :: [(k, r)]
     rs' = P.filter (\(c, _) -> not (any (\(c', _) -> c' == c) ls)) rs
 
@@ -225,8 +262,12 @@ union (Map ls) (Map rs) =
     Map (ls' ++ rs'')
 
 {-# INLINEABLE unionWith #-}
-
 -- | Combine two 'Map's with the given combination function.
+-- Note that well-formedness of the resulting map depends on the two input maps
+-- being well-formed.
+-- Also, as an implementation detail, in the case that the right map contains
+-- duplicate keys, and there exists a collision between the two maps,
+-- then only the left-most value of the right map will be kept.
 unionWith :: forall k a. (Eq k) => (a -> a -> a) -> Map k a -> Map k a -> Map k a
 unionWith merge (Map ls) (Map rs) =
   let
@@ -244,7 +285,6 @@ unionWith merge (Map ls) (Map rs) =
     Map (ls' ++ rs')
 
 {-# INLINEABLE mapThese #-}
-
 -- | A version of 'Data.Map.Lazy.mapEither' that works with 'These'.
 mapThese :: (v -> These a b) -> Map k v -> (Map k a, Map k b)
 mapThese f mps = (Map mpl, Map mpr)
@@ -280,7 +320,7 @@ filter f (Map m) = Map $ P.filter (f . snd) m
 
 {-# INLINEABLE elems #-}
 
--- | Return all elements of the map in the ascending order of their keys.
+-- | Return all elements of the map.
 elems :: Map k v -> [v]
 elems (Map xs) = P.fmap (\(_ :: k, v) -> v) xs