From 75b512bc6a8526f0a9f71af1cbdf769aa58da3f9 Mon Sep 17 00:00:00 2001
From: metagn <metagngn@gmail.com>
Date: Sat, 16 Nov 2024 12:43:58 +0300
Subject: [PATCH] prevent codegen of inactive case fields in VM object
 constructor nodes (#24442)

fixes #17571

Objects in the VM are represented as object constructor nodes that
contain every single field, including ones in different case branches.
This is so that every field has a unique invariant index in the object
constructor that can be written to and read from. However when
converting this node back into semantic code, fields from inactive case
branches can remain in the constructor which causes bad codegen,
generating assignments to fields from other case branches.

To fix this, fields from inactive branches are now detected in
`semmacrosanity.annotateType` (called in `fixupTypeAfterEval`) and
marked to prevent the codegen of their assignments. In #24441 these
fields were excluded from the resulting node, but this causes issues
when the node is directly supposed to go back into the VM, for example
as `const` values. I don't know if this is the only case where this
happens, so I wasn't sure about how to keep that implementation working.
---
 compiler/ccgexprs.nim       |  4 +++
 compiler/semmacrosanity.nim | 64 +++++++++++++++++++++++++++++++------
 tests/vm/tcaseobj.nim       | 19 +++++++++++
 3 files changed, 78 insertions(+), 9 deletions(-)
 create mode 100644 tests/vm/tcaseobj.nim

diff --git a/compiler/ccgexprs.nim b/compiler/ccgexprs.nim
index 1c4bbd885d..f5ce39f7a5 100644
--- a/compiler/ccgexprs.nim
+++ b/compiler/ccgexprs.nim
@@ -1835,6 +1835,10 @@ proc genObjConstr(p: BProc, e: PNode, d: var TLoc) =
   discard getTypeDesc(p.module, t)
   let ty = getUniqueType(t)
   for i in 1..<e.len:
+    if nfPreventCg in e[i].flags:
+      # this is an object constructor node generated by the VM and
+      # this field is in an inactive case branch, don't generate assignment
+      continue
     var check: PNode = nil
     if e[i].len == 3 and optFieldCheck in p.options:
       check = e[i][2]
diff --git a/compiler/semmacrosanity.nim b/compiler/semmacrosanity.nim
index 67cc86a14d..f246e38fb2 100644
--- a/compiler/semmacrosanity.nim
+++ b/compiler/semmacrosanity.nim
@@ -10,9 +10,28 @@
 ## Implements type sanity checking for ASTs resulting from macros. Lots of
 ## room for improvement here.
 
-import ast, msgs, types, options
+import ast, msgs, types, options, trees, nimsets
 
-proc ithField(n: PNode, field: var int): PSym =
+type
+  FieldTracker = object
+    index: int
+    remaining: int
+    constr: PNode
+    delete: bool # to delete fields from inactive case branches
+  FieldInfo = ref object
+    sym: PSym
+    delete: bool
+
+proc caseBranchMatchesExpr(branch, matched: PNode): bool =
+  # copied from sem
+  result = false
+  for i in 0 ..< branch.len-1:
+    if branch[i].kind == nkRange:
+      if overlap(branch[i], matched): return true
+    elif exprStructuralEquivalent(branch[i], matched):
+      return true
+
+proc ithField(n: PNode, field: var FieldTracker): FieldInfo =
   result = nil
   case n.kind
   of nkRecList:
@@ -23,18 +42,42 @@ proc ithField(n: PNode, field: var int): PSym =
     if n[0].kind != nkSym: return
     result = ithField(n[0], field)
     if result != nil: return
+    # value of the discriminator field, from (index - remaining - 1 + 1):
+    # - 1 because the `ithField` call above decreased it by 1,
+    # + 1 because the constructor node has an initial type child
+    let val = field.constr[field.index - field.remaining][1]
+    var branchFound = false
     for i in 1..<n.len:
+      let previousDelete = field.delete
       case n[i].kind
-      of nkOfBranch, nkElse:
+      of nkOfBranch:
+        if branchFound or previousDelete or
+            not caseBranchMatchesExpr(n[i], val):
+          # if this is not the active case branch,
+          # mark all fields inside as deleted
+          field.delete = true
+        else:
+          branchFound = true
         result = ithField(lastSon(n[i]), field)
         if result != nil: return
+        field.delete = previousDelete
+      of nkElse:
+        if branchFound:
+          # if this is not the active case branch,
+          # mark all fields inside as deleted
+          field.delete = true
+        result = ithField(lastSon(n[i]), field)
+        if result != nil: return
+        field.delete = previousDelete
       else: discard
   of nkSym:
-    if field == 0: result = n.sym
-    else: dec(field)
+    if field.remaining == 0:
+      result = FieldInfo(sym: n.sym, delete: field.delete)
+    else:
+      dec(field.remaining)
   else: discard
 
-proc ithField(t: PType, field: var int): PSym =
+proc ithField(t: PType, field: var FieldTracker): FieldInfo =
   var base = t.baseClass
   while base != nil:
     let b = skipTypes(base, skipPtrs)
@@ -53,13 +96,16 @@ proc annotateType*(n: PNode, t: PType; conf: ConfigRef) =
     n.typ() = t
     n[0].typ() = t
     for i in 1..<n.len:
-      var j = i-1
-      let field = x.ithField(j)
+      var tracker = FieldTracker(index: i-1, remaining: i-1, constr: n, delete: false)
+      let field = x.ithField(tracker)
       if field.isNil:
         globalError conf, n.info, "invalid field at index " & $i
       else:
         internalAssert(conf, n[i].kind == nkExprColonExpr)
-        annotateType(n[i][1], field.typ, conf)
+        annotateType(n[i][1], field.sym.typ, conf)
+        if field.delete:
+          # only codegen fields from active case branches
+          incl(n[i].flags, nfPreventCg)
   of nkPar, nkTupleConstr:
     if x.kind == tyTuple:
       n.typ() = t
diff --git a/tests/vm/tcaseobj.nim b/tests/vm/tcaseobj.nim
new file mode 100644
index 0000000000..80730f45a8
--- /dev/null
+++ b/tests/vm/tcaseobj.nim
@@ -0,0 +1,19 @@
+# issue #17571
+
+import std/[macros, objectdollar]
+
+type
+  MyEnum = enum
+    F, S, T
+  Foo = object
+    case o: MyEnum
+    of F:
+      f: string
+    of S:
+      s: string
+    of T:
+      t: string
+
+let val = static(Foo(o: F, f: "foo")).f
+doAssert val == "foo"
+doAssert $static(Foo(o: F, f: "foo")) == $Foo(o: F, f: "foo")