major improvements to `std/wrapnils`: optimal codegen, case objects, lvalue semantics (#18435)

* wrapnils now generates optimal code; also handles case objects
* changelog
* unsafeAddr => addr

1 files changed, 134 insertions, 51 deletions

diff --git a/lib/std/wrapnils.nim b/lib/std/wrapnils.nim
index ed0a79d79..facba85fa 100644
--- a/lib/std/wrapnils.nim
+++ b/lib/std/wrapnils.nim
@@ -1,9 +1,20 @@
-## This module allows chains of field-access and indexing where the LHS can be nil.
-## This simplifies code by reducing need for if-else branches around intermediate values
-## that may be nil.
+## This module allows evaluating expressions safely against the following conditions:
+## * nil dereferences
+## * field accesses with incorrect discriminant in case objects
+##
+## `default(T)` is returned in those cases when evaluating an expression of type `T`.
+## This simplifies code by reducing need for if-else branches.
 ##
 ## Note: experimental module, unstable API.
 
+#[
+TODO:
+consider handling indexing operations, eg:
+doAssert ?.default(seq[int])[3] == default(int)
+]#
+
+import macros
+
 runnableExamples:
   type Foo = ref object
     x1: string
@@ -24,8 +35,124 @@ runnableExamples:
 
   assert (?.f2.x2.x2).x3 == nil  # this terminates ?. early
 
+runnableExamples:
+  # ?. also allows case object
+  type B = object
+    b0: int
+    case cond: bool
+    of false: discard
+    of true:
+      b1: float
+
+  var b = B(cond: false, b0: 3)
+  doAssertRaises(FieldDefect): discard b.b1 # wrong discriminant
+  doAssert ?.b.b1 == 0.0 # safe
+  b = B(cond: true, b1: 4.5)
+  doAssert ?.b.b1 == 4.5
+
+  # lvalue semantics are preserved:
+  if (let p = ?.b.b1.addr; p != nil): p[] = 4.7
+  doAssert b.b1 == 4.7
+
+proc finalize(n: NimNode, lhs: NimNode, level: int): NimNode =
+  if level == 0:
+    result = quote: `lhs` = `n`
+  else:
+    result = quote: (let `lhs` = `n`)
+
+proc process(n: NimNode, lhs: NimNode, level: int): NimNode =
+  var n = n.copyNimTree
+  var it = n
+  let addr2 = bindSym"addr"
+  var old: tuple[n: NimNode, index: int]
+  while true:
+    if it.len == 0:
+      result = finalize(n, lhs, level)
+      break
+    elif it.kind == nnkCheckedFieldExpr:
+      let dot = it[0]
+      let obj = dot[0]
+      let objRef = quote do: `addr2`(`obj`)
+        # avoids a copy and preserves lvalue semantics, see tests
+      let check = it[1]
+      let okSet = check[1]
+      let kind1 = check[2]
+      let tmp = genSym(nskLet, "tmpCase")
+      let body = process(objRef, tmp, level + 1)
+      let tmp3 = nnkDerefExpr.newTree(tmp)
+      it[0][0] = tmp3
+      let dot2 = nnkDotExpr.newTree(@[tmp, dot[1]])
+      if old.n != nil: old.n[old.index] = dot2
+      else: n = dot2
+      let assgn = finalize(n, lhs, level)
+      result = quote do:
+        `body`
+        if `tmp3`.`kind1` notin `okSet`: break
+        `assgn`
+      break
+    elif it.kind in {nnkHiddenDeref, nnkDerefExpr}:
+      let tmp = genSym(nskLet, "tmp")
+      let body = process(it[0], tmp, level + 1)
+      it[0] = tmp
+      let assgn = finalize(n, lhs, level)
+      result = quote do:
+        `body`
+        if `tmp` == nil: break
+        `assgn`
+      break
+    elif it.kind == nnkCall: # consider extending to `nnkCallKinds`
+      # `copyNimTree` needed to avoid `typ = nil` issues
+      old = (it, 1)
+      it = it[1].copyNimTree
+    else:
+      old = (it, 0)
+      it = it[0]
+
+macro `?.`*(a: typed): auto =
+  ## Transforms `a` into an expression that can be safely evaluated even in
+  ## presence of intermediate nil pointers/references, in which case a default
+  ## value is produced.
+  let lhs = genSym(nskVar, "lhs")
+  let body = process(a, lhs, 0)
+  result = quote do:
+    var `lhs`: type(`a`)
+    block:
+      `body`
+    `lhs`
+
+# the code below is not needed for `?.`
 from options import Option, isSome, get, option, unsafeGet, UnpackDefect
 
