diff options
Diffstat (limited to 'lib/pure/collections/critbits.nim')
| -rw-r--r-- | lib/pure/collections/critbits.nim | 18 |
1 files changed, 14 insertions, 4 deletions
diff --git a/lib/pure/collections/critbits.nim b/lib/pure/collections/critbits.nim index 95fffdf88..5ae5e26b2 100644 --- a/lib/pure/collections/critbits.nim +++ b/lib/pure/collections/critbits.nim @@ -163,13 +163,13 @@ proc containsOrIncl*(c: var CritBitTree[void], key: string): bool = var n = rawInsert(c, key) result = c.count == oldCount -proc inc*(c: var CritBitTree[int]; key: string) = - ## counts the 'key'. +proc inc*(c: var CritBitTree[int]; key: string, val: int = 1) = + ## increments `c[key]` by `val`. let oldCount = c.count var n = rawInsert(c, key) - if c.count == oldCount: + if c.count == oldCount or oldCount == 0: # not a new key: - inc n.val + inc n.val, val proc incl*(c: var CritBitTree[void], key: string) = ## includes `key` in `c`. @@ -352,3 +352,13 @@ when isMainModule: assert toSeq(r.items) == @["abc", "definition", "prefix", "xyz"] assert toSeq(r.itemsWithPrefix("de")) == @["definition"] + var c = CritBitTree[int]() + + c.inc("a") + assert c["a"] == 1 + + c.inc("a", 4) + assert c["a"] == 5 + + c.inc("a", -5) + assert c["a"] == 0 |