Add a method to swap the type of a function in-place

This is motivated by the need to translate function across dialect which
requires morphing their type, as well as the Toy tutorial part on interprocedural
shape inference.

The alternative is cloning the function, but it is heavy and it seems like an
arbitrary restriction to forbid morphing the function type.

PiperOrigin-RevId: 240615755

diff --git a/mlir/include/mlir/IR/Function.h b/mlir/include/mlir/IR/Function.h
index 41435ad..5b877ae 100644 (file)
--- a/mlir/include/mlir/IR/Function.h
+++ b/mlir/include/mlir/IR/Function.h
@@ -61,6 +61,18 @@ public:
   /// Return the type of this function.
   FunctionType getType() { return type; }
 
+  /// Change the type of this function in place. This is an extremely dangerous
+  /// operation and it is up to the caller to ensure that this is legal for this
+  /// function, and to restore invariants:
+  ///   - the entry block args must be updated to match the function params.
+  ///  - the arguments attributes may need an update: if the new type has less
+  ///    parameters we drop the extra attributes, if there are more parameters
+  ///    they won't have any attributes.
+  void setType(FunctionType newType) {
+    type = newType;
+    argAttrs.resize(type.getNumInputs());
+  }
+
   MLIRContext *getContext();
   Module *getModule() { return module; }