From 06bdf8392bd3e597dbe3ef12c7a819b60e32d7ac Mon Sep 17 00:00:00 2001
From: Dmitry Atamanov <data-man@users.noreply.github.com>
Date: Fri, 18 May 2018 22:13:29 +0300
Subject: [PATCH] Fixes CritBitTree.inc's bug (#7838)

* Fixes CritBitTree.inc's bug

* Update changelog
---
 changelog.md                      |  1 +
 lib/pure/collections/critbits.nim | 18 ++++++++++++++----
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/changelog.md b/changelog.md
index cbefbc421e..f9daf55d13 100644
--- a/changelog.md
+++ b/changelog.md
@@ -74,6 +74,7 @@
   deprecated.
 - The `terminal` module now exports additional procs for generating ANSI color
   codes as strings.
+- Added the parameter ``val`` for the ``CritBitTree[int].inc`` proc.
 
 ### Language additions
 
diff --git a/lib/pure/collections/critbits.nim b/lib/pure/collections/critbits.nim
index 95fffdf88c..5ae5e26b21 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