Merge pull request #2139 from c-blake/devel

Add mgetOrPut to support just one probe chase for the common

3 files changed, 69 insertions, 23 deletions

diff --git a/lib/pure/collections/tables.nim b/lib/pure/collections/tables.nim
index c82f1230a..c75386cf1 100644
--- a/lib/pure/collections/tables.nim
+++ b/lib/pure/collections/tables.nim
@@ -231,31 +231,42 @@ template addImpl() {.dirty.} =
   rawInsert(t, t.data, key, val, hc, j)
   inc(t.counter)
 
+template maybeRehashPutImpl() {.dirty.} =
+  if mustRehash(len(t.data), t.counter):
+    enlarge(t)
+    index = rawGetKnownHC(t, key, hc)
+  index = -1 - index                  # important to transform for mgetOrPutImpl
+  rawInsert(t, t.data, key, val, hc, index)
+  inc(t.counter)
+
 template putImpl() {.dirty.} =
   var hc: THash
   var index = rawGet(t, key, hc)
-  if index >= 0:
-    t.data[index].val = val
-  else:
-    if mustRehash(len(t.data), t.counter):
-        enlarge(t)
-        index = rawGetKnownHC(t, key, hc)
-    rawInsert(t, t.data, key, val, hc, -1 - index)
-    inc(t.counter)
+  if index >= 0: t.data[index].val = val
+  else: maybeRehashPutImpl()
+
+template mgetOrPutImpl() {.dirty.} =
+  var hc: THash
+  var index = rawGet(t, key, hc)
+  if index < 0: maybeRehashPutImpl()    # not present: insert (flipping index)
+  result = t.data[index].val            # either way return modifiable val
+
+template hasKeyOrPutImpl() {.dirty.} =
+  var hc: THash
+  var index = rawGet(t, key, hc)
+  if index < 0:
+    result = false
+    maybeRehashPutImpl()
+  else: result = true
+
+proc mgetOrPut*[A, B](t: var Table[A, B], key: A, val: B): var B =
+  ## retrieves value at ``t[key]`` or puts ``val`` if not present, either way
+  ## returning a value which can be modified.
+  mgetOrPutImpl()
 
-when false:
-  # not yet used:
-  template hasKeyOrPutImpl() {.dirty.} =
-    var hc: THash
-    var index = rawGet(t, key, hc)
-    if index >= 0:
-      t.data[index].val = val
-      result = true
-    else:
-      if mustRehash(len(t.data), t.counter): enlarge(t)
-      rawInsert(t, t.data, key, val)
-      inc(t.counter)
-      result = false
+proc hasKeyOrPut*[A, B](t: var Table[A, B], key: A, val: B): bool =
+  ## returns true iff `key` is in the table, otherwise inserts `value`.
+  hasKeyOrPutImpl()
 
 proc `[]=`*[A, B](t: var Table[A, B], key: A, val: B) =
   ## puts a (key, value)-pair into `t`.
@@ -383,6 +394,15 @@ proc mget*[A, B](t: TableRef[A, B], key: A): var B =
   ## If `key` is not in `t`, the ``EInvalidKey`` exception is raised.
   t[].mget(key)
 
+proc mgetOrPut*[A, B](t: TableRef[A, B], key: A, val: B): var B =
+  ## retrieves value at ``t[key]`` or puts ``val`` if not present, either way
+  ## returning a value which can be modified.
+  t[].mgetOrPut(key, val)
+
+proc hasKeyOrPut*[A, B](t: var TableRef[A, B], key: A, val: B): bool =
+  ## returns true iff `key` is in the table, otherwise inserts `value`.
+  t[].hasKeyOrPut(key, val)
+
 proc hasKey*[A, B](t: TableRef[A, B], key: A): bool =
   ## returns true iff `key` is in the table `t`.
   result = t[].hasKey(key)
@@ -539,6 +559,15 @@ proc add*[A, B](t: var OrderedTable[A, B], key: A, val: B) =
   ## puts a new (key, value)-pair into `t` even if ``t[key]`` already exists.
   addImpl()
 
+proc mgetOrPut*[A, B](t: var OrderedTable[A, B], key: A, val: B): var B =
+  ## retrieves value at ``t[key]`` or puts ``value`` if not present, either way
+  ## returning a value which can be modified.
+  mgetOrPutImpl()
+
+proc hasKeyOrPut*[A, B](t: var OrderedTable[A, B], key: A, val: B): bool =
+  ## returns true iff `key` is in the table, otherwise inserts `value`.
+  hasKeyOrPutImpl()
+
 proc initOrderedTable*[A, B](initialSize=64): OrderedTable[A, B] =
   ## creates a new ordered hash table that is empty.
   ##
@@ -658,6 +687,15 @@ proc mget*[A, B](t: OrderedTableRef[A, B], key: A): var B =
   ## If `key` is not in `t`, the ``EInvalidKey`` exception is raised.
   result = t[].mget(key)
 
+proc mgetOrPut*[A, B](t: OrderedTableRef[A, B], key: A, val: B): var B =
+  ## retrieves value at ``t[key]`` or puts ``val`` if not present, either way
+  ## returning a value which can be modified.
+  result = t[].mgetOrPut(key, val)
+
+proc hasKeyOrPut*[A, B](t: var OrderedTableRef[A, B], key: A, val: B): bool =
+  ## returns true iff `key` is in the table, otherwise inserts `val`.
+  result = t[].hasKeyOrPut(key, val)
+
 proc hasKey*[A, B](t: OrderedTableRef[A, B], key: A): bool =
   ## returns true iff `key` is in the table `t`.
   result = t[].hasKey(key)
diff --git a/tests/collections/ttables.nim b/tests/collections/ttables.nim
index 8534e6767..7af14a75b 100644
--- a/tests/collections/ttables.nim
+++ b/tests/collections/ttables.nim
@@ -65,7 +65,15 @@ block tableTest2:
   
   for key, val in items(data): t[key] = val.toFloat
   for key, val in items(data): assert t[key] == val.toFloat
-  
+
+  assert(not t.hasKeyOrPut("456", 4.0))     # test absent key
+  assert t.hasKeyOrPut("012", 3.0)          # test present key
+  var x = t.mgetOrPut("111", 1.5)           # test absent key
+  x = x * 2
+  assert x == 3.0
+  x = t.mgetOrPut("test", 1.5)              # test present key
+  x = x * 2
+  assert x == 2 * 1.2345
 
 block orderedTableTest1:
   var t = initOrderedTable[string, int](2)
diff --git a/tests/collections/ttablesref.nim b/tests/collections/ttablesref.nim
index e666c7852..b57aedf4a 100644
--- a/tests/collections/ttablesref.nim
+++ b/tests/collections/ttablesref.nim
@@ -47,7 +47,7 @@ block tableTest1:
     for y in 0..1:
       assert t[(x,y)] == $x & $y
   assert($t == 
-    "{(x: 0, y: 0): 00, (x: 0, y: 1): 01, (x: 1, y: 0): 10, (x: 1, y: 1): 11}")
+    "{(x: 0, y: 1): 01, (x: 0, y: 0): 00, (x: 1, y: 0): 10, (x: 1, y: 1): 11}")
 
 block tableTest2:
   var t = newTable[string, float]()