From fce529fc6e83dc045fab1715734aecdde962ec2e Mon Sep 17 00:00:00 2001
From: Christian Sigg <csigg@google.com>
Date: Thu, 4 Nov 2021 09:05:53 +0100
Subject: [PATCH] Fix `insertFunctionArguments()` block argument order.

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D113171
---
 mlir/lib/IR/FunctionSupport.cpp        |  2 +-
 mlir/test/IR/test-func-insert-arg.mlir | 38 +++++++++++++++++-----------------
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/mlir/lib/IR/FunctionSupport.cpp b/mlir/lib/IR/FunctionSupport.cpp
index 4f6f76c..d67d277 100644
--- a/mlir/lib/IR/FunctionSupport.cpp
+++ b/mlir/lib/IR/FunctionSupport.cpp
@@ -164,7 +164,7 @@ void mlir::function_like_impl::insertFunctionArguments(
   // Update the function type and any entry block arguments.
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
   for (unsigned i = 0, e = argIndices.size(); i < e; ++i)
-    entry.insertArgument(argIndices[i], argTypes[i],
+    entry.insertArgument(argIndices[i] + i, argTypes[i],
                          argLocs.empty() ? Optional<Location>{} : argLocs[i]);
 }
 
diff --git a/mlir/test/IR/test-func-insert-arg.mlir b/mlir/test/IR/test-func-insert-arg.mlir
index 2de6c66..61308a1 100644
--- a/mlir/test/IR/test-func-insert-arg.mlir
+++ b/mlir/test/IR/test-func-insert-arg.mlir
@@ -1,49 +1,49 @@
 // RUN: mlir-opt %s -test-func-insert-arg -split-input-file | FileCheck %s
 
-// CHECK: func @f(%arg0: f32 {test.A})
+// CHECK: func @f(%arg0: i1 {test.A})
 func @f() attributes {test.insert_args = [
-  [0, f32, {test.A}]]} {
+  [0, i1, {test.A}]]} {
   return
 }
 
 // -----
 
-// CHECK: func @f(%arg0: f32 {test.A}, %arg1: f32 {test.B})
-func @f(%arg0: f32 {test.B}) attributes {test.insert_args = [
-  [0, f32, {test.A}]]} {
+// CHECK: func @f(%arg0: i1 {test.A}, %arg1: i2 {test.B})
+func @f(%arg0: i2 {test.B}) attributes {test.insert_args = [
+  [0, i1, {test.A}]]} {
   return
 }
 
 // -----
 
-// CHECK: func @f(%arg0: f32 {test.A}, %arg1: f32 {test.B})
-func @f(%arg0: f32 {test.A}) attributes {test.insert_args = [
-  [1, f32, {test.B}]]} {
+// CHECK: func @f(%arg0: i1 {test.A}, %arg1: i2 {test.B})
+func @f(%arg0: i1 {test.A}) attributes {test.insert_args = [
+  [1, i2, {test.B}]]} {
   return
 }
 
 // -----
 
-// CHECK: func @f(%arg0: f32 {test.A}, %arg1: f32 {test.B}, %arg2: f32 {test.C})
-func @f(%arg0: f32 {test.A}, %arg1: f32 {test.C}) attributes {test.insert_args = [
-  [1, f32, {test.B}]]} {
+// CHECK: func @f(%arg0: i1 {test.A}, %arg1: i2 {test.B}, %arg2: i3 {test.C})
+func @f(%arg0: i1 {test.A}, %arg1: i3 {test.C}) attributes {test.insert_args = [
+  [1, i2, {test.B}]]} {
   return
 }
 
 // -----
 
-// CHECK: func @f(%arg0: f32 {test.A}, %arg1: f32 {test.B}, %arg2: f32 {test.C})
-func @f(%arg0: f32 {test.B}) attributes {test.insert_args = [
-  [0, f32, {test.A}],
-  [1, f32, {test.C}]]} {
+// CHECK: func @f(%arg0: i1 {test.A}, %arg1: i2 {test.B}, %arg2: i3 {test.C})
+func @f(%arg0: i2 {test.B}) attributes {test.insert_args = [
+  [0, i1, {test.A}],
+  [1, i3, {test.C}]]} {
   return
 }
 
 // -----
 
-// CHECK: func @f(%arg0: f32 {test.A}, %arg1: f32 {test.B}, %arg2: f32 {test.C})
-func @f(%arg0: f32 {test.C}) attributes {test.insert_args = [
-  [0, f32, {test.A}],
-  [0, f32, {test.B}]]} {
+// CHECK: func @f(%arg0: i1 {test.A}, %arg1: i2 {test.B}, %arg2: i3 {test.C})
+func @f(%arg0: i3 {test.C}) attributes {test.insert_args = [
+  [0, i1, {test.A}],
+  [0, i2, {test.B}]]} {
   return
 }
-- 
2.7.4