From 40d0566ebfd18824b93b2dc9b2721e823165d0f7 Mon Sep 17 00:00:00 2001
From: Zahary Karadjov <zahary@gmail.com>
Date: Fri, 7 Aug 2020 21:39:46 +0300
Subject: [PATCH] Add tables.mgetOrPutLazy

---
 stew/shims/tables.nim | 11 +++++++++++
 tests/all_tests.nim   |  5 +++--
 tests/test_tables.nim | 15 +++++++++++++++
 3 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_tables.nim

diff --git a/stew/shims/tables.nim b/stew/shims/tables.nim
index ee06bc8..1988f07 100644
--- a/stew/shims/tables.nim
+++ b/stew/shims/tables.nim
@@ -9,4 +9,15 @@ template init*[A, B](T: type OrderedTableRef[A, B]): auto = newOrderedTable[A, B
 template init*[A](T: type CountTable[A]): auto = initCountTable[A]()
 template init*[A](T: type CountTableRef[A]): auto = newCountTable[A]()
 
+template mgetOrPutLazy*[A, B](t: Table[A, B], key: A, val: B): var B =
+  type R = B
+
+  proc setter(loc: var R): var R =
+    if loc == default(R):
+      loc = val
+    loc
+
+  setter(mgetOrPut(t, key, default(R)))
+
 export tables, objects
+
diff --git a/tests/all_tests.nim b/tests/all_tests.nim
index f6c25ac..f6f46c3 100644
--- a/tests/all_tests.nim
+++ b/tests/all_tests.nim
@@ -15,10 +15,11 @@ import
   test_bitops2,
   test_bitseqs,
   test_byteutils,
+  test_ctops,
   test_endians2,
   test_macros,
   test_objects,
   test_ptrops,
   test_results,
-  test_varints,
-  test_ctops
+  test_tables,
+  test_varints
diff --git a/tests/test_tables.nim b/tests/test_tables.nim
new file mode 100644
index 0000000..a2c9c94
--- /dev/null
+++ b/tests/test_tables.nim
@@ -0,0 +1,15 @@
+import
+  unittest,
+  ../stew/shims/tables
+
+suite "tables":
+  test "mgetOrPutLazy":
+    var t = {"a": 10, "b": 20}.toTable
+    check t.mgetOrPutLazy("a", 30) == 10
+    check t.mgetOrPutLazy("c", 40) == 40
+
+    check t.mgetOrPutLazy("c", 20) == 40
+
+    t.mgetOrPutLazy("c", 20) = 15
+    check t.mgetOrPutLazy("c", 40) == 15
+