Added optimized SortedMap.lookup/member/delete

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 #-}

plutus-tx/test/PlutusTx/SortedMap/Tests.hs

@@ -40,17 +40,31 @@ 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 =
     testGroup "SortedMap"
         [ 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
         ]