[Mlir-commits] [mlir] [mlir] Validate function insertions (PR #206009)

[llvmlistbot at llvm.org]

https://github.com/mygitljf created https://github.com/llvm/llvm-project/pull/206009

I added validation before function argument/result insertion updates the signature. Invalid insertion positions now fail cleanly with diagnostics instead of reaching lower-level type list handling and crashing.
Added focused tests for both argument and result insertion, covering out-of-range and unsorted insertion positions.
Fixes #205743

>From 4dba83c1f2fd6b12e0a530ac66cb96bdd308944b Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 26 Jun 2026 16:54:11 +0000
Subject: [PATCH] [mlir] Validate function insertions

---
 .../mlir/Interfaces/FunctionInterfaces.h        | 13 +++++++++++++
 .../mlir/Interfaces/FunctionInterfaces.td       | 11 +++++++++++
 mlir/test/IR/test-func-insert-arg-invalid.mlir  | 17 +++++++++++++++++
 .../IR/test-func-insert-result-invalid.mlir     | 13 +++++++++++++
 4 files changed, 54 insertions(+)
 create mode 100644 mlir/test/IR/test-func-insert-arg-invalid.mlir
 create mode 100644 mlir/test/IR/test-func-insert-result-invalid.mlir

diff --git a/mlir/include/mlir/Interfaces/FunctionInterfaces.h b/mlir/include/mlir/Interfaces/FunctionInterfaces.h
index e10e9bd342702..8d10157af049b 100644
--- a/mlir/include/mlir/Interfaces/FunctionInterfaces.h
+++ b/mlir/include/mlir/Interfaces/FunctionInterfaces.h
@@ -28,6 +28,19 @@ class FunctionOpInterface;
 
 namespace function_interface_impl {
 
+inline bool areValidInsertionIndices(ArrayRef<unsigned> indices,
+                                     unsigned originalSize) {
+  unsigned previousIndex = 0;
+  bool isFirstIndex = true;
+  for (unsigned index : indices) {
+    if (index > originalSize || (!isFirstIndex && index < previousIndex))
+      return false;
+    previousIndex = index;
+    isFirstIndex = false;
+  }
+  return true;
+}
+
 /// Returns the dictionary attribute corresponding to the argument at 'index'.
 /// If there are no argument attributes at 'index', a null attribute is
 /// returned.
diff --git a/mlir/include/mlir/Interfaces/FunctionInterfaces.td b/mlir/include/mlir/Interfaces/FunctionInterfaces.td
index f701e828ed641..009b1fc08afa9 100644
--- a/mlir/include/mlir/Interfaces/FunctionInterfaces.td
+++ b/mlir/include/mlir/Interfaces/FunctionInterfaces.td
@@ -273,6 +273,12 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface", [
         ::llvm::ArrayRef<::mlir::DictionaryAttr> argAttrs,
         ::llvm::ArrayRef<::mlir::Location> argLocs) {
       unsigned originalNumArgs = $_op.getNumArguments();
+      if (argIndices.size() != argTypes.size() ||
+          argIndices.size() != argLocs.size() ||
+          (!argAttrs.empty() && argIndices.size() != argAttrs.size()) ||
+          !::mlir::function_interface_impl::areValidInsertionIndices(
+              argIndices, originalNumArgs))
+        return ::llvm::failure();
       ::mlir::Type newType = $_op.getTypeWithArgsAndResults(
           argIndices, argTypes, /*resultIndices=*/{}, /*resultTypes=*/{});
       if (!newType)
@@ -301,6 +307,11 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface", [
         ::mlir::TypeRange resultTypes,
         ::llvm::ArrayRef<::mlir::DictionaryAttr> resultAttrs) {
       unsigned originalNumResults = $_op.getNumResults();
+      if (resultIndices.size() != resultTypes.size() ||
+          (!resultAttrs.empty() && resultIndices.size() != resultAttrs.size()) ||
+          !::mlir::function_interface_impl::areValidInsertionIndices(
+              resultIndices, originalNumResults))
+        return ::llvm::failure();
       ::mlir::Type newType = $_op.getTypeWithArgsAndResults(
         /*argIndices=*/{}, /*argTypes=*/{}, resultIndices, resultTypes);
       if (!newType)
diff --git a/mlir/test/IR/test-func-insert-arg-invalid.mlir b/mlir/test/IR/test-func-insert-arg-invalid.mlir
new file mode 100644
index 0000000000000..d4c2769035552
--- /dev/null
+++ b/mlir/test/IR/test-func-insert-arg-invalid.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s -test-func-insert-arg -split-input-file -verify-diagnostics
+
+// expected-error @below {{failed to insert arguments}}
+func.func @f() attributes {test.insert_args = [
+  [0, i1, {test.A}],
+  [1, i2, {test.B}]]} {
+  return
+}
+
+// -----
+
+// expected-error @below {{failed to insert arguments}}
+func.func @f(%arg0: i1 {test.A}) attributes {test.insert_args = [
+  [1, i2, {test.B}],
+  [0, i3, {test.C}]]} {
+  return
+}
diff --git a/mlir/test/IR/test-func-insert-result-invalid.mlir b/mlir/test/IR/test-func-insert-result-invalid.mlir
new file mode 100644
index 0000000000000..176bd94aa1575
--- /dev/null
+++ b/mlir/test/IR/test-func-insert-result-invalid.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s -test-func-insert-result -split-input-file -verify-diagnostics
+
+// expected-error @below {{failed to insert results}}
+func.func private @f() attributes {test.insert_results = [
+  [0, f32, {test.A}],
+  [1, f32, {test.B}]]}
+
+// -----
+
+// expected-error @below {{failed to insert results}}
+func.func private @f() -> (f32 {test.A}) attributes {test.insert_results = [
+  [1, f32, {test.B}],
+  [0, f32, {test.C}]]}