Skip to content

Commit

Permalink
table.mgetOrPut without default val (#22994)
Browse files Browse the repository at this point in the history
RFC: nim-lang/RFCs#539

- ~~mgetOrPutDefaultImpl template into `tableimpl.nim` to avoid macros~~
- mgetOrPut for `Table`, `TableRef`, `OrderedTable`, `OrderedTableRef`
- `tests/stdlib/tmget.nim` tests update

---------

Co-authored-by: inv2004 <>
  • Loading branch information
inv2004 authored Nov 30, 2023
1 parent beeacc8 commit 0f7ebb4
Show file tree
Hide file tree
Showing 3 changed files with 81 additions and 4 deletions.
21 changes: 17 additions & 4 deletions lib/pure/collections/tableimpl.nim
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ template addImpl(enlarge) {.dirty.} =
rawInsert(t, t.data, key, val, hc, j)
inc(t.counter)

template maybeRehashPutImpl(enlarge) {.dirty.} =
template maybeRehashPutImpl(enlarge, val) {.dirty.} =
checkIfInitialized()
if mustRehash(t):
enlarge(t)
Expand All @@ -59,25 +59,38 @@ template putImpl(enlarge) {.dirty.} =
var hc: Hash = default(Hash)
var index = rawGet(t, key, hc)
if index >= 0: t.data[index].val = val
else: maybeRehashPutImpl(enlarge)
else: maybeRehashPutImpl(enlarge, val)

template mgetOrPutImpl(enlarge) {.dirty.} =
checkIfInitialized()
var hc: Hash = default(Hash)
var index = rawGet(t, key, hc)
if index < 0:
# not present: insert (flipping index)
maybeRehashPutImpl(enlarge)
when declared(val):
maybeRehashPutImpl(enlarge, val)
else:
maybeRehashPutImpl(enlarge, default(B))
# either way return modifiable val
result = t.data[index].val

# template mgetOrPutDefaultImpl(enlarge) {.dirty.} =
# checkIfInitialized()
# var hc: Hash = default(Hash)
# var index = rawGet(t, key, hc)
# if index < 0:
# # not present: insert (flipping index)
# maybeRehashPutImpl(enlarge, default(B))
# # either way return modifiable val
# result = t.data[index].val

template hasKeyOrPutImpl(enlarge) {.dirty.} =
checkIfInitialized()
var hc: Hash = default(Hash)
var index = rawGet(t, key, hc)
if index < 0:
result = false
maybeRehashPutImpl(enlarge)
maybeRehashPutImpl(enlarge, val)
else: result = true

# delImplIdx is KnuthV3 Algo6.4R adapted to i=i+1 (from i=i-1) which has come to
Expand Down
48 changes: 48 additions & 0 deletions lib/pure/collections/tables.nim
Original file line number Diff line number Diff line change
Expand Up @@ -474,6 +474,18 @@ proc mgetOrPut*[A, B](t: var Table[A, B], key: A, val: B): var B =

mgetOrPutImpl(enlarge)

proc mgetOrPut*[A, B](t: var Table[A, B], key: A): var B =
## Retrieves the value at `t[key]` or puts the
## default initialization value for type `B` (e.g. 0 for any
## integer type).
runnableExamples:
var a = {'a': 5}.newTable
doAssert a.mgetOrPut('a') == 5
a.mgetOrPut('z').inc
doAssert a == {'a': 5, 'z': 1}.newTable

mgetOrPutImpl(enlarge)

proc len*[A, B](t: Table[A, B]): int =
## Returns the number of keys in `t`.
runnableExamples:
Expand Down Expand Up @@ -1013,6 +1025,18 @@ proc mgetOrPut*[A, B](t: TableRef[A, B], key: A, val: B): var B =
doAssert t[25] == @[25, 35]
t[].mgetOrPut(key, val)

proc mgetOrPut*[A, B](t: TableRef[A, B], key: A): var B =
## Retrieves the value at `t[key]` or puts the
## default initialization value for type `B` (e.g. 0 for any
## integer type).
runnableExamples:
var a = {'a': 5}.newTable
doAssert a.mgetOrPut('a') == 5
a.mgetOrPut('z').inc
doAssert a == {'a': 5, 'z': 1}.newTable

t[].mgetOrPut(key)

proc len*[A, B](t: TableRef[A, B]): int =
## Returns the number of keys in `t`.
runnableExamples:
Expand Down Expand Up @@ -1503,6 +1527,18 @@ proc mgetOrPut*[A, B](t: var OrderedTable[A, B], key: A, val: B): var B =

mgetOrPutImpl(enlarge)

proc mgetOrPut*[A, B](t: var OrderedTable[A, B], key: A): var B =
## Retrieves the value at `t[key]` or puts the
## default initialization value for type `B` (e.g. 0 for any
## integer type).
runnableExamples:
var a = {'a': 5}.toOrderedTable
doAssert a.mgetOrPut('a') == 5
a.mgetOrPut('z').inc
doAssert a == {'a': 5, 'z': 1}.toOrderedTable

mgetOrPutImpl(enlarge)

proc len*[A, B](t: OrderedTable[A, B]): int {.inline.} =
## Returns the number of keys in `t`.
runnableExamples:
Expand Down Expand Up @@ -1992,6 +2028,18 @@ proc mgetOrPut*[A, B](t: OrderedTableRef[A, B], key: A, val: B): var B =

result = t[].mgetOrPut(key, val)

proc mgetOrPut*[A, B](t: OrderedTableRef[A, B], key: A): var B =
## Retrieves the value at `t[key]` or puts the
## default initialization value for type `B` (e.g. 0 for any
## integer type).
runnableExamples:
var a = {'a': 5}.toOrderedTable
doAssert a.mgetOrPut('a') == 5
a.mgetOrPut('z').inc
doAssert a == {'a': 5, 'z': 1}.toOrderedTable

t[].mgetOrPut(key)

proc len*[A, B](t: OrderedTableRef[A, B]): int {.inline.} =
## Returns the number of keys in `t`.
runnableExamples:
Expand Down
16 changes: 16 additions & 0 deletions tests/stdlib/tmget.nim
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,19 @@ discard """
output: '''Can't access 6
10
11
2
Can't access 6
10
11
2
Can't access 6
10
11
2
Can't access 6
10
11
2
0
10
11
Expand Down Expand Up @@ -41,6 +45,9 @@ block:
x[5] += 1
var c = x[5]
echo c
x.mgetOrPut(7).inc
x.mgetOrPut(7).inc
echo x[7]

block:
var x = newTable[int, int]()
Expand All @@ -53,6 +60,9 @@ block:
x[5] += 1
var c = x[5]
echo c
x.mgetOrPut(7).inc
x.mgetOrPut(7).inc
echo x[7]

block:
var x = initOrderedTable[int, int]()
Expand All @@ -65,6 +75,9 @@ block:
x[5] += 1
var c = x[5]
echo c
x.mgetOrPut(7).inc
x.mgetOrPut(7).inc
echo x[7]

block:
var x = newOrderedTable[int, int]()
Expand All @@ -77,6 +90,9 @@ block:
x[5] += 1
var c = x[5]
echo c
x.mgetOrPut(7).inc
x.mgetOrPut(7).inc
echo x[7]

block:
var x = initCountTable[int]()
Expand Down

0 comments on commit 0f7ebb4

Please sign in to comment.