[Mlir-commits] [llvm] [mlir] [mlir][IR] Add SymbolUserTypeInterface (PR #198435)

[Jacques Pienaar]

================
@@ -476,6 +476,58 @@ raw_ostream &mlir::operator<<(raw_ostream &os,
 // SymbolTable Trait Types
 //===----------------------------------------------------------------------===//
 
+/// Collect `type` and any nested type parameters reachable from it.
+static void collectSymbolUseTypes(Type type, SetVector<Type> &types) {
+  type.walk<WalkOrder::PreOrder>([&](Type nestedType) {
+    types.insert(nestedType);
+    return WalkResult::advance();
+  });
+}
+
+/// Collect types directly owned by `op`, including nested type parameters.
+static void collectSymbolUseTypesInOpTypes(Operation *op,
+                                           SetVector<Type> &types) {
+  for (Type type : op->getOperandTypes())
+    collectSymbolUseTypes(type, types);
+  for (Type type : op->getResultTypes())
+    collectSymbolUseTypes(type, types);
+  for (Region &region : op->getRegions()) {
+    for (Block &block : region) {
+      for (BlockArgument argument : block.getArguments())
+        collectSymbolUseTypes(argument.getType(), types);
+    }
+  }
----------------
jpienaar wrote:

Not on you, but we should perhaps have this as a utility method somewhere (like an AttrOrTypeInOpWalker thing) as I've seen this a few times now.

https://github.com/llvm/llvm-project/pull/198435