[mlir] Promote the SubElementInterfaces to a core Attribute/Type construct

This commit restructures the sub element infrastructure to be a core part
of attributes and types, instead of being relegated to an interface. This
establishes sub element walking/replacement as something "always there",
which makes it easier to rely on for correctness/etc (which various bits of
infrastructure want, such as Symbols).

Attribute/Type now have `walk` and `replace` methods directly
accessible, which provide power API for interacting with sub elements. As
part of this, a new AttrTypeWalker class is introduced that supports caching
walked attributes/types, and a friendlier API (see the simplification of symbol
walking in SymbolTable.cpp).

Differential Revision: https://reviews.llvm.org/D142272

5 files changed, 29 insertions, 73 deletions

diff --git a/mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp b/mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp
index 295c58f925e8..f39093498e09 100644
--- a/mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp
+++ b/mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp
@@ -8,7 +8,6 @@
 
 #include "LLVMTestBase.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
-#include "mlir/IR/SubElementInterfaces.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
@@ -31,33 +30,24 @@ TEST_F(LLVMIRTest, MutualReferencedSubElementTypes) {
   Type barBody[] = {LLVMPointerType::get(fooStructTy)};
   ASSERT_TRUE(succeeded(barStructTy.setBody(barBody, /*isPacked=*/false)));
 
-  auto subElementInterface = fooStructTy.dyn_cast<SubElementTypeInterface>();
-  ASSERT_TRUE(bool(subElementInterface));
   // Test if walkSubElements goes into infinite loops.
   SmallVector<Type, 4> subElementTypes;
-  subElementInterface.walkSubElements(
-      [](Attribute attr) {},
-      [&](Type type) { subElementTypes.push_back(type); });
-  // We don't record LLVMPointerType (because it's immutable), thus
-  // !llvm.ptr<struct<"bar",...>> will be visited twice.
-  ASSERT_EQ(subElementTypes.size(), 5U);
+  fooStructTy.walk([&](Type type) { subElementTypes.push_back(type); });
+  ASSERT_EQ(subElementTypes.size(), 4U);
 
-  // !llvm.ptr<struct<"bar",...>>
+  // !llvm.ptr<struct<"foo",...>>
   ASSERT_TRUE(subElementTypes[0].isa<LLVMPointerType>());
 
-  // !llvm.struct<"foo",...>
+  // !llvm.struct<"bar",...>
   auto structType = subElementTypes[1].dyn_cast<LLVMStructType>();
   ASSERT_TRUE(bool(structType));
-  ASSERT_TRUE(structType.getName().equals("foo"));
+  ASSERT_TRUE(structType.getName().equals("bar"));
 
-  // !llvm.ptr<struct<"foo",...>>
+  // !llvm.ptr<struct<"bar",...>>
   ASSERT_TRUE(subElementTypes[2].isa<LLVMPointerType>());
 
-  // !llvm.struct<"bar",...>
+  // !llvm.struct<"foo",...>
   structType = subElementTypes[3].dyn_cast<LLVMStructType>();
   ASSERT_TRUE(bool(structType));
-  ASSERT_TRUE(structType.getName().equals("bar"));
-
-  // !llvm.ptr<struct<"bar",...>>
-  ASSERT_TRUE(subElementTypes[4].isa<LLVMPointerType>());
+  ASSERT_TRUE(structType.getName().equals("foo"));
 }
diff --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp
index 13c7762563b1..7c0572ee89fe 100644
--- a/mlir/unittests/IR/AttributeTest.cpp
+++ b/mlir/unittests/IR/AttributeTest.cpp
@@ -422,4 +422,25 @@ TEST(SparseElementsAttrTest, GetZero) {
   EXPECT_TRUE(zeroStringValue.getType() == stringTy);
 }
 
