Fix unused declaration member resolution

changelogs/unreleased/fix__issue-562-member-resolution.yaml

@@ -0,0 +1,3 @@
+fixed:
+- Preserve read/write-referenced struct members during unused declaration
+  elimination when component roots are nested.

lib/Transforms/LLZKUnusedDeclarationEliminationPass.cpp

@@ -18,6 +18,7 @@
 #include "llzk/Util/SymbolLookup.h"
 
 #include <mlir/IR/BuiltinOps.h>
+#include <mlir/IR/SymbolTable.h>
 
 #include <llvm/ADT/SmallVector.h>
 #include <llvm/Support/Debug.h>
@@ -36,12 +37,6 @@ using namespace llzk::component;
 
 namespace {
 
-/// @brief Get the fully-qualified member symbol.
-SymbolRefAttr getFullMemberSymbol(MemberRefOpInterface op) {
-  SymbolRefAttr structSym = op.getStructType().getNameRef(); // this is fully qualified
-  return appendLeaf(structSym, op.getMemberNameAttr());
-}
-
 class PassImpl : public llzk::impl::UnusedDeclarationEliminationPassBase<PassImpl> {
   using Base = UnusedDeclarationEliminationPassBase<PassImpl>;
   using Base::Base;
@@ -73,7 +68,10 @@ class PassImpl : public llzk::impl::UnusedDeclarationEliminationPassBase<PassImp
     PassContext ctx = PassContext::populate(getOperation());
     // First, remove unused members. This may allow more structs to be removed,
     // if their final remaining uses are as types for unused members.
-    removeUnusedMembers(ctx);
+    if (failed(removeUnusedMembers(ctx))) {
+      signalPassFailure();
+      return;
+    }
 
     // Last, remove unused structs if configured
     if (removeStructs) {
@@ -85,31 +83,49 @@ class PassImpl : public llzk::impl::UnusedDeclarationEliminationPassBase<PassImp
   /// @brief Removes unused members.
   /// A member is unused if it is never read from (only written to).
   /// @param structDef
-  void removeUnusedMembers(PassContext &ctx) {
+  LogicalResult removeUnusedMembers(PassContext &ctx) {
     ModuleOp modOp = getOperation();
 
-    // Map fully-qualified member symbols -> member ops
-    DenseMap<SymbolRefAttr, MemberDefOp> members;
-    for (auto &[structDef, structSym] : ctx.structToSymbol) {
+    SymbolTableCollection symbolTables;
+
+    // Candidate members are tracked by resolved definitions instead of by the
+    // symbol text used at each read/write site, which may be relative.
+    DenseSet<MemberDefOp> members;
+    for (const auto &entry : ctx.structToSymbol) {
+      StructDefOp structDef = entry.first;
       bool notMain = !structDef.isMainComponent();
-      structDef.walk([notMain, &structSym, &members](MemberDefOp member) {
+      structDef.walk([notMain, &members](MemberDefOp member) {
         // We don't consider public members in the Main component for removal, as these are output
         // values and removing them would result in modifying the overall circuit interface.
         if (notMain || !member.hasPublicAttr()) {
-          SymbolRefAttr memberSym =
-              appendLeaf(structSym, FlatSymbolRefAttr::get(member.getSymNameAttr()));
-          members[memberSym] = member;
+          members.insert(member);
         }
       });
     }
 
     // Remove all members that are read.
-    modOp.walk([&members](MemberReadOp readm) { members.erase(getFullMemberSymbol(readm)); });
+    WalkResult readWalk = modOp.walk([&](MemberReadOp readm) {
+      FailureOr<SymbolLookupResult<MemberDefOp>> memberDef = readm.getMemberDefOp(symbolTables);
+      if (failed(memberDef)) {
+        return WalkResult::interrupt();
+      }
+
+      members.erase(memberDef->get());
+      return WalkResult::advance();
+    });
+    if (readWalk.wasInterrupted()) {
+      return failure();
+    }
 
     // Remove all writes that reference the remaining members, as these writes
     // are now known to only update write-only members.
-    modOp.walk([&members](MemberWriteOp writem) {
-      SymbolRefAttr writtenMember = getFullMemberSymbol(writem);
+    WalkResult writeWalk = modOp.walk([&](MemberWriteOp writem) {
+      FailureOr<SymbolLookupResult<MemberDefOp>> memberDef = writem.getMemberDefOp(symbolTables);
+      if (failed(memberDef)) {
+        return WalkResult::interrupt();
+      }
+
+      MemberDefOp writtenMember = memberDef->get();
       if (members.contains(writtenMember)) {
         // We need not check the users of a writem, since it produces no results.
         LLVM_DEBUG(
@@ -118,13 +134,20 @@ class PassImpl : public llzk::impl::UnusedDeclarationEliminationPassBase<PassImp
         );
         writem.erase();
       }
+
+      return WalkResult::advance();
     });
+    if (writeWalk.wasInterrupted()) {
+      return failure();
+    }
 
     // Finally, erase the remaining members.
-    for (auto &[_, memberDef] : members) {
+    for (MemberDefOp memberDef : members) {
       LLVM_DEBUG(llvm::dbgs() << "Removing member " << memberDef << '\n');
       memberDef->erase();
     }
+
+    return success();
   }
 
   /// @brief Remove unused structs by looking for any uses of the struct's fully-qualified

test/Transforms/RedundantAndUnusedElim/issue-562-unused-declaration-member-resolution.llzk

@@ -0,0 +1,54 @@
+// RUN: llzk-opt -verify-each -llzk-unused-declaration-elim %s >/dev/null
+// Regression for #562: nested component roots must not let live member references
+// lose their resolved member declarations during unused declaration elimination.
+
+module attributes {llzk.lang, llzk.main = !struct.type<@std::@C1>} {
+  module @std attributes {llzk.lang = "zirgen"} {
+    module @std attributes {llzk.lang = "zirgen"} {
+      module @risc0 {
+        struct.def @ValU32 {
+          struct.member @low : !felt.type {llzk.pub}
+          struct.member @high : !felt.type {llzk.pub}
+          function.def @compute(%arg0 : !felt.type, %arg1 : !felt.type) -> !struct.type<@risc0::@ValU32> attributes {function.allow_witness} {
+            %self = struct.new : !struct.type<@risc0::@ValU32>
+            struct.writem %self[@low] = %arg0 : !struct.type<@risc0::@ValU32>, !felt.type
+            struct.writem %self[@high] = %arg1 : !struct.type<@risc0::@ValU32>, !felt.type
+            function.return %self : !struct.type<@risc0::@ValU32>
+          }
+          function.def @constrain(%arg0 : !struct.type<@risc0::@ValU32>, %arg1 : !felt.type, %arg2 : !felt.type) attributes {function.allow_constraint} {
+            function.return
+          }
+        }
+        struct.def @Reg {
+          struct.member @reg : !felt.type {llzk.pub}
+          function.def @compute(%arg0 : !felt.type) -> !struct.type<@risc0::@Reg> attributes {function.allow_witness} {
+            %self = struct.new : !struct.type<@risc0::@Reg>
+            struct.writem %self[@reg] = %arg0 : !struct.type<@risc0::@Reg>, !felt.type
+            function.return %self : !struct.type<@risc0::@Reg>
+          }
+          function.def @constrain(%arg0 : !struct.type<@risc0::@Reg>, %arg1 : !felt.type) attributes {function.allow_constraint} {
+            %0 = struct.readm %arg0[@reg] : !struct.type<@risc0::@Reg>, !felt.type
+            constrain.eq %arg1, %0 : !felt.type, !felt.type
+            function.return
+          }
+        }
+      }
+    }
+    struct.def @C1 {
+      struct.member @z : !struct.type<@std::@risc0::@Reg> {llzk.pub}
+      function.def @compute(%arg0 : !felt.type, %arg1 : !felt.type) -> !struct.type<@C1> attributes {function.allow_witness} {
+        %self = struct.new : !struct.type<@C1>
+        %sum = felt.add %arg0, %arg1 : !felt.type, !felt.type
+        %reg = function.call @std::@risc0::@Reg::@compute(%sum) : (!felt.type) -> !struct.type<@std::@risc0::@Reg>
+        struct.writem %self[@z] = %reg : !struct.type<@C1>, !struct.type<@std::@risc0::@Reg>
+        function.return %self : !struct.type<@C1>
+      }
+      function.def @constrain(%arg0 : !struct.type<@C1>, %arg1 : !felt.type, %arg2 : !felt.type) attributes {function.allow_constraint} {
+        %z = struct.readm %arg0[@z] : !struct.type<@C1>, !struct.type<@std::@risc0::@Reg>
+        %sum = felt.add %arg1, %arg2 : !felt.type, !felt.type
+        function.call @std::@risc0::@Reg::@constrain(%z, %sum) : (!struct.type<@std::@risc0::@Reg>, !felt.type) -> ()
+        function.return
+      }
+    }
+  }
+}