[MLIR] Properly add operations to blocks during createOrFold (#70010)

Fixes #68884.

Author: mortbopet

mlir/include/mlir/IR/Builders.h

@@ -504,18 +504,20 @@ class OpBuilder : public Builder {
   template <typename OpTy, typename... Args>
   void createOrFold(SmallVectorImpl<Value> &results, Location location,
                     Args &&...args) {
-    // Create the operation without using 'create' as we don't want to
-    // insert it yet.
+    // Create the operation without using 'create' as we want to control when
+    // the listener is notified.
     OperationState state(location,
                          getCheckRegisteredInfo<OpTy>(location.getContext()));
     OpTy::build(*this, state, std::forward<Args>(args)...);
     Operation *op = Operation::create(state);
+    if (block)
+      block->getOperations().insert(insertPoint, op);
 
-    // Fold the operation. If successful destroy it, otherwise insert it.
+    // Fold the operation. If successful erase it, otherwise notify.
     if (succeeded(tryFold(op, results)))
-      op->destroy();
-    else
-      insert(op);
+      op->erase();
+    else if (listener)
+      listener->notifyOperationInserted(op);
   }
 
   /// Overload to create or fold a single result operation.

mlir/test/lib/Dialect/Test/TestDialect.cpp

@@ -529,6 +529,11 @@ LogicalResult TestOpWithVariadicResultsAndFolder::fold(
 }
 
 OpFoldResult TestOpInPlaceFold::fold(FoldAdaptor adaptor) {
+  // Exercise the fact that an operation created with createOrFold should be
+  // allowed to access its parent block.
+  assert(getOperation()->getBlock() &&
+         "expected that operation is not unlinked");
+
   if (adaptor.getOp() && !getProperties().attr) {
     // The folder adds "attr" if not present.
     getProperties().attr = dyn_cast_or_null<IntegerAttr>(adaptor.getOp());