Re-add condition and add tests

Author: sdesmalen-arm

clang/lib/AST/Decl.cpp

@@ -5988,9 +5988,11 @@ bool clang::IsArmStreamingFunction(const FunctionDecl *FD,
     if (FD->hasAttr<ArmLocallyStreamingAttr>())
       return true;
 
-  if (const auto *FPT = FD->getType()->getAs<FunctionProtoType>())
-    if (FPT->getAArch64SMEAttributes() & FunctionType::SME_PStateSMEnabledMask)
-      return true;
+  if (const Type *Ty = FD->getType().getTypePtrOrNull())
+    if (const auto *FPT = Ty->getAs<FunctionProtoType>())
+      if (FPT->getAArch64SMEAttributes() &
+          FunctionType::SME_PStateSMEnabledMask)
+        return true;
 
   return false;
 }

clang/test/Sema/aarch64-sme-attrs-without-sve.cpp

@@ -13,7 +13,7 @@ void test_non_streaming(svint32_t *out, svint32_t *in) {
 
 // This previously led to a diagnostic that '&a' could not be used in a non-streaming function,
 // even though all functions are streaming.
-void test_both_streaming(int32_t *out) __arm_streaming{
+void test_both_streaming(int32_t *out) __arm_streaming {
   svint32_t a;
   [&a, &out]() __arm_streaming {
     a = svdup_s32(1);
@@ -29,3 +29,26 @@ void test_lambda_streaming(int32_t *out) {
   }();
 }
 
+void test_lambda_non_streaming_capture_do_nothing() __arm_streaming {
+  svint32_t a;
+  [&a] {
+    // Do nothing.
+  }();
+}
+
+// Error: Non-streaming function attempts to dereference capture:
+void test_lambda_non_streaming_capture_return_vector() __arm_streaming {
+  svint32_t a;
+  [&a] {
+    return a; // expected-error {{SVE vector type 'svint32_t' (aka '__SVInt32_t') cannot be used in a non-streaming function}}
+  }();
+}
+
+// By reference capture, only records and uses the address of `a`:
+// FIXME: This should be okay.
+void test_lambda_non_streaming_capture_return_address() __arm_streaming {
+  svint32_t a;
+  [&a] {
+    return &a; // expected-error {{SVE vector type 'svint32_t' (aka '__SVInt32_t') cannot be used in a non-streaming function}}
+  }();
+}