Merge pull request #7075 from lcrees/get_default

Add default argument to tables.getOrDefault

1 files changed, 80 insertions, 13 deletions

diff --git a/lib/pure/collections/tables.nim b/lib/pure/collections/tables.nim
index c97846f92..7b508b294 100644
--- a/lib/pure/collections/tables.nim
+++ b/lib/pure/collections/tables.nim
@@ -158,6 +158,12 @@ template getOrDefaultImpl(t, key): untyped =
   var index = rawGet(t, key, hc)
   if index >= 0: result = t.data[index].val
 
+template getOrDefaultImpl(t, key, default: untyped): untyped =
+  mixin rawGet
+  var hc: Hash
+  var index = rawGet(t, key, hc)
+  result = if index >= 0: t.data[index].val else: default
+
 proc `[]`*[A, B](t: Table[A, B], key: A): B {.deprecatedGet.} =
   ## retrieves the value at ``t[key]``. If `key` is not in `t`, the
   ## ``KeyError`` exception is raised. One can check with ``hasKey`` whether
@@ -175,10 +181,18 @@ proc mget*[A, B](t: var Table[A, B], key: A): var B {.deprecated.} =
   ## instead.
   get(t, key)
 
-proc getOrDefault*[A, B](t: Table[A, B], key: A): B = getOrDefaultImpl(t, key)
+proc getOrDefault*[A, B](t: Table[A, B], key: A): B =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, the
+  ## default initialization value for type `B` is returned (e.g. 0 for any
+  ## integer type).
+  getOrDefaultImpl(t, key)
 
-template withValue*[A, B](t: var Table[A, B], key: A,
-                          value, body: untyped) =
+proc getOrDefault*[A, B](t: Table[A, B], key: A, default: B): B =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, `default`
+  ## is returned.
+  getOrDefaultImpl(t, key, default)
+
+template withValue*[A, B](t: var Table[A, B], key: A, value, body: untyped) =
   ## retrieves the value at ``t[key]``.
   ## `value` can be modified in the scope of the ``withValue`` call.
   ##
@@ -325,8 +339,7 @@ proc initTable*[A, B](initialSize=64): Table[A, B] =
   result.counter = 0
   newSeq(result.data, initialSize)
 
-proc toTable*[A, B](pairs: openArray[(A,
-                    B)]): Table[A, B] =
+proc toTable*[A, B](pairs: openArray[(A, B)]): Table[A, B] =
   ## creates a new hash table that contains the given `pairs`.
   result = initTable[A, B](rightSize(pairs.len))
   for key, val in items(pairs): result[key] = val
@@ -410,7 +423,16 @@ proc mget*[A, B](t: TableRef[A, B], key: A): var B {.deprecated.} =
   ## Use ```[]``` instead.
   t[][key]
 
-proc getOrDefault*[A, B](t: TableRef[A, B], key: A): B = getOrDefault(t[], key)
+proc getOrDefault*[A, B](t: TableRef[A, B], key: A): B =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, the
+  ## default initialization value for type `B` is returned (e.g. 0 for any
+  ## integer type).
+  getOrDefault(t[], key)
+
+proc getOrDefault*[A, B](t: TableRef[A, B], key: A, default: B): B =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, `default`
+  ## is returned.
+  getOrDefault(t[], key, default)
 
 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
@@ -562,8 +584,15 @@ proc mget*[A, B](t: var OrderedTable[A, B], key: A): var B {.deprecated.} =
   get(t, key)
 
 proc getOrDefault*[A, B](t: OrderedTable[A, B], key: A): B =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, the
+  ## default initialization value for type `B` is returned (e.g. 0 for any
+  ## integer type).
   getOrDefaultImpl(t, key)
 
+proc getOrDefault*[A, B](t: OrderedTable[A, B], key: A, default: B): B =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, `default`
+  ## is returned.
+  getOrDefaultImpl(t, key, default)
 
 proc hasKey*[A, B](t: OrderedTable[A, B], key: A): bool =
   ## returns true iff `key` is in the table `t`.
@@ -630,8 +659,7 @@ proc initOrderedTable*[A, B](initialSize=64): OrderedTable[A, B] =
   result.last = -1
   newSeq(result.data, initialSize)
 
-proc toOrderedTable*[A, B](pairs: openArray[(A,
-                           B)]): OrderedTable[A, B] =
+proc toOrderedTable*[A, B](pairs: openArray[(A, B)]): OrderedTable[A, B] =
   ## creates a new ordered hash table that contains the given `pairs`.
   result = initOrderedTable[A, B](rightSize(pairs.len))
   for key, val in items(pairs): result[key] = val
@@ -657,8 +685,7 @@ proc `==`*[A, B](s, t: OrderedTable[A, B]): bool =
     hs = nxts
   return true
 
