[mlir] Extract DictionaryAttr sort method

Enable calling the sort, as expected by getWithSorted, into static member function so that callers can get same sorting behavior.

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

mlir/include/mlir/IR/Attributes.h

@@ -303,6 +303,11 @@ public:
   bool empty() const { return size() == 0; }
   size_t size() const;
 
+  /// Sorts the NamedAttributes in the array ordered by name as expected by
+  /// getWithSorted.
+  /// Requires: uniquely named attributes.
+  static void sort(SmallVectorImpl<NamedAttribute> &array);
+
   /// Methods for supporting type inquiry through isa, cast, and dyn_cast.
   static bool kindof(unsigned kind) {
     return kind == StandardAttributes::Dictionary;

mlir/lib/IR/Attributes.cpp

@@ -101,45 +101,40 @@ static bool compareNamedAttributeWithName(const NamedAttribute &attr,
   return strncmp(attr.first.data(), name.data(), name.size()) < 0;
 }
 
-DictionaryAttr DictionaryAttr::get(ArrayRef<NamedAttribute> value,
-                                   MLIRContext *context) {
-  assert(llvm::all_of(value,
-                      [](const NamedAttribute &attr) { return attr.second; }) &&
-         "value cannot have null entries");
-
-  // We need to sort the element list to canonicalize it, but we also don't want
-  // to do a ton of work in the super common case where the element list is
-  // already sorted.
-  SmallVector<NamedAttribute, 8> storage;
+/// Helper function that does either an in place sort or sorts from source array
+/// into destination. If inPlace then storage is both the source and the
+/// destination, else value is the source and storage destination. Returns
+/// whether source was sorted.
+template <bool inPlace>
+static bool dictionaryAttrSort(ArrayRef<NamedAttribute> value,
+                               SmallVectorImpl<NamedAttribute> &storage) {
+  // Specialize for the common case.
   switch (value.size()) {
   case 0:
-    break;
   case 1:
-    // A single element is already sorted.
+    // Zero or one elements are already sorted.
     break;
   case 2:
     assert(value[0].first != value[1].first &&
            "DictionaryAttr element names must be unique");
-
-    // Don't invoke a general sort for two element case.
     if (compareNamedAttributes(&value[0], &value[1]) > 0) {
-      storage.push_back(value[1]);
-      storage.push_back(value[0]);
-      value = storage;
+      if (inPlace)
+        std::swap(storage[0], storage[1]);
+      else
+        storage.append({value[1], value[0]});
+      return true;
     }
     break;
   default:
     // Check to see they are sorted already.
-    bool isSorted = true;
-    for (unsigned i = 0, e = value.size() - 1; i != e; ++i) {
-      if (compareNamedAttributes(&value[i], &value[i + 1]) > 0) {
-        isSorted = false;
-        break;
-      }
-    }
-    // If not, do a general sort.
+    bool isSorted =
+        llvm::is_sorted(value, [](NamedAttribute l, NamedAttribute r) {
+          return compareNamedAttributes(&l, &r) < 0;
+        });
     if (!isSorted) {
-      storage.append(value.begin(), value.end());
+      // If not, do a general sort.
+      if (!inPlace)
+        storage.append(value.begin(), value.end());
       llvm::array_pod_sort(storage.begin(), storage.end(),
                            compareNamedAttributes);
       value = storage;
@@ -151,7 +146,28 @@ DictionaryAttr DictionaryAttr::get(ArrayRef<NamedAttribute> value,
                                 return l.first == r.first;
                               }) == value.end() &&
            "DictionaryAttr element names must be unique");
+    return !isSorted;
   }
+  return false;
+}
+
+/// Sorts the NamedAttributes in the array ordered by name as expected by
+/// getWithSorted.
+/// Requires: uniquely named attributes.
+void DictionaryAttr::sort(SmallVectorImpl<NamedAttribute> &array) {
+  dictionaryAttrSort</*inPlace=*/true>(array, array);
+}
+
+DictionaryAttr DictionaryAttr::get(ArrayRef<NamedAttribute> value,
+                                   MLIRContext *context) {
+  assert(llvm::all_of(value,
+                      [](const NamedAttribute &attr) { return attr.second; }) &&
+         "value cannot have null entries");
+
+  // We need to sort the element list to canonicalize it.
+  SmallVector<NamedAttribute, 8> storage;
+  if (dictionaryAttrSort</*inPlace=*/false>(value, storage))
+    value = storage;
 
   return Base::get(context, StandardAttributes::Dictionary, value);
 }