From 4887e455462f262e389f68b1f322a632ecbfd9df Mon Sep 17 00:00:00 2001
From: Lei Zhang <antiagainst@google.com>
Date: Mon, 25 Feb 2019 04:03:02 -0800
Subject: [PATCH] [TableGen] Fix infinite loop in SubstLeaves substitution

Previously we have `auto pos = std::string::find(...) != std::string::npos` as
if condition to control substring substitution. Instead of the position for the
found substring, `pos` will be a boolean value indicating found nor not. Then
used as the replace start position, we were always replacing starting from 0 or
1. If the replaced substring also has the pattern to be matched, we'll see
an infinite loop.

PiperOrigin-RevId: 235504681
---
 mlir/lib/TableGen/Predicate.cpp    | 11 +++++++----
 mlir/test/mlir-tblgen/predicate.td | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 4 deletions(-)
 create mode 100644 mlir/test/mlir-tblgen/predicate.td

diff --git a/mlir/lib/TableGen/Predicate.cpp b/mlir/lib/TableGen/Predicate.cpp
index 20ee266b7566..c8f4066e4de2 100644
--- a/mlir/lib/TableGen/Predicate.cpp
+++ b/mlir/lib/TableGen/Predicate.cpp
@@ -150,11 +150,14 @@ static PredNode *buildPredicateTree(const tblgen::Pred &root,
     rootNode->expr = root.getCondition();
     // Apply all parent substitutions from innermost to outermost.
     for (const auto &subst : llvm::reverse(substitutions)) {
-      size_t start = 0;
-      while (auto pos =
-                 rootNode->expr.find(subst.first, start) != std::string::npos) {
+      auto pos = rootNode->expr.find(subst.first);
+      while (pos != std::string::npos) {
         rootNode->expr.replace(pos, subst.first.size(), subst.second);
-        start = pos + subst.second.size();
+        // Skip the newly inserted substring, which itself may consider the
+        // pattern to match.
+        pos += subst.second.size();
+        // Find the next possible match position.
+        pos = rootNode->expr.find(subst.first, pos);
       }
     }
     return rootNode;
diff --git a/mlir/test/mlir-tblgen/predicate.td b/mlir/test/mlir-tblgen/predicate.td
new file mode 100644
index 000000000000..347424e91d5b
--- /dev/null
+++ b/mlir/test/mlir-tblgen/predicate.td
@@ -0,0 +1,17 @@
+// RUN: mlir-tblgen -gen-op-definitions -I %S/../../include %s | FileCheck %s
+
+include "mlir/IR/op_base.td"
+
+def I32OrF32 : Type<CPred<"{0}.isInteger(32) || {0}.isF32()">,
+                    "32-bit integer or floating-point type">;
+def I32OrF32Tensor : TypedTensor<I32OrF32>;
+
+def Identity : Op<"identity", []> {
+  let arguments = (ins
+      I32OrF32Tensor:$x);
+  let results = (outs
+      I32OrF32Tensor:$y);
+}
+
+// CHECK: this->getInstruction()->getOperand(0)->getType().cast<TensorType>().getElementType().isInteger(32) ||
+// CHECK-SAME: this->getInstruction()->getOperand(0)->getType().cast<TensorType>().getElementType().isF32()