Add counterpart to containsOrIncl for excl (#6360)

3 files changed, 93 insertions, 37 deletions

diff --git a/lib/pure/collections/critbits.nim b/lib/pure/collections/critbits.nim
index 5f84f3101..f70a12843 100644
--- a/lib/pure/collections/critbits.nim
+++ b/lib/pure/collections/critbits.nim
@@ -110,6 +110,42 @@ proc rawInsert[T](c: var CritBitTree[T], key: string): Node[T] =
     wherep[] = inner
   inc c.count
 
+proc exclImpl[T](c: var CritBitTree[T], key: string) : int =
+  var p = c.root
+  var wherep = addr(c.root)
+  var whereq: ptr Node[T] = nil
+  if p == nil: return c.count
+  var dir = 0
+  var q: Node[T]
+  while not p.isLeaf:
+    whereq = wherep
+    q = p
+    let ch = if p.byte < key.len: key[p.byte] else: '\0'
+    dir = (1 + (ch.ord or p.otherBits.ord)) shr 8
+    wherep = addr(p.child[dir])
+    p = wherep[]
+  if p.key == key:
+    # else: not in tree at all
+    if whereq == nil:
+      c.root = nil
+    else:
+      whereq[] = q.child[1 - dir]
+    dec c.count
+
+  return c.count
+
+proc excl*[T](c: var CritBitTree[T], key: string) =
+  ## removes `key` (and its associated value) from the set `c`.
+  ## If the `key` does not exist, nothing happens.
+  discard exclImpl(c, key)
+
+proc missingOrExcl*[T](c: var CritBitTree[T], key: string): bool =
+  ## Returns true iff `c` does not contain the given `key`. If the key
+  ## does exist, c.excl(key) is performed. 
+  let oldCount = c.count 
+  var n = exclImpl(c, key)
+  result = c.count == oldCount
+
 proc containsOrIncl*[T](c: var CritBitTree[T], key: string, val: T): bool =
   ## returns true iff `c` contains the given `key`. If the key does not exist
   ## ``c[key] = val`` is performed.
@@ -171,30 +207,6 @@ proc mget*[T](c: var CritBitTree[T], key: string): var T {.inline, deprecated.}
   ## Use ```[]``` instead.
   get(c, key)
 
-proc excl*[T](c: var CritBitTree[T], key: string) =
-  ## removes `key` (and its associated value) from the set `c`.
-  ## If the `key` does not exist, nothing happens.
-  var p = c.root
-  var wherep = addr(c.root)
-  var whereq: ptr Node[T] = nil
-  if p == nil: return
-  var dir = 0
-  var q: Node[T]
-  while not p.isLeaf:
-    whereq = wherep
-    q = p
-    let ch = if p.byte < key.len: key[p.byte] else: '\0'
-    dir = (1 + (ch.ord or p.otherBits.ord)) shr 8
-    wherep = addr(p.child[dir])
-    p = wherep[]
-  if p.key == key:
-    # else: not in tree at all
-    if whereq == nil:
-      c.root = nil
-    else:
-      whereq[] = q.child[1 - dir]
-    dec c.count
-
 iterator leaves[T](n: Node[T]): Node[T] =
   if n != nil:
     # XXX actually we could compute the necessary stack size in advance:
@@ -326,10 +338,15 @@ when isMainModule:
   r.incl "def"
   r.incl "definition"
   r.incl "prefix"
+  r.incl "foo"
 
   doAssert r.contains"def"
 
   r.excl "def"
+  assert r.missingOrExcl("foo") == false
+  assert "foo" notin toSeq(r.items)
+
+  assert r.missingOrExcl("foo") == true
 
   assert toSeq(r.items) == @["abc", "definition", "prefix", "xyz"]
 
diff --git a/lib/pure/collections/intsets.nim b/lib/pure/collections/intsets.nim
index 334e33f2e..085232564 100644
--- a/lib/pure/collections/intsets.nim
+++ b/lib/pure/collections/intsets.nim
@@ -131,8 +131,7 @@ proc incl*(s: var IntSet, key: int) =
     # fall through:
   bitincl(s, key)
 
-proc excl*(s: var IntSet, key: int) =
-  ## excludes `key` from the set `s`.
+proc exclImpl(s: var IntSet, key: int) =
   if s.elems <= s.a.len:
     for i in 0..<s.elems:
       if s.a[i] == key:
@@ -146,6 +145,17 @@ proc excl*(s: var IntSet, key: int) =
       t.bits[`shr`(u, IntShift)] = t.bits[`shr`(u, IntShift)] and
           not `shl`(1, u and IntMask)
 
+proc excl*(s: var IntSet, key: int) =
+  ## excludes `key` from the set `s`.
+  exclImpl(s, key)
+
+proc missingOrExcl*(s: var IntSet, key: int) : bool =
+  ## returns true if `s` does not contain `key`, otherwise
+  ## `key` is removed from `s` and false is returned.
+  var count = s.elems
+  exclImpl(s, key)
+  result = count == s.elems 
+
 proc containsOrIncl*(s: var IntSet, key: int): bool =
   ## returns true if `s` contains `key`, otherwise `key` is included in `s`
   ## and false is returned.
@@ -270,6 +280,17 @@ when isMainModule:
   x.incl(7)
   x.incl(1056)
 
+  x.incl(1044)
+  x.excl(1044) 
+
+  assert x.containsOrIncl(888) == false
+  assert 888 in x
+  assert x.containsOrIncl(888) == true
+
+  assert x.missingOrExcl(888) == false
+  assert 888 notin x
+  assert x.missingOrExcl(888) == true
+
   var xs = toSeq(items(x))
   xs.sort(cmp[int])
   assert xs == @[1, 2, 7, 1056]
diff --git a/lib/pure/collections/sets.nim b/lib/pure/collections/sets.nim
index c0ffcb19c..d51a5c388 100644
--- a/lib/pure/collections/sets.nim
+++ b/lib/pure/collections/sets.nim
@@ -278,21 +278,15 @@ template default[T](t: typedesc[T]): T =
   var v: T
   v
 
-proc excl*[A](s: var HashSet[A], key: A) =
-  ## Excludes `key` from the set `s`.
-  ##
-  ## This doesn't do anything if `key` is not found in `s`. Example:
-  ##
-  ## .. code-block::
-  ##   var s = toSet([2, 3, 6, 7])
-  ##   s.excl(2)
-  ##   s.excl(2)
-  ##   assert s.len == 3
+proc exclImpl[A](s: var HashSet[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
     s.data[i].hcode = 0
     s.data[i].key = default(type(s.data[i].key))
     dec(s.counter)
@@ -308,6 +302,30 @@ proc excl*[A](s: var HashSet[A], key: A) =
         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
 
+proc missingOrExcl*[A](s: var HashSet[A], key: A): bool =
+  ## Excludes `key` in the set `s` and tells if `key` was removed from `s`.
+  ##
+  ## The difference with regards to the `excl() <#excl,TSet[A],A>`_ proc is
+  ## that this proc returns `true` if `key` was not present in `s`. Example:
+  ##
+  ## .. code-block::
+  ##  var s = toSet([2, 3, 6, 7])
+  ##  assert s.missingOrExcl(4) == true
+  ##  assert s.missingOrExcl(6) == false
+  exclImpl(s, key)
+
+proc excl*[A](s: var HashSet[A], key: A) =
+  ## Excludes `key` from the set `s`.
+  ##
+  ## This doesn't do anything if `key` is not found in `s`. Example:
+  ##
+  ## .. code-block::
+  ##   var s = toSet([2, 3, 6, 7])
+  ##   s.excl(2)
+  ##   s.excl(2)
+  ##   assert s.len == 3
+  discard exclImpl(s, key)
+
 proc excl*[A](s: var HashSet[A], other: HashSet[A]) =
   ## Excludes everything in `other` from `s`.
   ##
@@ -322,7 +340,7 @@ proc excl*[A](s: var HashSet[A], other: HashSet[A]) =
   ##   # --> {1, 3, 5}
   assert s.isValid, "The set `s` needs to be initialized."
   assert other.isValid, "The set `other` needs to be initialized."
-  for item in other: excl(s, item)
+  for item in other: discard exclImpl(s, item)
 
 proc containsOrIncl*[A](s: var HashSet[A], key: A): bool =
   ## Includes `key` in the set `s` and tells if `key` was added to `s`.