[llvm] [TableGen] Fix crashes on invalid instruction patterns (PR #115600)

[Sergey Barannikov via llvm-commits]

https://github.com/s-barannikov updated https://github.com/llvm/llvm-project/pull/115600

>From 8615d94f56983f35da43e7ec0d1ac76e9bcc8aba Mon Sep 17 00:00:00 2001
From: Sergei Barannikov <barannikov88 at gmail.com>
Date: Sat, 9 Nov 2024 15:00:37 +0300
Subject: [PATCH] [TableGen] Fix crashes on invalid instruction patterns

There were out-of-bounds accesses due to invalid indexing
of instruction operand list.
---
 .../TableGen/dag-isel-invalid-patterns.td     | 42 +++++++++++++++++++
 .../TableGen/Common/CodeGenDAGPatterns.cpp    | 42 +++++++++++--------
 2 files changed, 67 insertions(+), 17 deletions(-)
 create mode 100644 llvm/test/TableGen/dag-isel-invalid-patterns.td

diff --git a/llvm/test/TableGen/dag-isel-invalid-patterns.td b/llvm/test/TableGen/dag-isel-invalid-patterns.td
new file mode 100644
index 00000000000000..04889d25be997c
--- /dev/null
+++ b/llvm/test/TableGen/dag-isel-invalid-patterns.td
@@ -0,0 +1,42 @@
+// RUN: not llvm-tblgen -gen-dag-isel -I %p/../../include %s 2>&1 >/dev/null | FileCheck %s
+
+include "llvm/Target/Target.td"
+
+def R0 : Register<"gr0">;
+def GPR : RegisterClass<"My", [i32], 32, (add R0)>;
+
+def my_node_1 : SDNode<"MYISD::NODE1", SDTypeProfile<1, 0, []>>;
+def my_node_2 : SDNode<"MYISD::NODE2", SDTypeProfile<2, 0, []>>;
+def my_node_3 : SDNode<"MYISD::NODE3", SDTypeProfile<2, 1, []>>;
+
+// CHECK: In Instr1: instruction output operand $rd does not appear in pattern outputs
+def Instr1 : Instruction {
+  let Namespace = "My";
+  let OutOperandList = (outs GPR:$rd);
+  let InOperandList = (ins GPR:$rs);
+  let Pattern = [(set GPR:$rt, (my_node_1))];
+}
+
+// CHECK: In Instr2: pattern output operand $rs does not appear in instruction outputs
+def Instr2 : Instruction {
+  let Namespace = "My";
+  let OutOperandList = (outs GPR:$rd);
+  let InOperandList = (ins);
+  let Pattern = [(set GPR:$rd, GPR:$rs, (my_node_2))];
+}
+
+// CHECK: In Instr3: pattern output operand $rs1 does not appear in instruction outputs
+def Instr3 : Instruction {
+  let Namespace = "My";
+  let OutOperandList = (outs GPR:$rd);
+  let InOperandList = (ins GPR:$rs1, GPR:$rs2);
+  let Pattern = [(set GPR:$rd, GPR:$rs1, (my_node_3 GPR:$rs2))];
+}
+
+def MyInstrInfo : InstrInfo;
+
+def My : Target {
+  let InstructionSet = MyInstrInfo;
+}
+
+// CHECK: llvm-tblgen: 3 errors.
diff --git a/llvm/utils/TableGen/Common/CodeGenDAGPatterns.cpp b/llvm/utils/TableGen/Common/CodeGenDAGPatterns.cpp
index 31bf9a98943e56..46baef6e09fc22 100644
--- a/llvm/utils/TableGen/Common/CodeGenDAGPatterns.cpp
+++ b/llvm/utils/TableGen/Common/CodeGenDAGPatterns.cpp
@@ -3803,25 +3803,21 @@ void CodeGenDAGPatterns::parseInstructionPattern(CodeGenInstruction &CGI,
   std::vector<const Record *> Results;
   std::vector<unsigned> ResultIndices;
   SmallVector<TreePatternNodePtr, 2> ResNodes;
-  for (unsigned i = 0; i != NumResults; ++i) {
-    if (i == CGI.Operands.size()) {
-      const std::string &OpName =
-          llvm::find_if(
-              InstResults,
-              [](const std::pair<std::string, TreePatternNodePtr> &P) {
-                return P.second;
-              })
-              ->first;
-
-      I.error("'" + OpName + "' set but does not appear in operand list!");
-    }
-
+  for (unsigned i = 0; i != CGI.Operands.NumDefs; ++i) {
+    const CGIOperandList::OperandInfo &Op = CGI.Operands[i];
     const std::string &OpName = CGI.Operands[i].Name;
 
     // Check that it exists in InstResults.
-    auto InstResultIter = InstResults.find(OpName);
-    if (InstResultIter == InstResults.end() || !InstResultIter->second)
-      I.error("Operand $" + OpName + " does not exist in operand list!");
+    auto InstResultIter = InstResults.find(Op.Name);
+    if (InstResultIter == InstResults.end() || !InstResultIter->second) {
+      // If this operand is optional and has a default value, ignore it.
+      if (Op.Rec->isSubClassOf("OptionalDefOperand") &&
+          !getDefaultOperand(Op.Rec).DefaultOps.empty())
+        continue;
+      I.error("instruction output operand $" + OpName +
+              " does not appear in pattern outputs");
+      continue;
+    }
 
     TreePatternNodePtr RNode = InstResultIter->second;
     const Record *R = cast<DefInit>(RNode->getLeafValue())->getDef();
@@ -3844,10 +3840,19 @@ void CodeGenDAGPatterns::parseInstructionPattern(CodeGenInstruction &CGI,
     InstResultIter->second = nullptr;
   }
 
+  auto FirstExtraResult = llvm::find_if(
+      InstResults, [](const std::pair<std::string, TreePatternNodePtr> &P) {
+        return P.second;
+      });
+  if (FirstExtraResult != InstResults.end())
+    I.error("pattern output operand $" + FirstExtraResult->first +
+            " does not appear in instruction outputs");
+
   // Loop over the inputs next.
   std::vector<TreePatternNodePtr> ResultNodeOperands;
   std::vector<const Record *> Operands;
-  for (unsigned i = NumResults, e = CGI.Operands.size(); i != e; ++i) {
+  for (unsigned i = CGI.Operands.NumDefs, e = CGI.Operands.size(); i != e;
+       ++i) {
     CGIOperandList::OperandInfo &Op = CGI.Operands[i];
     const std::string &OpName = Op.Name;
     if (OpName.empty()) {
@@ -3904,6 +3909,9 @@ void CodeGenDAGPatterns::parseInstructionPattern(CodeGenInstruction &CGI,
     I.error("Input operand $" + InstInputs.begin()->first +
             " occurs in pattern but not in operands list!");
 
+  if (I.hasError())
+    return;
+
   TreePatternNodePtr ResultPattern = makeIntrusiveRefCnt<TreePatternNode>(
       I.getRecord(), std::move(ResultNodeOperands),
       GetNumNodeResults(I.getRecord(), *this));