Nested `with` blocks (#22042)

* Implemented with-nesting in underscoredCalls()

* Add tests for nested with

3 files changed, 40 insertions, 6 deletions

diff --git a/lib/std/private/underscored_calls.nim b/lib/std/private/underscored_calls.nim
index f0bcbcc74..8b0392641 100644
--- a/lib/std/private/underscored_calls.nim
+++ b/lib/std/private/underscored_calls.nim
@@ -12,6 +12,8 @@
 
 import macros
 
+proc underscoredCalls*(result, calls, arg0: NimNode)
+
 proc underscoredCall(n, arg0: NimNode): NimNode =
   proc underscorePos(n: NimNode): int =
     for i in 1 ..< n.len:
@@ -19,13 +21,19 @@ proc underscoredCall(n, arg0: NimNode): NimNode =
     return 0
 
   if n.kind in nnkCallKinds:
-    result = copyNimNode(n)
-    result.add n[0]
+    if n[0].kind in {nnkIdent, nnkSym} and n[0].eqIdent("with"):
+      expectKind n[1], {nnkIdent, nnkSym}
 
-    let u = underscorePos(n)
-    for i in 1..u-1: result.add n[i]
-    result.add arg0
-    for i in u+1..n.len-1: result.add n[i]
+      result = newStmtList()
+      underscoredCalls(result, n[2 .. ^1].newStmtList, newDotExpr(arg0, n[1]))
+    else:
+      result = copyNimNode(n)
+      result.add n[0]
+
+      let u = underscorePos(n)
+      for i in 1..u-1: result.add n[i]
+      result.add arg0
+      for i in u+1..n.len-1: result.add n[i]
   elif n.kind in {nnkAsgn, nnkExprEqExpr}:
     var field = n[0]
     if n[0].kind == nnkDotExpr and n[0][0].eqIdent("_"):
diff --git a/lib/std/with.nim b/lib/std/with.nim
index c7338b4e4..7d14a4100 100644
--- a/lib/std/with.nim
+++ b/lib/std/with.nim
@@ -35,5 +35,15 @@ macro with*(arg: typed; calls: varargs[untyped]): untyped =
       -= 5
     doAssert a == 43
 
+    # Nesting works for object types too!
+    var foo = (bar: 1, qux: (baz: 2))
+    with foo:
+      bar = 2
+      with qux:
+        baz = 3
+    doAssert foo.bar == 2
+    doAssert foo.qux.baz == 3
+
   result = newNimNode(nnkStmtList, arg)
   underscoredCalls(result, calls, arg)
+  echo result.astGenRepr
diff --git a/tests/stdlib/twith.nim b/tests/stdlib/twith.nim
index ea3f3d99f..ceadbe7bf 100644
--- a/tests/stdlib/twith.nim
+++ b/tests/stdlib/twith.nim
@@ -26,3 +26,19 @@ with f:
 doAssert f.col == "(2, 3, 4)"
 doAssert f.pos == "(0.0, 1.0)"
 doAssert f.name == "bar"
+
+type
+  Baz* = object
+    a*, b*: int
+  Bar* = object
+    x*: int
+    baz*: Baz
+
+var bar: Bar
+with bar:
+  x = 1
+  with baz:
+    a = 2
+
+doAssert bar.x == 1
+doAssert bar.baz.a == 2