diff --git a/lib/pure/collections/critbits.nim b/lib/pure/collections/critbits.nim index 5f84f31011..f70a128435 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 334e33f2e8..0852325645 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..= 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`. diff --git a/tests/sets/tsets2.nim b/tests/sets/tsets2.nim index 9c73dbe031..f288228403 100644 --- a/tests/sets/tsets2.nim +++ b/tests/sets/tsets2.nim @@ -37,16 +37,26 @@ block setTest2: t.incl("111") t.incl("123") t.excl("111") - t.incl("012") t.incl("123") # test duplicates assert "123" in t assert "111" notin t # deleted + assert t.missingOrExcl("000") == true + assert "000" notin t + assert t.missingOrExcl("012") == false + assert "012" notin t + + assert t.containsOrIncl("012") == false + assert t.containsOrIncl("012") == true + assert "012" in t # added back + for key in items(data): t.incl(key) for key in items(data): assert key in t + for key in items(data): t.excl(key) + for key in items(data): assert key notin t block orderedSetTest1: var t = data.toOrderedSet