+macro `??.`*(a: typed): Option =
+  ## Same as `?.` but returns an `Option`.
+  runnableExamples:
+    import std/options
+    type Foo = ref object
+      x1: ref int
+      x2: int
+    # `?.` can't distinguish between a valid vs invalid default value, but `??.` can:
+    var f1 = Foo(x1: int.new, x2: 2)
+    doAssert (??.f1.x1[]).get == 0 # not enough to tell when the chain was valid.
+    doAssert (??.f1.x1[]).isSome # a nil didn't occur in the chain
+    doAssert (??.f1.x2).get == 2
+
+    var f2: Foo
+    doAssert not (??.f2.x1[]).isSome # f2 was nil
+
+    doAssertRaises(UnpackDefect): discard (??.f2.x1[]).get
+    doAssert ?.f2.x1[] == 0 # in contrast, this returns default(int)
+
+  let lhs = genSym(nskVar, "lhs")
+  let lhs2 = genSym(nskVar, "lhs")
+  let body = process(a, lhs2, 0)
+  result = quote do:
+    var `lhs`: Option[type(`a`)]
+    block:
+      var `lhs2`: type(`a`)
+      `body`
+      `lhs` = option(`lhs2`)
+    `lhs`
+
 template fakeDot*(a: Option, b): untyped =
   ## See top-level example.
   let a1 = a # to avoid double evaluations
@@ -58,51 +185,7 @@ func `[]`*[U](a: Option[U]): auto {.inline.} =
     if a2 != nil:
       result = option(a2[])
 
-import macros
-
-func replace(n: NimNode): NimNode =
-  if n.kind == nnkDotExpr:
-    result = newCall(bindSym"fakeDot", replace(n[0]), n[1])
-  elif n.kind == nnkPar:
-    doAssert n.len == 1
-    result = newCall(bindSym"option", n[0])
-  elif n.kind in {nnkCall, nnkObjConstr}:
-    result = newCall(bindSym"option", n)
-  elif n.len == 0:
-    result = newCall(bindSym"option", n)
-  else:
-    n[0] = replace(n[0])
-    result = n
-
-proc safeGet[T](a: Option[T]): T {.inline.} =
-  get(a, default(T))
-
-macro `?.`*(a: untyped): auto =
-  ## Transforms `a` into an expression that can be safely evaluated even in
-  ## presence of intermediate nil pointers/references, in which case a default
-  ## value is produced.
-  result = replace(a)
-  result = quote do:
-    # `result`.val # TODO: expose a way to do this directly in std/options, e.g.: `getAsIs`
-    safeGet(`result`)
-
-macro `??.`*(a: untyped): Option =
-  ## Same as `?.` but returns an `Option`.
-  runnableExamples:
-    import std/options
-    type Foo = ref object
-      x1: ref int
-      x2: int
-    # `?.` can't distinguish between a valid vs invalid default value, but `??.` can:
-    var f1 = Foo(x1: int.new, x2: 2)
-    doAssert (??.f1.x1[]).get == 0 # not enough to tell when the chain was valid.
-    doAssert (??.f1.x1[]).isSome # a nil didn't occur in the chain
-    doAssert (??.f1.x2).get == 2
-
-    var f2: Foo
-    doAssert not (??.f2.x1[]).isSome # f2 was nil
-    from std/options import UnpackDefect
-    doAssertRaises(UnpackDefect): discard (??.f2.x1[]).get
-    doAssert ?.f2.x1[] == 0 # in contrast, this returns default(int)
-
-  result = replace(a)
+when false:
+  # xxx: expose a way to do this directly in std/options, e.g.: `getAsIs`
+  proc safeGet[T](a: Option[T]): T {.inline.} =
+    get(a, default(T))