[Attributor] Allow lookupAAFor to return null on invalid state

[jhuber6]

Summary

This patch adds an option to lookupAAFor that allows it to return a
nullptr if the state of the looked up attribute is invalid. This is so
future passes can use this to query other attributes with the guarantee
that they are valid.

[jhuber6]

Forgot to format

[jdoerfert]

I'd change the default and pass a non-default value in getAAFor

[jhuber6]

Changing default.

[jdoerfert]

LG, one nit

llvm/include/llvm/Transforms/IPO/Attributor.h

1140–1141

/* AllowInvalidState */

llvm/include/llvm/Transforms/IPO/Attributor.h

struct Attributor {
   /// NOTE: ForceUpdate is ignored in any stage other than the update stage.
   template <typename AAType>
   const AAType &
   getOrCreateAAFor(IRPosition IRP, const AbstractAttribute *QueryingAA,
                    DepClassTy DepClass, bool ForceUpdate = false) {
     if (!shouldPropagateCallBaseContext(IRP))
       IRP = IRP.stripCallBaseContext();
 
-    if (AAType *AAPtr = lookupAAFor<AAType>(IRP, QueryingAA, DepClass)) {
+    if (AAType *AAPtr = lookupAAFor<AAType>(IRP, QueryingAA, DepClass,
+                                            /* AllowInvalidState */ true)) {
       if (ForceUpdate && Phase == AttributorPhase::UPDATE)
         updateAA(*AAPtr);
       return *AAPtr;
     }
 
     // No matching attribute found, create one.
     // Use the static create method.
     auto &AA = AAType::createForPosition(IRP, *this);
getOrCreateAAFor(IRPosition IRP, const AbstractAttribute *QueryingAA,
     return AA;
   }
   template <typename AAType>
   const AAType &getOrCreateAAFor(const IRPosition &IRP) {
     return getOrCreateAAFor<AAType>(IRP, /* QueryingAA */ nullptr,
                                     DepClassTy::NONE);
   }
 
-  /// Return the attribute of \p AAType for \p IRP if existing. This also allows
-  /// non-AA users lookup.
+  /// Return the attribute of \p AAType for \p IRP if existing and valid. This
+  /// also allows non-AA users lookup.
   template <typename AAType>
   AAType *lookupAAFor(const IRPosition &IRP,
                       const AbstractAttribute *QueryingAA = nullptr,
-                      DepClassTy DepClass = DepClassTy::OPTIONAL) {
+                      DepClassTy DepClass = DepClassTy::OPTIONAL,
+                      bool AllowInvalidState = false) {
     static_assert(std::is_base_of<AbstractAttribute, AAType>::value,
                   "Cannot query an attribute with a type not derived from "
                   "'AbstractAttribute'!");
     // Lookup the abstract attribute of type AAType. If found, return it after
     // registering a dependence of QueryingAA on the one returned attribute.
     AbstractAttribute *AAPtr = AAMap.lookup({&AAType::ID, IRP});
     if (!AAPtr)
       return nullptr;
 
     AAType *AA = static_cast<AAType *>(AAPtr);
 
     // Do not register a dependence on an attribute with an invalid state.
     if (DepClass != DepClassTy::NONE && QueryingAA &&
         AA->getState().isValidState())
       recordDependence(*AA, const_cast<AbstractAttribute &>(*QueryingAA),
                        DepClass);
+
+    // Return nullptr if this attribute has an invalid state.
+    if (!AllowInvalidState && !AA->getState().isValidState())
+      return nullptr;
     return AA;
   }
 
   /// Explicitly record a dependence from \p FromAA to \p ToAA, that is if
   /// \p FromAA changes \p ToAA should be updated as well.
   ///
   /// This method should be used in conjunction with the `getAAFor` method and
   /// with the DepClass enum passed to the method set to None. This can