[flang-commits] [flang] [flang][MIF] Adding common features related to coarray (PR #215576)

[Dan Bonachea via flang-commits]

================
@@ -346,6 +346,66 @@ mlir::Value genTerminationOperationWrapper(fir::FirOpBuilder &builder,
   return fir::AddrOfOp::create(builder, loc, funcType, symbolRef);
 }
 
+// Generates the image index relative to the initial team, regardless of which
+// team is selected. Generates a call to the `prif_initial_team_index` function
+// (analogous to `prif_image_index`) if `cosubcripts` contains at least one
+// value; otherwise, it generates a call to `prif_this_image_no_coarray` without
+// the `team` argument.
+[[maybe_unused]] static mlir::Value
+getInitialTeamIndex(fir::FirOpBuilder &builder, mlir::Location loc,
+                    mlir::Value coarrayHandle,
+                    llvm::SmallVector<mlir::Value> cosubscripts) {
+  mlir::Type boxTy = fir::BoxType::get(builder.getNoneType());
+  mlir::Type i32Ty = builder.getI32Type();
+  mlir::Type i64Ty = builder.getI64Type();
+  mlir::Type boxArrTy = genBoxedSequenceType(i64Ty);
+  mlir::Value index = builder.createTemporary(loc, i32Ty);
+
+  // If there are no subscripts, the current image index is used.
+  if (cosubscripts.size() == 0) {
+    mlir::FunctionType ftype = mlir::FunctionType::get(
+        builder.getContext(),
+        /*inputs*/ {boxTy, builder.getRefType(i32Ty)}, /*results*/ {});
+    mlir::Value teamArg = fir::AbsentOp::create(builder, loc, boxTy);
+    mlir::func::FuncOp funcOp = builder.createFunction(
+        loc, getPRIFProcName("this_image_no_coarray"), ftype);
+    llvm::SmallVector<mlir::Value> args =
+        fir::runtime::createArguments(builder, loc, ftype, teamArg, index);
+    fir::CallOp::create(builder, loc, funcOp, args);
+    return index;
+  }
+
+  mlir::FunctionType ftype = mlir::FunctionType::get(
+      builder.getContext(),
+      /*inputs*/
+      {boxTy, boxArrTy, builder.getRefType(i32Ty), builder.getRefType(i32Ty)},
+      /*results*/ {});
+  mlir::func::FuncOp funcOp =
+      builder.createFunction(loc, getPRIFProcName("initial_team_index"), ftype);
+
+  // Creation of sub
+  unsigned corank = cosubscripts.size();
+  mlir::Type indexType = builder.getIndexType();
+  mlir::Type arrayType = fir::SequenceType::get(
+      {static_cast<fir::SequenceType::Extent>(corank)}, i64Ty);
+  mlir::Value sub = builder.createTemporary(loc, arrayType);
+  mlir::Type addrType = builder.getRefType(i64Ty);
+  for (unsigned i = 0; i < corank; ++i) {
+    mlir::Value cs = builder.createConvert(loc, i64Ty, cosubscripts[i]);
+    auto index = builder.createIntegerConstant(loc, indexType, i);
+    auto addr = fir::CoordinateOp::create(builder, loc, addrType, sub, index);
----------------
bonachea wrote:

Minor: Shadowed variable declaration is confusing

```suggestion
    auto cs_index = builder.createIntegerConstant(loc, indexType, i);
    auto addr = fir::CoordinateOp::create(builder, loc, addrType, sub, cs_index);
```

https://github.com/llvm/llvm-project/pull/215576