Reland [AArch64][AsmParser] Directives should clear transitively implied features (#106625) (#106850)

Relands 2497739 addressing the buffer overflow caused when
dereferencing an iterator past the end of ExtensionMap.

Author: labrinea

llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp

@@ -6947,18 +6947,22 @@ static void ExpandCryptoAEK(const AArch64::ArchInfo &ArchInfo,
   }
 }
 
+static SMLoc incrementLoc(SMLoc L, int Offset) {
+  return SMLoc::getFromPointer(L.getPointer() + Offset);
+}
+
 /// parseDirectiveArch
 ///   ::= .arch token
 bool AArch64AsmParser::parseDirectiveArch(SMLoc L) {
-  SMLoc ArchLoc = getLoc();
+  SMLoc CurLoc = getLoc();
 
   StringRef Arch, ExtensionString;
   std::tie(Arch, ExtensionString) =
       getParser().parseStringToEndOfStatement().trim().split('+');
 
   const AArch64::ArchInfo *ArchInfo = AArch64::parseArch(Arch);
   if (!ArchInfo)
-    return Error(ArchLoc, "unknown arch name");
+    return Error(CurLoc, "unknown arch name");
 
   if (parseToken(AsmToken::EndOfStatement))
     return true;
@@ -6978,27 +6982,29 @@ bool AArch64AsmParser::parseDirectiveArch(SMLoc L) {
     ExtensionString.split(RequestedExtensions, '+');
 
   ExpandCryptoAEK(*ArchInfo, RequestedExtensions);
+  CurLoc = incrementLoc(CurLoc, Arch.size());
 
-  FeatureBitset Features = STI.getFeatureBits();
-  setAvailableFeatures(ComputeAvailableFeatures(Features));
   for (auto Name : RequestedExtensions) {
+    // Advance source location past '+'.
+    CurLoc = incrementLoc(CurLoc, 1);
+
     bool EnableFeature = !Name.consume_front_insensitive("no");
 
-    for (const auto &Extension : ExtensionMap) {
-      if (Extension.Name != Name)
-        continue;
+    auto It = llvm::find_if(ExtensionMap, [&Name](const auto &Extension) {
+      return Extension.Name == Name;
+    });
 
-      if (Extension.Features.none())
-        report_fatal_error("unsupported architectural extension: " + Name);
+    if (It == std::end(ExtensionMap))
+      return Error(CurLoc, "unsupported architectural extension: " + Name);
 
-      FeatureBitset ToggleFeatures =
-          EnableFeature
-              ? STI.SetFeatureBitsTransitively(~Features & Extension.Features)
-              : STI.ToggleFeature(Features & Extension.Features);
-      setAvailableFeatures(ComputeAvailableFeatures(ToggleFeatures));
-      break;
-    }
+    if (EnableFeature)
+      STI.SetFeatureBitsTransitively(It->Features);
+    else
+      STI.ClearFeatureBitsTransitively(It->Features);
+    CurLoc = incrementLoc(CurLoc, Name.size());
   }
+  FeatureBitset Features = ComputeAvailableFeatures(STI.getFeatureBits());
+  setAvailableFeatures(Features);
   return false;
 }
 
@@ -7018,28 +7024,21 @@ bool AArch64AsmParser::parseDirectiveArchExtension(SMLoc L) {
     Name = Name.substr(2);
   }
 
-  MCSubtargetInfo &STI = copySTI();
-  FeatureBitset Features = STI.getFeatureBits();
-  for (const auto &Extension : ExtensionMap) {
-    if (Extension.Name != Name)
-      continue;
-
-    if (Extension.Features.none())
-      return Error(ExtLoc, "unsupported architectural extension: " + Name);
-
-    FeatureBitset ToggleFeatures =
-        EnableFeature
-            ? STI.SetFeatureBitsTransitively(~Features & Extension.Features)
-            : STI.ToggleFeature(Features & Extension.Features);
-    setAvailableFeatures(ComputeAvailableFeatures(ToggleFeatures));
-    return false;
-  }
+  auto It = llvm::find_if(ExtensionMap, [&Name](const auto &Extension) {
+    return Extension.Name == Name;
+  });
 
-  return Error(ExtLoc, "unknown architectural extension: " + Name);
-}
+  if (It == std::end(ExtensionMap))
+    return Error(ExtLoc, "unsupported architectural extension: " + Name);
 