+//===----------------------------------------------------------------------===//
+// SubElements
+//===----------------------------------------------------------------------===//
+
+TEST(SubElementTest, Nested) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  BoolAttr trueAttr = builder.getBoolAttr(true);
+  BoolAttr falseAttr = builder.getBoolAttr(false);
+  ArrayAttr boolArrayAttr = builder.getArrayAttr({trueAttr, falseAttr});
+  StringAttr strAttr = builder.getStringAttr("array");
+  DictionaryAttr dictAttr =
+      builder.getDictionaryAttr(builder.getNamedAttr(strAttr, boolArrayAttr));
+
+  SmallVector<Attribute> subAttrs;
+  dictAttr.walk([&](Attribute attr) { subAttrs.push_back(attr); });
+  EXPECT_EQ(llvm::ArrayRef(subAttrs),
+            ArrayRef<Attribute>(
+                {strAttr, trueAttr, falseAttr, boolArrayAttr, dictAttr}));
+}
 } // namespace
diff --git a/mlir/unittests/IR/CMakeLists.txt b/mlir/unittests/IR/CMakeLists.txt
index 9dcfb71744be..7d49283c59c3 100644
--- a/mlir/unittests/IR/CMakeLists.txt
+++ b/mlir/unittests/IR/CMakeLists.txt
@@ -8,7 +8,6 @@ add_mlir_unittest(MLIRIRTests
   OperationSupportTest.cpp
   PatternMatchTest.cpp
   ShapedTypeTest.cpp
-  SubElementInterfaceTest.cpp
   TypeTest.cpp
 
   DEPENDS
diff --git a/mlir/unittests/IR/InterfaceTest.cpp b/mlir/unittests/IR/InterfaceTest.cpp
index e77e8794d696..9b20d3c1219e 100644
--- a/mlir/unittests/IR/InterfaceTest.cpp
+++ b/mlir/unittests/IR/InterfaceTest.cpp
@@ -41,24 +41,6 @@ TEST(InterfaceTest, OpInterfaceDenseMapKey) {
   EXPECT_FALSE(opSet.contains(op3));
 }
 
-TEST(InterfaceTest, AttrInterfaceDenseMapKey) {
-  MLIRContext context;
-  context.loadDialect<test::TestDialect>();
-
-  OpBuilder builder(&context);
-
-  DenseSet<SubElementAttrInterface> attrSet;
-  auto attr1 = builder.getArrayAttr({});
-  auto attr2 = builder.getI32ArrayAttr({0});
-  auto attr3 = builder.getI32ArrayAttr({1});
-  attrSet.insert(attr1);
-  attrSet.insert(attr2);
-  attrSet.erase(attr1);
-  EXPECT_FALSE(attrSet.contains(attr1));
-  EXPECT_TRUE(attrSet.contains(attr2));
-  EXPECT_FALSE(attrSet.contains(attr3));
-}
-
 TEST(InterfaceTest, TypeInterfaceDenseMapKey) {
   MLIRContext context;
   context.loadDialect<test::TestDialect>();
diff --git a/mlir/unittests/IR/SubElementInterfaceTest.cpp b/mlir/unittests/IR/SubElementInterfaceTest.cpp
deleted file mode 100644
index ab461f4dc340..000000000000
--- a/mlir/unittests/IR/SubElementInterfaceTest.cpp
+++ /dev/null
@@ -1,36 +0,0 @@
-//===- SubElementInterfaceTest.cpp - SubElementInterface unit tests -------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/IR/Builders.h"
-#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/SubElementInterfaces.h"
-#include "gtest/gtest.h"
-#include <cstdint>
-
-using namespace mlir;
-using namespace mlir::detail;
-
-namespace {
-TEST(SubElementInterfaceTest, Nested) {
-  MLIRContext context;
-  Builder builder(&context);
-
-  BoolAttr trueAttr = builder.getBoolAttr(true);
-  BoolAttr falseAttr = builder.getBoolAttr(false);
-  ArrayAttr boolArrayAttr = builder.getArrayAttr({trueAttr, falseAttr});
-  StringAttr strAttr = builder.getStringAttr("array");
-  DictionaryAttr dictAttr =
-      builder.getDictionaryAttr(builder.getNamedAttr(strAttr, boolArrayAttr));
-
-  SmallVector<Attribute> subAttrs;
-  dictAttr.walkSubAttrs([&](Attribute attr) { subAttrs.push_back(attr); });
-  EXPECT_EQ(llvm::ArrayRef(subAttrs),
-            ArrayRef<Attribute>({strAttr, trueAttr, falseAttr, boolArrayAttr}));
-}
-
-} // namespace