Sets enhancements, fixes #2467 (#6158)

1 files changed, 94 insertions, 4 deletions

diff --git a/lib/pure/collections/sets.nim b/lib/pure/collections/sets.nim
index d51a5c388..dbdf17514 100644
--- a/lib/pure/collections/sets.nim
+++ b/lib/pure/collections/sets.nim
@@ -287,8 +287,6 @@ proc exclImpl[A](s: var HashSet[A], key: A) : bool {. inline .} =
 
   if i >= 0:
     result = false
-    s.data[i].hcode = 0
-    s.data[i].key = default(type(s.data[i].key))
     dec(s.counter)
     while true:         # KnuthV3 Algo6.4R adapted for i=i+1 instead of i=i-1
       var j = i         # The correctness of this depends on (h+1) in nextTry,
@@ -300,7 +298,7 @@ proc exclImpl[A](s: var HashSet[A], key: A) : bool {. inline .} =
         if isEmpty(s.data[i].hcode):   # end of collision cluster; So all done
           return
         r = s.data[i].hcode and msk    # "home" location of key@i
-      shallowCopy(s.data[j], s.data[i]) # data[j] will be marked EMPTY next loop
+      shallowCopy(s.data[j], s.data[i]) # data[i] will be marked EMPTY next loop
 
 proc missingOrExcl*[A](s: var HashSet[A], key: A): bool =
   ## Excludes `key` in the set `s` and tells if `key` was removed from `s`.
@@ -662,9 +660,12 @@ proc card*[A](s: OrderedSet[A]): int {.inline.} =
 
 template forAllOrderedPairs(yieldStmt: untyped) {.dirty.} =
   var h = s.first
+  var idx = 0
   while h >= 0:
     var nxt = s.data[h].next
-    if isFilled(s.data[h].hcode): yieldStmt
+    if isFilled(s.data[h].hcode):
+      yieldStmt
+      inc(idx)
     h = nxt
 
 iterator items*[A](s: OrderedSet[A]): A =
@@ -689,6 +690,11 @@ iterator items*[A](s: OrderedSet[A]): A =
   forAllOrderedPairs:
     yield s.data[h].key
 
+iterator pairs*[A](s: OrderedSet[A]): tuple[a: int, b: A] =
+  assert s.isValid, "The set needs to be initialized"
+  forAllOrderedPairs:
+    yield (idx, s.data[h].key)
+
 proc rawGetKnownHC[A](s: OrderedSet[A], key: A, hc: Hash): int {.inline.} =
   rawGetKnownHCImpl()
 
@@ -760,6 +766,67 @@ proc incl*[A](s: var HashSet[A], other: OrderedSet[A]) =
   assert other.isValid, "The set `other` needs to be initialized."
   for item in other: incl(s, item)
 
+proc exclImpl[A](s: var OrderedSet[A], key: A) : bool {. inline .} =
+  assert s.isValid, "The set needs to be initialized."
+  var hc: Hash
+  var i = rawGet(s, key, hc)
+  var msk = high(s.data)
+  result = true
+
+  if i >= 0:
+    result = false
+    # Fix ordering
+    if s.first == i:
+      s.first = s.data[i].next
+    else:
+      var itr = s.first
+      while true:
+        if (s.data[itr].next == i):
+          s.data[itr].next = s.data[i].next
+          if s.last == i:
+            s.last = itr
+          break
+        itr = s.data[itr].next
+
+    dec(s.counter)
+    while true:         # KnuthV3 Algo6.4R adapted for i=i+1 instead of i=i-1
+      var j = i         # The correctness of this depends on (h+1) in nextTry,
+      var r = j         # though may be adaptable to other simple sequences.
+      s.data[i].hcode = 0              # mark current EMPTY
+      s.data[i].key = default(type(s.data[i].key))
+      s.data[i].next = 0
+      doWhile((i >= r and r > j) or (r > j and j > i) or (j > i and i >= r)):
+        i = (i + 1) and msk            # increment mod table size
+        if isEmpty(s.data[i].hcode):   # end of collision cluster; So all done
+          return
+        r = s.data[i].hcode and msk    # "home" location of key@i
+      shallowCopy(s.data[j], s.data[i]) # data[i] will be marked EMPTY next loop
+
+proc missingOrExcl*[A](s: var OrderedSet[A], key: A): bool =
+  ## Excludes `key` in the set `s` and tells if `key` was removed from `s`. Efficiency: O(n).
+  ##
+  ## The difference with regards to the `excl() <#excl,TOrderedSet[A],A>`_ proc is
+  ## that this proc returns `true` if `key` was not present in `s`. Example:
+  ##
+  ## .. code-block::
+  ##  var s = toOrderedSet([2, 3, 6, 7])
+  ##  assert s.missingOrExcl(4) == true
+  ##  assert s.missingOrExcl(6) == false
+  exclImpl(s, key)
+
+
+proc excl*[A](s: var OrderedSet[A], key: A) =
+  ## Excludes `key` from the set `s`. Efficiency: O(n).
+  ##
+  ## This doesn't do anything if `key` is not found in `s`. Example:
+  ##
+  ## .. code-block::
+  ##   var s = toOrderedSet([2, 3, 6, 7])
+  ##   s.excl(2)
+  ##   s.excl(2)
+  ##   assert s.len == 3
+  discard exclImpl(s, key)
+
 proc containsOrIncl*[A](s: var OrderedSet[A], key: A): bool =
   ## Includes `key` in the set `s` and tells if `key` was added to `s`.
   ##
@@ -986,6 +1053,24 @@ when isMainModule and not defined(release):
       assert a.len == b.card
       assert a.len == 2
 
+    block setPairsIterator:
+      var s = toOrderedSet([1, 3, 5, 7])
+      var items = newSeq[tuple[a: int, b: int]]()
+      for idx, item in s: items.add((idx, item))
+      assert items == @[(0, 1), (1, 3), (2, 5), (3, 7)]
+
+    block exclusions:
+      var s = toOrderedSet([1, 2, 3, 6, 7, 4])
+
+      s.excl(3)
+      s.excl(3)
+      s.excl(1)
+      s.excl(4)
+
+      var items = newSeq[int]()
+      for item in s: items.add item
+      assert items == @[2, 6, 7]
+
     #block orderedSetIterator:
     #  var a = initOrderedSet[int]()
     #  for value in [9, 2, 1, 5, 1, 8, 4, 2]:
@@ -1030,6 +1115,11 @@ when isMainModule and not defined(release):
       if s <= i or mustRehash(s, i):
         echo "performance issue: rightSize() will not elide enlarge() at ", i
 
+    block missingOrExcl:
+      var s = toOrderedSet([2, 3, 6, 7])
+      assert s.missingOrExcl(4) == true
+      assert s.missingOrExcl(6) == false
+
     when not defined(testing):
       echo "Micro tests run successfully."