-proc sort*[A, B](t: var OrderedTable[A, B],
-                 cmp: proc (x,y: (A, B)): int) =
+proc sort*[A, B](t: var OrderedTable[A, B], cmp: proc (x,y: (A, B)): int) =
   ## sorts `t` according to `cmp`. This modifies the internal list
   ## that kept the insertion order, so insertion order is lost after this
   ## call but key lookup and insertions remain possible after `sort` (in
@@ -748,8 +775,16 @@ proc mget*[A, B](t: OrderedTableRef[A, B], key: A): var B {.deprecated.} =
   result = t[][key]
 
 proc getOrDefault*[A, B](t: OrderedTableRef[A, B], key: A): B =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, the
+  ## default initialization value for type `B` is returned (e.g. 0 for any
+  ## integer type).
   getOrDefault(t[], key)
 
+proc getOrDefault*[A, B](t: OrderedTableRef[A, B], key: A, default: B): B =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, `default`
+  ## is returned.
+  getOrDefault(t[], key, default)
+
 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.
@@ -802,8 +837,7 @@ proc `==`*[A, B](s, t: OrderedTableRef[A, B]): bool =
   elif isNil(t): result = false
   else: result = s[] == t[]
 
-proc sort*[A, B](t: OrderedTableRef[A, B],
-                 cmp: proc (x,y: (A, B)): int) =
+proc sort*[A, B](t: OrderedTableRef[A, B], cmp: proc (x,y: (A, B)): int) =
   ## sorts `t` according to `cmp`. This modifies the internal list
   ## that kept the insertion order, so insertion order is lost after this
   ## call but key lookup and insertions remain possible after `sort` (in
@@ -916,9 +950,17 @@ proc mget*[A](t: var CountTable[A], key: A): var int {.deprecated.} =
   ctget(t, key)
 
 proc getOrDefault*[A](t: CountTable[A], key: A): int =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, 0 (the
+  ## default initialization value of `int`), is returned.
   var index = rawGet(t, key)
   if index >= 0: result = t.data[index].val
 
+proc getOrDefault*[A](t: CountTable[A], key: A, default: int): int =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, the
+  ## integer value of `default` is returned.
+  var index = rawGet(t, key)
+  result = if index >= 0: t.data[index].val else: default
+
 proc hasKey*[A](t: CountTable[A], key: A): bool =
   ## returns true iff `key` is in the table `t`.
   result = rawGet(t, key) >= 0
@@ -1073,8 +1115,15 @@ proc mget*[A](t: CountTableRef[A], key: A): var int {.deprecated.} =
   result = t[][key]
 
 proc getOrDefault*[A](t: CountTableRef[A], key: A): int =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, 0 (the
+  ## default initialization value of `int`), is returned.
   result = t[].getOrDefault(key)
 
+proc getOrDefault*[A](t: CountTableRef[A], key: A, default: int): int =
+  ## retrieves the value at ``t[key]`` iff `key` is in `t`. Otherwise, the
+  ## integer value of `default` is returned.
+  result = t[].getOrDefault(key, default)
+
 proc hasKey*[A](t: CountTableRef[A], key: A): bool =
   ## returns true iff `key` is in the table `t`.
   result = t[].hasKey(key)
@@ -1267,7 +1316,7 @@ when isMainModule:
     #lib/pure/collections/tables.nim(117, 21) template/generic instantiation from here
     #lib/pure/collections/tableimpl.nim(32, 27) Error: undeclared field: 'hcode
     doAssert 0 == t.getOrDefault(testKey)
-    t.inc(testKey,3)
+    t.inc(testKey, 3)
     doAssert 3 == t.getOrDefault(testKey)
 
   block:
@@ -1334,3 +1383,21 @@ when isMainModule:
   block: # CountTable.smallest
     let t = toCountTable([0, 0, 5, 5, 5])
     doAssert t.smallest == (0, 2)
+
+  block:
+    var tp: Table[string, string] = initTable[string, string]()
+    doAssert "test1" == tp.getOrDefault("test1", "test1")
+    tp["test2"] = "test2"
+    doAssert "test2" == tp.getOrDefault("test2", "test1")
+    var tr: TableRef[string, string] = newTable[string, string]()
+    doAssert "test1" == tr.getOrDefault("test1", "test1")
+    tr["test2"] = "test2"
+    doAssert "test2" == tr.getOrDefault("test2", "test1")
+    var op: OrderedTable[string, string] = initOrderedTable[string, string]()
+    doAssert "test1" == op.getOrDefault("test1", "test1")
+    op["test2"] = "test2"
+    doAssert "test2" == op.getOrDefault("test2", "test1")
+    var orf: OrderedTableRef[string, string] = newOrderedTable[string, string]()
+    doAssert "test1" == orf.getOrDefault("test1", "test1")
+    orf["test2"] = "test2"
+    doAssert "test2" == orf.getOrDefault("test2", "test1")