From 0450d93a32f1e4fb9f2e50961342eb7f0e10e2b9 Mon Sep 17 00:00:00 2001
From: Nikolaos Bezirgiannis <bezirg@users.noreply.github.com>
Date: Wed, 28 Feb 2024 22:40:15 +0100
Subject: [PATCH] Added optimized SortedMap.lookup/member/delete

---
 plutus-tx/src/PlutusTx/SortedMap.hs        | 40 ++++++++++++++++++----
 plutus-tx/test/PlutusTx/SortedMap/Tests.hs | 16 ++++++++-
 2 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/plutus-tx/src/PlutusTx/SortedMap.hs b/plutus-tx/src/PlutusTx/SortedMap.hs
index 6b4384eb624..2014ddc332a 100644
--- a/plutus-tx/src/PlutusTx/SortedMap.hs
+++ b/plutus-tx/src/PlutusTx/SortedMap.hs
@@ -7,6 +7,7 @@
 {-# LANGUAGE TemplateHaskell       #-}
 {-# LANGUAGE UndecidableInstances  #-}
 {-# LANGUAGE ViewPatterns          #-}
+{-# LANGUAGE ScopedTypeVariables          #-}
 {-| This is like `PlutusTx.AssocMap` but with the extra invariant that the entries are kept
 sorted on the KEYS upon construction, insertion, and deletion.
 As with "PlutusTx.AssocMap", using the unsafe API may lead to duplicate entries, whereas using the
@@ -18,6 +19,8 @@ module PlutusTx.SortedMap
     , insert
     , delete
     , lookup
+    , member
+    , null
     , minViewWithKey
     , isSortedDeDuped
     , isValid
@@ -126,14 +129,37 @@ empty :: SortedMap k v
 empty = SortedMap AssocMap.empty
 
 {-# INLINABLE delete #-}
--- OPTIMIZE: Using AssocMap.delete works, but a version that stops earlier would be better
-delete :: (Eq k) => k -> SortedMap k v -> SortedMap k v
-delete k = fromMapUnsafe . AssocMap.delete k . toMap
+delete :: forall k v. (Ord k) => k -> SortedMap k v -> SortedMap k v
+delete c (toList -> alist) = fromListUnsafe (go alist)
+  where
+    go [] = []
+    go ms@(hd@(c', _) : rest) = case c `compare` c' of
+                              LT -> ms
+                              EQ -> rest
+                              GT -> hd : go rest
+
+{-# INLINABLE member #-}
+-- | Is the key a member of the map?
+member :: forall k v. (Ord k) => k -> SortedMap k v -> Bool
+member k m = isJust (lookup k m)
+
+-- | Is the map empty?
+null :: SortedMap k v -> Bool
+null = Tx.null . toList
 
 {-# INLINABLE lookup #-}
--- OPTIMIZE: a manual implementation can stop earlier than AssocMap.lookup, upon missing keys only
-lookup :: (Eq k) => k -> SortedMap k v -> Haskell.Maybe v
-lookup k = AssocMap.lookup k . toMap
+-- | Find an entry in a 'SortedMap'.
+lookup :: forall k v. (Tx.Ord k) => k -> SortedMap k v -> Haskell.Maybe v
+lookup c (toList -> alist) =
+  let
+    go :: [(k, v)] -> Maybe v
+    go []              = Nothing
+    go ((c', i) : xs') = case c `compare` c' of
+        LT -> Nothing
+        EQ -> Just i
+        GT -> go xs'
+   in
+    go alist
 
 {-# INLINABLE minViewWithKey #-}
 -- | Assumes that the SortedMap is valid.
@@ -154,8 +180,8 @@ insertInternal (k,v) = go
       go = \case
         [] -> [(k,v)]
         ms@((k',v'):ms') -> case k `compare` k' of
-            EQ -> (k,v) : ms'
             LT -> (k,v) : ms
+            EQ -> (k,v) : ms'
             GT -> (k',v') : go ms'
 
 {-# INLINABLE isSortedDeDuped #-}
diff --git a/plutus-tx/test/PlutusTx/SortedMap/Tests.hs b/plutus-tx/test/PlutusTx/SortedMap/Tests.hs
index d4ce2487ba0..aff57249ec8 100644
--- a/plutus-tx/test/PlutusTx/SortedMap/Tests.hs
+++ b/plutus-tx/test/PlutusTx/SortedMap/Tests.hs
@@ -40,12 +40,24 @@ mergeSortPreservesDuplicatesPoly alist =
     in before === False .&&.
        after === False
 
+deletesAllPoly :: (Tx.Ord a) => [(a, b)] -> Property
+deletesAllPoly alist =
+    let smap = SortedMap.fromListSafe alist
+    in property $ SortedMap.null $ Tx.foldr (\(k,_) acc -> SortedMap.delete k acc) smap alist
+
+membersAllPoly :: (Tx.Ord a) => [(a, b)] -> Property
+membersAllPoly alist =
+    let smap = SortedMap.fromListSafe alist
+    in property $ Tx.and $ Tx.map ((`SortedMap.member` smap) . fst) alist
+
 -- need to monomorphize to test with QC
-prop_equivSortsNoDupl, prop_insertionSortFixesValidity, prop_mergeSortPreservesDuplicates :: [(Integer, Integer)] -> Property
+prop_equivSortsNoDupl, prop_insertionSortFixesValidity, prop_mergeSortPreservesDuplicates, prop_deletesAll, prop_membersAll :: [(Integer, Integer)] -> Property
 
 prop_equivSortsNoDupl = equivSortsNoDuplPoly
 prop_insertionSortFixesValidity = insertionSortFixesValidityPoly
 prop_mergeSortPreservesDuplicates = mergeSortPreservesDuplicatesPoly
+prop_deletesAll = deletesAllPoly
+prop_membersAll = membersAllPoly
 
 propertyTests :: TestTree
 propertyTests =
@@ -53,4 +65,6 @@ propertyTests =
         [ testProperty "merge-sort is equiv to insertion-sort for non-duplicated maps" prop_equivSortsNoDupl
         , testProperty "insertion-sort turns an invalid assocmap to a valid one" prop_insertionSortFixesValidity
         , testProperty "merge-sort preserves the (in)validity" prop_mergeSortPreservesDuplicates
+        , testProperty "folding delete over all elements should make the Map empty" prop_deletesAll
+        , testProperty "folding member over all elements succeeds" prop_membersAll
         ]