[NFC][MLIR] Simplify lookup of nested symbols (#179362)
Consolidate the handling of nested symbols in the loop over nested
references. This is an NFC refactor to make it simpler to enforce symbol
visibility here.
diff --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index 9f5dd2c..4f8418f 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -413,30 +413,22 @@
assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());
// Lookup the root reference for this symbol.
- symbolTableOp = lookupSymbolFn(symbolTableOp, symbol.getRootReference());
- if (!symbolTableOp)
+ auto *symbolOp = lookupSymbolFn(symbolTableOp, symbol.getRootReference());
+ if (!symbolOp)
return failure();
- symbols.push_back(symbolTableOp);
+ symbols.push_back(symbolOp);
- // If there are no nested references, just return the root symbol directly.
- ArrayRef<FlatSymbolRefAttr> nestedRefs = symbol.getNestedReferences();
- if (nestedRefs.empty())
- return success();
-
- // Verify that the root is also a symbol table.
- if (!symbolTableOp->hasTrait<OpTrait::SymbolTable>())
- return failure();
-
- // Otherwise, lookup each of the nested non-leaf references and ensure that
- // each corresponds to a valid symbol table.
- for (FlatSymbolRefAttr ref : nestedRefs.drop_back()) {
- symbolTableOp = lookupSymbolFn(symbolTableOp, ref.getAttr());
- if (!symbolTableOp || !symbolTableOp->hasTrait<OpTrait::SymbolTable>())
+ // Lookup each of the nested references.
+ for (FlatSymbolRefAttr ref : symbol.getNestedReferences()) {
+ // Check that we have a valid symbol table to lookup ref.
+ if (!symbolOp->hasTrait<OpTrait::SymbolTable>())
return failure();
- symbols.push_back(symbolTableOp);
+ symbolOp = lookupSymbolFn(symbolOp, ref.getAttr());
+ if (!symbolOp)
+ return failure();
+ symbols.push_back(symbolOp);
}
- symbols.push_back(lookupSymbolFn(symbolTableOp, symbol.getLeafReference()));
- return success(symbols.back());
+ return success();
}
LogicalResult