[MLIR] Add callback builder for `shape.assuming` op

[bkramer]

looks good

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_AssumingAllOp : Shape_Op<"assuming_all", [Commutative, NoSideEffect]> {
   let assemblyFormat = "$inputs attr-dict";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 
   let verifier = [{ return ::verify(*this); }];
 }
 
-def Shape_AssumingOp : Shape_Op<"assuming",
-                           [SingleBlockImplicitTerminator<"AssumingYieldOp">,
-                            DeclareOpInterfaceMethods<RegionBranchOpInterface>,
-                            RecursiveSideEffects]> {
+def Shape_AssumingOp : Shape_Op<"assuming", [
+    SingleBlockImplicitTerminator<"AssumingYieldOp">,
   let summary = "Execute the region";
   let description = [{
     Executes the region assuming all witnesses are true.
 
     "assuming" operations represent an execution order restriction to the
     compiler, information for dependent code to rely on (by assuming), and
     nothing else. They should not exist after a program is fully lowered and
     ready to execute.
def Shape_AssumingOp : Shape_Op<"assuming", [
   let extraClassDeclaration = [{
     // Inline the region into the region containing the AssumingOp and delete
     // the AssumingOp.
     //
     // This does no checks on the inputs to the AssumingOp.
     static void inlineRegionIntoParent(AssumingOp &op, PatternRewriter &rewriter);
   }];
 
+  let builders = [
+    OpBuilder<(ins "Value":$witness,
+        CArg<"function_ref<SmallVector<Value, 2>(OpBuilder &, Location)>">)>
+  ];
+
   let hasCanonicalizer = 1;
 }
 
 def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
        [NoSideEffect, ReturnLike, Terminator, HasParent<"AssumingOp">]> {
   let summary = "Yield operation";
   let description = [{
     This yield operation represents a return operation within the
     `shape.assuming` operation region. The operation takes variable number of
     operands and produces no results. The operand number and types must match
     the number and types of parent `shape.assuming` results.
   }];
 
   let arguments = (ins Variadic<AnyType>:$operands);
 
-  let builders = [OpBuilder<(ins), [{ /* nothing to do */ }]>];
+  let builders = [
+    OpBuilder<(ins), [{ /* nothing to do */ }]>,
+  ];
 
   let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
 }
 
 def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable",
                                          [Commutative, InferTypeOpInterface]> {
   let summary = "Determines if 2+ shapes can be successfully broadcasted";
   let description = [{

mlir/lib/Dialect/Shape/IR/Shape.cpp

void AssumingOp::inlineRegionIntoParent(AssumingOp &op,
   rewriter.eraseOp(&yieldOp);
 
   // Merge blocks together as there was no branching behavior from the
   // AssumingOp.
   rewriter.mergeBlocks(assumingBlock, blockBeforeAssuming);
   rewriter.mergeBlocks(blockAfterAssuming, blockBeforeAssuming);
 }
 
+void AssumingOp::build(
+    OpBuilder &builder, OperationState &result, Value witness,
+    function_ref<SmallVector<Value, 2>(OpBuilder &, Location)> bodyBuilder) {
+
+  result.addOperands(witness);
+  Region *bodyRegion = result.addRegion();
+  bodyRegion->push_back(new Block);
+  Block &bodyBlock = bodyRegion->front();
+
+  // Build body.
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(&bodyBlock);
+  SmallVector<Value, 2> yieldValues = bodyBuilder(builder, result.location);
+  builder.create<AssumingYieldOp>(result.location, yieldValues);
+
+  SmallVector<Type, 2> assumingTypes;
+  for (Value v : yieldValues)
+    assumingTypes.push_back(v.getType());
+  result.addTypes(assumingTypes);
+}
+
 //===----------------------------------------------------------------------===//
 // AssumingAllOp
 //===----------------------------------------------------------------------===//
 
 void AssumingAllOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                 MLIRContext *context) {
   patterns.add<AssumingAllOneOp>(context);
 }