-static SMLoc incrementLoc(SMLoc L, int Offset) {
-  return SMLoc::getFromPointer(L.getPointer() + Offset);
+  MCSubtargetInfo &STI = copySTI();
+  if (EnableFeature)
+    STI.SetFeatureBitsTransitively(It->Features);
+  else
+    STI.ClearFeatureBitsTransitively(It->Features);
+  FeatureBitset Features = ComputeAvailableFeatures(STI.getFeatureBits());
+  setAvailableFeatures(Features);
+  return false;
 }
 
 /// parseDirectiveCPU
@@ -7075,30 +7074,21 @@ bool AArch64AsmParser::parseDirectiveCPU(SMLoc L) {
 
     bool EnableFeature = !Name.consume_front_insensitive("no");
 
-    bool FoundExtension = false;
-    for (const auto &Extension : ExtensionMap) {
-      if (Extension.Name != Name)
-        continue;
+    auto It = llvm::find_if(ExtensionMap, [&Name](const auto &Extension) {
+      return Extension.Name == Name;
+    });
 
-      if (Extension.Features.none())
-        report_fatal_error("unsupported architectural extension: " + Name);
-
-      FeatureBitset Features = STI.getFeatureBits();
-      FeatureBitset ToggleFeatures =
-          EnableFeature
-              ? STI.SetFeatureBitsTransitively(~Features & Extension.Features)
-              : STI.ToggleFeature(Features & Extension.Features);
-      setAvailableFeatures(ComputeAvailableFeatures(ToggleFeatures));
-      FoundExtension = true;
-
-      break;
-    }
-
-    if (!FoundExtension)
-      Error(CurLoc, "unsupported architectural extension");
+    if (It == std::end(ExtensionMap))
+      return Error(CurLoc, "unsupported architectural extension: " + Name);
 
+    if (EnableFeature)
+      STI.SetFeatureBitsTransitively(It->Features);
+    else
+      STI.ClearFeatureBitsTransitively(It->Features);
     CurLoc = incrementLoc(CurLoc, Name.size());
   }
+  FeatureBitset Features = ComputeAvailableFeatures(STI.getFeatureBits());
+  setAvailableFeatures(Features);
   return false;
 }
 

llvm/test/MC/AArch64/SVE/directive-arch-negative.s

@@ -0,0 +1,8 @@
+// RUN: not llvm-mc -triple aarch64 -filetype asm -o - %s 2>&1 | FileCheck %s
+
+// Check that setting +nosve implies +nosve2
+.arch armv9-a+nosve
+
+adclb z0.s, z1.s, z31.s
+// CHECK: error: instruction requires: sve2
+// CHECK-NEXT: adclb z0.s, z1.s, z31.s

llvm/test/MC/AArch64/SVE/directive-arch_extension-negative.s

@@ -1,7 +1,12 @@
 // RUN: not llvm-mc -triple aarch64 -filetype asm -o - %s 2>&1 | FileCheck %s
 
-.arch_extension nosve
+.arch_extension sve2+nosve
 
 ptrue   p0.b, pow2
 // CHECK: error: instruction requires: sve or sme
 // CHECK-NEXT: ptrue   p0.b, pow2
+
+// Check that setting +nosve implies +nosve2
+adclb z0.s, z1.s, z31.s
+// CHECK: error: instruction requires: sve2
+// CHECK-NEXT: adclb z0.s, z1.s, z31.s

llvm/test/MC/AArch64/SVE/directive-cpu-negative.s

@@ -1,6 +1,11 @@
 // RUN: not llvm-mc -triple aarch64 -filetype asm -o - %s 2>&1 | FileCheck %s
 
-.cpu generic+sve+nosve
+.cpu generic+sve2+nosve
 ptrue   p0.b, pow2
 // CHECK: error: instruction requires: sve or sme
 // CHECK-NEXT: ptrue   p0.b, pow2
+
+// Check that setting +nosve implies +nosve2
+adclb z0.s, z1.s, z31.s
+// CHECK: error: instruction requires: sve2
+// CHECK-NEXT: adclb z0.s, z1.s, z31.s

llvm/test/MC/AArch64/directive-arch-negative.s

@@ -12,10 +12,13 @@
 # CHECK-NEXT: 	aese v0.8h, v1.8h
 # CHECK-NEXT:	^
 
