Merge pull request #2139 from c-blake/devel

Add mgetOrPut to support just one probe chase for the common

1 files changed, 59 insertions, 21 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)