From 0e18d5ed21c829220849bb2adcc17b2f8077bbae Mon Sep 17 00:00:00 2001
From: River Riddle <riddleriver@gmail.com>
Date: Sat, 5 Nov 2022 16:36:17 -0700
Subject: [PATCH] [mlir][SubElements] Re-add null guards to better enable
 downstream adoption

We used to allow this, and it can break clients that still rely on it.
---
 mlir/lib/IR/SubElementInterfaces.cpp | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/mlir/lib/IR/SubElementInterfaces.cpp b/mlir/lib/IR/SubElementInterfaces.cpp
index ae0223f..fd05b9d 100644
--- a/mlir/lib/IR/SubElementInterfaces.cpp
+++ b/mlir/lib/IR/SubElementInterfaces.cpp
@@ -27,6 +27,11 @@ static void walkSubElementsImpl(InterfaceT interface,
                                 DenseSet<Type> &visitedTypes) {
   interface.walkImmediateSubElements(
       [&](Attribute attr) {
+        // Guard against potentially null inputs. This removes the need for the
+        // derived attribute/type to do it.
+        if (!attr)
+          return;
+
         // Avoid infinite recursion when visiting sub attributes later, if this
         // is a mutable attribute.
         if (LLVM_UNLIKELY(attr.hasTrait<AttributeTrait::IsMutable>())) {
@@ -43,6 +48,11 @@ static void walkSubElementsImpl(InterfaceT interface,
         walkAttrsFn(attr);
       },
       [&](Type type) {
+        // Guard against potentially null inputs. This removes the need for the
+        // derived attribute/type to do it.
+        if (!type)
+          return;
+
         // Avoid infinite recursion when visiting sub types later, if this
         // is a mutable type.
         if (LLVM_UNLIKELY(type.hasTrait<TypeTrait::IsMutable>())) {
@@ -93,6 +103,10 @@ static void updateSubElementImpl(
     return;
   newElements.push_back(element);
 
+  // Guard against potentially null inputs. We always map null to null.
+  if (!element)
+    return;
+
   // Check for an existing mapping for this element, and walk it if we haven't
   // yet.
   T *mappedElement = &visited[element];
-- 
2.7.4