From 21bb463e9639719f1aae9535825a40732eda487b Mon Sep 17 00:00:00 2001
From: Mogball <jeffniu22@gmail.com>
Date: Tue, 19 Oct 2021 00:48:24 +0000
Subject: [PATCH] [mlir] fix bugs with NamedAttrList

- `assign` with ArrayRef was calling `append`
- `assign` with empty ArrayRef was not clearing storage

Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D112043
---
 mlir/include/mlir/IR/OperationSupport.h    |  2 +-
 mlir/lib/IR/BuiltinAttributes.cpp          |  2 ++
 mlir/unittests/IR/OperationSupportTest.cpp | 44 ++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index 0ac793f..1e85b4f 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -301,7 +301,7 @@ public:
 
   /// Replaces the attributes with new list of attributes.
   void assign(ArrayRef<NamedAttribute> range) {
-    append(range.begin(), range.end());
+    assign(range.begin(), range.end());
   }
 
   bool empty() const { return attrs.empty(); }
diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index 27d851c..9bb90a0 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/IR/BuiltinAttributes.cpp
@@ -68,6 +68,8 @@ static bool dictionaryAttrSort(ArrayRef<NamedAttribute> value,
   switch (value.size()) {
   case 0:
     // Zero already sorted.
+    if (!inPlace)
+      storage.clear();
     break;
   case 1:
     // One already sorted but may need to be copied.
diff --git a/mlir/unittests/IR/OperationSupportTest.cpp b/mlir/unittests/IR/OperationSupportTest.cpp
index 3d03329..6939064 100644
--- a/mlir/unittests/IR/OperationSupportTest.cpp
+++ b/mlir/unittests/IR/OperationSupportTest.cpp
@@ -225,4 +225,48 @@ TEST(OperationFormatPrintTest, CanUseVariadicFormat) {
   ASSERT_STREQ(str.c_str(), "\"foo.bar\"() : () -> ()");
 }
 
+TEST(NamedAttrListTest, TestAppendAssign) {
+  MLIRContext ctx;
+  NamedAttrList attrs;
+  Builder b(&ctx);
+
+  attrs.append("foo", b.getStringAttr("bar"));
+  attrs.append("baz", b.getStringAttr("boo"));
+
+  {
+    auto it = attrs.begin();
+    EXPECT_EQ(it->first, b.getIdentifier("foo"));
+    EXPECT_EQ(it->second, b.getStringAttr("bar"));
+    ++it;
+    EXPECT_EQ(it->first, b.getIdentifier("baz"));
+    EXPECT_EQ(it->second, b.getStringAttr("boo"));
+  }
+
+  attrs.append("foo", b.getStringAttr("zoo"));
+  {
+    auto dup = attrs.findDuplicate();
+    ASSERT_TRUE(dup.hasValue());
+  }
+
+  SmallVector<NamedAttribute> newAttrs = {
+      b.getNamedAttr("foo", b.getStringAttr("f")),
+      b.getNamedAttr("zoo", b.getStringAttr("z")),
+  };
+  attrs.assign(newAttrs);
+
+  auto dup = attrs.findDuplicate();
+  ASSERT_FALSE(dup.hasValue());
+
+  {
+    auto it = attrs.begin();
+    EXPECT_EQ(it->first, b.getIdentifier("foo"));
+    EXPECT_EQ(it->second, b.getStringAttr("f"));
+    ++it;
+    EXPECT_EQ(it->first, b.getIdentifier("zoo"));
+    EXPECT_EQ(it->second, b.getStringAttr("z"));
+  }
+
+  attrs.assign({});
+  ASSERT_TRUE(attrs.empty());
+}
 } // end namespace
-- 
2.7.4