[mlir-c] Add RewriterBase insertion point save/restore (#206531)

Exposes `OpBuilder::saveInsertionPoint` / `restoreInsertionPoint` through the MLIR C API for `MlirRewriterBase`, continuing the buildout of the rewrite/conversion C bindings.

Assisted by: Claude

mlir/include/mlir-c/Rewrite.h

@@ -78,8 +78,8 @@ mlirRewriterBaseGetContext(MlirRewriterBase rewriter);
 //===----------------------------------------------------------------------===//
 
 // These do not include functions using Block::iterator or Region::iterator, as
-// they are not exposed by the C API yet. Similarly for methods using
-// `InsertPoint` directly.
+// they are not exposed by the C API yet. `InsertPoint` is exposed as
+// `MlirRewriterBaseInsertPoint` below.
 
 /// Reset the insertion point to no location.  Creating an operation without a
 /// set insertion point is an error, but this can still be useful when the
@@ -133,6 +133,25 @@ mlirRewriterBaseGetBlock(MlirRewriterBase rewriter);
 MLIR_CAPI_EXPORTED MlirOperation
 mlirRewriterBaseGetOperationAfterInsertion(MlirRewriterBase rewriter);
 
+/// A saved insertion point: a (block, operationAfter) pair. `operationAfter` is
+/// the operation that subsequent insertions go before. If `operationAfter` is
+/// null, the insertion point is at the end of `block`. If `block` is null, the
+/// insertion point is not set (cleared).
+typedef struct MlirRewriterBaseInsertPoint {
+  MlirBlock block;
+  MlirOperation operationAfter;
+} MlirRewriterBaseInsertPoint;
+
+/// Returns the current insertion point of the rewriter so that it can be
+/// restored later with mlirRewriterBaseRestoreInsertionPoint.
+MLIR_CAPI_EXPORTED MlirRewriterBaseInsertPoint
+mlirRewriterBaseSaveInsertionPoint(MlirRewriterBase rewriter);
+
+/// Restores a previously saved insertion point.
+MLIR_CAPI_EXPORTED void
+mlirRewriterBaseRestoreInsertionPoint(MlirRewriterBase rewriter,
+                                      MlirRewriterBaseInsertPoint insertPoint);
+
 //===----------------------------------------------------------------------===//
 /// Block and operation creation/insertion/cloning
 //===----------------------------------------------------------------------===//

mlir/lib/CAPI/Transforms/Rewrite.cpp

@@ -87,6 +87,32 @@ mlirRewriterBaseGetOperationAfterInsertion(MlirRewriterBase rewriter) {
   return wrap(std::addressof(*it));
 }
 
+MlirRewriterBaseInsertPoint
+mlirRewriterBaseSaveInsertionPoint(MlirRewriterBase rewriter) {
+  OpBuilder::InsertPoint ip = unwrap(rewriter)->saveInsertionPoint();
+  if (!ip.isSet())
+    return {{nullptr}, {nullptr}};
+  Block *block = ip.getBlock();
+  MlirOperation operationAfter = ip.getPoint() == block->end()
+                                     ? MlirOperation{nullptr}
+                                     : wrap(&*ip.getPoint());
+  return {wrap(block), operationAfter};
+}
+
+void mlirRewriterBaseRestoreInsertionPoint(
+    MlirRewriterBase rewriter, MlirRewriterBaseInsertPoint insertPoint) {
+  if (mlirBlockIsNull(insertPoint.block)) {
+    unwrap(rewriter)->clearInsertionPoint();
+    return;
+  }
+  Block *block = unwrap(insertPoint.block);
+  if (mlirOperationIsNull(insertPoint.operationAfter))
+    unwrap(rewriter)->setInsertionPointToEnd(block);
+  else
+    unwrap(rewriter)->setInsertionPoint(
+        block, Block::iterator(unwrap(insertPoint.operationAfter)));
+}
+
 //===----------------------------------------------------------------------===//
 /// Block and operation creation/insertion/cloning
 //===----------------------------------------------------------------------===//

mlir/test/CAPI/rewrite.c

@@ -623,6 +623,82 @@ void testCloneWithMapping(MlirContext ctx) {
   fprintf(stderr, "testCloneWithMapping: PASSED\n");
 }
 
+void testInsertionPointSaveRestore(MlirContext ctx) {
+  // CHECK-LABEL: @testInsertionPointSaveRestore
+  fprintf(stderr, "@testInsertionPointSaveRestore\n");
+
+  const char *moduleString = "\"dialect.op1\"() : () -> ()\n"
+                             "\"dialect.op2\"() : () -> ()\n";
+  MlirModule module =
+      mlirModuleCreateParse(ctx, mlirStringRefCreateFromCString(moduleString));
+  MlirOperation op = mlirModuleGetOperation(module);
+  MlirBlock body = mlirModuleGetBody(module);
+  MlirOperation op1 = mlirBlockGetFirstOperation(body);
+  MlirOperation op2 = mlirOperationGetNextInBlock(op1);
+
+  MlirRewriterBase rewriter = mlirIRRewriterCreate(ctx);
+
+  // Save an insertion point that points right before op2.
+  mlirRewriterBaseSetInsertionPointBefore(rewriter, op2);
+  MlirRewriterBaseInsertPoint saved =
+      mlirRewriterBaseSaveInsertionPoint(rewriter);
+  assert(!mlirBlockIsNull(saved.block));
+  assert(mlirOperationEqual(saved.operationAfter, op2));
+
+  // Move the insertion point to the end of the block. An end-of-block insertion
+  // point round-trips with a null `operationAfter`.
+  mlirRewriterBaseSetInsertionPointToEnd(rewriter, body);
+  MlirRewriterBaseInsertPoint endIp =
+      mlirRewriterBaseSaveInsertionPoint(rewriter);
+  assert(!mlirBlockIsNull(endIp.block));
+  assert(mlirOperationIsNull(endIp.operationAfter));
+
+  // Restoring the first saved insertion point makes subsequent insertions land
+  // before op2 again, not at the end where we just were.
+  mlirRewriterBaseRestoreInsertionPoint(rewriter, saved);
+  MlirOperation opRestored =
+      createOperationWithName(ctx, "dialect.op_restored");
+  mlirRewriterBaseInsert(rewriter, opRestored);
+
+  // Restoring the null-`operationAfter` point re-establishes end-of-block, even
+  // though the insertion point currently sits in the middle of the block.
+  mlirRewriterBaseRestoreInsertionPoint(rewriter, endIp);
+  assert(!mlirBlockIsNull(mlirRewriterBaseGetInsertionBlock(rewriter)));
+  assert(mlirOperationIsNull(
+      mlirRewriterBaseGetOperationAfterInsertion(rewriter)));
+  MlirOperation opEnd = createOperationWithName(ctx, "dialect.op_end");
+  mlirRewriterBaseInsert(rewriter, opEnd);
+
+  // A cleared insertion point round-trips as a null block.
+  mlirRewriterBaseClearInsertionPoint(rewriter);
+  MlirRewriterBaseInsertPoint clearedIp =
+      mlirRewriterBaseSaveInsertionPoint(rewriter);
+  assert(mlirBlockIsNull(clearedIp.block));
+  assert(mlirOperationIsNull(clearedIp.operationAfter));
+
+  // Restoring a cleared insertion point clears the current one.
+  mlirRewriterBaseSetInsertionPointToStart(rewriter, body);
+  assert(!mlirBlockIsNull(mlirRewriterBaseGetInsertionBlock(rewriter)));
+  mlirRewriterBaseRestoreInsertionPoint(rewriter, clearedIp);
+  assert(mlirBlockIsNull(mlirRewriterBaseGetInsertionBlock(rewriter)));
+
+  mlirOperationDump(op);
+  // clang-format off
+  // CHECK:      module {
+  // CHECK-NEXT:   "dialect.op1"() : () -> ()
+  // CHECK-NEXT:   %{{.*}} = "dialect.op_restored"() : () -> index
+  // CHECK-NEXT:   "dialect.op2"() : () -> ()
+  // CHECK-NEXT:   %{{.*}} = "dialect.op_end"() : () -> index
+  // CHECK-NEXT: }
+  // clang-format on
+
+  mlirIRRewriterDestroy(rewriter);
+  mlirModuleDestroy(module);
+
+  // CHECK: testInsertionPointSaveRestore: PASSED
+  fprintf(stderr, "testInsertionPointSaveRestore: PASSED\n");
+}
+
 static MlirConversionTargetLegality dynamicLegalityAlwaysLegal(MlirOperation op,
                                                                void *userData) {
   (void)op;
@@ -1899,6 +1975,7 @@ int main(void) {
   testReplaceUses(ctx);
   testGreedyRewriteDriverConfig(ctx);
   testCloneWithMapping(ctx);
+  testInsertionPointSaveRestore(ctx);
   testConversionTargetDynamicLegality(ctx);
   testTypeConverterSourceMaterialization(ctx);
   testTypeConverterTargetMaterialization(ctx);