[flang] Allow embox's source_box to be a !fir.box. (#148305)

In order to create temporary copies of assumed-type arrays
(e.g. for `-frepack-arrays`), we have to allow the source_box
to be a !fir.box.

This patch replaces #147618.

Author: vzakhari

flang/include/flang/Optimizer/Dialect/FIRCG/CGOps.td

@@ -60,7 +60,8 @@ def fircg_XEmboxOp : fircg_Op<"ext_embox", [AttrSizedOperandSegments]> {
       Variadic<AnyIntegerType>:$slice,
       Variadic<AnyCoordinateType>:$subcomponent,
       Variadic<AnyIntegerType>:$substr, Variadic<AnyIntegerType>:$lenParams,
-      Optional<fir_ClassType>:$sourceBox, OptionalAttr<I32Attr>:$allocator_idx);
+      Optional<BoxOrClassType>:$sourceBox,
+      OptionalAttr<I32Attr>:$allocator_idx);
   let results = (outs BoxOrClassType);
 
   let assemblyFormat = [{

flang/include/flang/Optimizer/Dialect/FIROps.td

@@ -834,15 +834,11 @@ def fir_EmboxOp : fir_Op<"embox", [NoMemoryEffect, AttrSizedOperandSegments]> {
         - allocator_idx: specify special allocator to use.
   }];
 
-  let arguments = (ins
-    AnyReferenceLike:$memref,
-    Optional<AnyShapeType>:$shape,
-    Optional<fir_SliceType>:$slice,
-    Variadic<AnyIntegerType>:$typeparams,
-    Optional<fir_ClassType>:$sourceBox,
-    OptionalAttr<AffineMapAttr>:$accessMap,
-    OptionalAttr<I32Attr>:$allocator_idx
-  );
+  let arguments = (ins AnyReferenceLike:$memref, Optional<AnyShapeType>:$shape,
+      Optional<fir_SliceType>:$slice, Variadic<AnyIntegerType>:$typeparams,
+      Optional<BoxOrClassType>:$sourceBox,
+      OptionalAttr<AffineMapAttr>:$accessMap,
+      OptionalAttr<I32Attr>:$allocator_idx);
 
   let results = (outs BoxOrClassType);
 

flang/include/flang/Optimizer/Dialect/FIRTypes.td

@@ -626,8 +626,10 @@ def AnyBoxLike : TypeConstraint<Or<[fir_BoxType.predicate,
     fir_BoxCharType.predicate, fir_BoxProcType.predicate,
     fir_ClassType.predicate]>, "any box">;
 
-def BoxOrClassType : TypeConstraint<Or<[fir_BoxType.predicate,
-    fir_ClassType.predicate]>, "box or class">;
+def BoxOrClassLike
+    : TypeConstraint<Or<[fir_BoxType.predicate, fir_ClassType.predicate]>,
+                     "box or class">;
+def BoxOrClassType : Type<BoxOrClassLike.predicate, "box or class">;
 
 def AnyRefOrBoxLike : TypeConstraint<Or<[AnyReferenceLike.predicate,
     AnyBoxLike.predicate, FunctionType.predicate]>,

flang/test/Fir/fir-ops.fir

@@ -843,6 +843,16 @@ func.func @embox_tdesc(%arg0: !fir.class<!fir.array<10x!fir.type<derived_poly{a:
   return
 }
 
+// CHECK-LABEL:   func.func @embox_with_source_box(
+// CHECK-SAME:      %[[ARG0:.*]]: !fir.heap<!fir.array<?xnone>>,
+// CHECK-SAME:      %[[ARG1:.*]]: !fir.shape<1>,
+// CHECK-SAME:      %[[ARG2:.*]]: !fir.box<!fir.array<?xnone>>) -> !fir.class<!fir.heap<!fir.array<?xnone>>> {
+func.func @embox_with_source_box(%arg0: !fir.heap<!fir.array<?xnone>>, %arg1: !fir.shape<1>, %arg2: !fir.box<!fir.array<?xnone>>) -> !fir.class<!fir.heap<!fir.array<?xnone>>> {
+// CHECK:           %[[VAL_0:.*]] = fir.embox %[[ARG0]](%[[ARG1]]) source_box %[[ARG2]] : (!fir.heap<!fir.array<?xnone>>, !fir.shape<1>, !fir.box<!fir.array<?xnone>>) -> !fir.class<!fir.heap<!fir.array<?xnone>>>
+  %0 = fir.embox %arg0(%arg1) source_box %arg2 : (!fir.heap<!fir.array<?xnone>>, !fir.shape<1>, !fir.box<!fir.array<?xnone>>) -> !fir.class<!fir.heap<!fir.array<?xnone>>>
+  return %0 : !fir.class<!fir.heap<!fir.array<?xnone>>>
+}
+
 func.func @test_fortran_var_attrs() {
   %0 = fir.alloca !fir.heap<f32> {fortran_attrs = #fir.var_attrs<allocatable>}
   %1 = fir.alloca f32 {fortran_attrs = #fir.var_attrs<asynchronous>}