-// We silently ignore invalid features.
-	.arch armv8+foo
+	.arch armv8+foo+nobar
 	aese v0.8h, v1.8h
 
+# CHECK: error: unsupported architectural extension: foo
+# CHECK-NEXT:   .arch armv8+foo+nobar
+# CHECK-NEXT:               ^
+
 # CHECK: error: invalid operand for instruction
 # CHECK-NEXT:	aese v0.8h, v1.8h
 # CHECK-NEXT:	^

llvm/test/MC/AArch64/directive-arch_extension-negative.s

@@ -4,7 +4,7 @@
 // RUN: -filetype asm -o - %s 2>&1 | FileCheck %s
 
 .arch_extension axp64
-// CHECK: error: unknown architectural extension: axp64
+// CHECK: error: unsupported architectural extension: axp64
 // CHECK-NEXT: .arch_extension axp64
 
 crc32cx w0, w1, x3
@@ -49,6 +49,8 @@ fminnm d0, d0, d1
 // CHECK: [[@LINE-1]]:1: error: instruction requires: fp
 // CHECK-NEXT: fminnm d0, d0, d1
 
+// nofp implied nosimd, so reinstate it
+.arch_extension simd
 addp v0.4s, v0.4s, v0.4s
 // CHECK-NOT: [[@LINE-1]]:1: error: instruction requires: neon
 .arch_extension nosimd
@@ -70,6 +72,8 @@ casa w5, w7, [x20]
 // CHECK: [[@LINE-1]]:1: error: instruction requires: lse
 // CHECK-NEXT: casa w5, w7, [x20]
 
+// nolse implied nolse128, so reinstate it
+.arch_extension lse128
 swpp x0, x2, [x3]
 // CHECK-NOT: [[@LINE-1]]:1: error: instruction requires: lse128
 .arch_extension nolse128
@@ -84,6 +88,8 @@ cfp rctx, x0
 // CHECK: [[@LINE-1]]:5: error: CFPRCTX requires: predres
 // CHECK-NEXT: cfp rctx, x0
 
+// nopredres implied nopredres2, so reinstate it
+.arch_extension predres2
 cosp rctx, x0
 // CHECK-NOT: [[@LINE-1]]:6: error: COSP requires: predres2
 .arch_extension nopredres2
@@ -133,6 +139,8 @@ ldapr x0, [x1]
 // CHECK: [[@LINE-1]]:1: error: instruction requires: rcpc
 // CHECK-NEXT: ldapr x0, [x1]
 
+// norcpc implied norcpc3, so reinstate it
+.arch_extension rcpc3
 stilp w24, w0, [x16, #-8]!
 // CHECK-NOT: [[@LINE-1]]:1: error: instruction requires: rcpc3
 .arch_extension norcpc3
@@ -169,6 +177,8 @@ cpyfp [x0]!, [x1]!, x2!
 // CHECK: [[@LINE-1]]:1: error: instruction requires: mops
 // CHECK-NEXT: cpyfp [x0]!, [x1]!, x2!
 
+// nolse128 implied nod128, so reinstate it
+.arch_extension d128
 // This needs to come before `.arch_extension nothe` as it uses an instruction
 // that requires both the and d128
 sysp #0, c2, c0, #0, x0, x1
@@ -204,6 +214,8 @@ umax x0, x1, x2
 // CHECK: [[@LINE-1]]:1: error: instruction requires: cssc
 // CHECK-NEXT: umax x0, x1, x2
 
+// noras implied norasv2, so reinstate it
+.arch_extension rasv2
 mrs x0, ERXGSR_EL1
 // CHECK-NOT: [[@LINE-1]]:9: error: expected readable system register
 .arch_extension norasv2

llvm/test/MC/AArch64/directive-cpu-err.s

@@ -4,9 +4,10 @@
     .cpu invalid
     // CHECK: error: unknown CPU name
 
-    .cpu generic+wibble+nowobble
-    // CHECK: :[[@LINE-1]]:18: error: unsupported architectural extension
-    // CHECK: :[[@LINE-2]]:25: error: unsupported architectural extension
+    .cpu generic+fp+wibble+nowobble
+    // CHECK: error: unsupported architectural extension: wibble
+    // CHECK-NEXT: .cpu generic+fp+wibble+nowobble
+    // CHECK-NEXT:                 ^
 
     .cpu generic+nofp
     fminnm d0, d0, d1