[flang] Avoid raising a TODO in fir.boxproc rewrite when not needed (#1560)

[clementval]

Summary

The pass was raising TODOs when a function both had a fir.boxproc<> argument
and a fir.type<> argument (even if the fir.type<> did not contain a
fir.boxproc itself).

Prevent the TODO from firing when a fir.type<> does not actually contain
a fir.boxproc. Add the location for the remaining TODO (it will be
needed when procedure pointer components are supported in lowering).

FYI, I actually tried to just implement the TODO, but I there is a funny
issue. When creating the new fir::RecordType, since the name and context
are the same as the type being translated, fir::RecordType:get just
returns the existing type, and there is no way to change it (finalize()
does nothing since it is already finalized). So this will require to add
the ability to mutate the existing type, and I am not sure what are the
MLIR constraints here, so I escaped and left the TODO for that case.

This patch is part of the upstreaming effort from fir-dev branch.

Co-authored-by: Jean Perier <[e-mail]>

[jeanPerier]

Thanks

[PeteSteinfeld]

All builds and tests correctly and looks good.

flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp

bool needsConversion(mlir::Type ty) {
     }
     if (auto tupleTy = ty.dyn_cast<mlir::TupleType>()) {
       for (auto t : tupleTy.getTypes())
         if (needsConversion(t))
           return true;
       return false;
     }
     if (auto recTy = ty.dyn_cast<RecordType>()) {
+      if (llvm::any_of(visitedTypes,
+                       [&](mlir::Type rt) { return rt == recTy; }))
+        return false;
       bool result = false;
       visitedTypes.push_back(recTy);
       for (auto t : recTy.getTypeList()) {
-        if (llvm::any_of(visitedTypes,
-                         [&](mlir::Type rt) { return rt == recTy; }))
-          continue;
         if (needsConversion(t.second)) {
           result = true;
           break;
         }
       }
       visitedTypes.pop_back();
       return result;
     }
     if (auto boxTy = ty.dyn_cast<BoxType>())
       return needsConversion(boxTy.getEleTy());
     if (isa_ref_type(ty))
       return needsConversion(unwrapRefType(ty));
     if (auto t = ty.dyn_cast<SequenceType>())
       return needsConversion(unwrapSequenceType(ty));
     return false;
   }
 
-  BoxprocTypeRewriter() {
+  BoxprocTypeRewriter(mlir::Location location) : loc{location} {
     addConversion([](mlir::Type ty) { return ty; });
-    addConversion([](BoxProcType boxproc) { return boxproc.getEleTy(); });
+    addConversion(
+        [&](BoxProcType boxproc) { return convertType(boxproc.getEleTy()); });
     addConversion([&](mlir::TupleType tupTy) {
       llvm::SmallVector<mlir::Type> memTys;
       for (auto ty : tupTy.getTypes())
         memTys.push_back(convertType(ty));
       return mlir::TupleType::get(tupTy.getContext(), memTys);
     });
     addConversion([&](mlir::FunctionType funcTy) {
       llvm::SmallVector<mlir::Type> inTys;
BoxprocTypeRewriter(mlir::Location location) : loc{location} {
     addConversion(
         [&](HeapType ty) { return HeapType::get(convertType(ty.getEleTy())); });
     addConversion(
         [&](BoxType ty) { return BoxType::get(convertType(ty.getEleTy())); });
     addConversion([&](SequenceType ty) {
       // TODO: add ty.getLayoutMap() as needed.
       return SequenceType::get(ty.getShape(), convertType(ty.getEleTy()));
     });
-    addConversion([&](RecordType ty) {
+    addConversion([&](RecordType ty) -> mlir::Type {
+      if (!needsConversion(ty))
+        return ty;
       // FIR record types can have recursive references, so conversion is a bit
       // more complex than the other types. This conversion is not needed
       // presently, so just emit a TODO message. Need to consider the uniqued
-      // name of the record, etc.
+      // name of the record, etc. Also, fir::RecordType::get returns the
+      // existing type being translated. So finalize() will not change it, and
+      // the translation would not do anything. So the type needs to be mutated,
+      // and this might require special care to comply with MLIR infrastructure.
+
+      // TODO: this will be needed to support derived type containing procedure
+      // pointer components.
       fir::emitFatalError(
-          mlir::UnknownLoc::get(ty.getContext()),
-          "not yet implemented: record type with a boxproc type");
+          loc, "not yet implemented: record type with a boxproc type");
       return RecordType::get(ty.getContext(), "*fixme*");
     });
     addArgumentMaterialization(materializeProcedure);
     addSourceMaterialization(materializeProcedure);
     addTargetMaterialization(materializeProcedure);
   }
 
   static mlir::Value materializeProcedure(mlir::OpBuilder &builder,
                                           BoxProcType type,
                                           mlir::ValueRange inputs,
                                           mlir::Location loc) {
     assert(inputs.size() == 1);
     return builder.create<ConvertOp>(loc, unwrapRefType(type.getEleTy()),
                                      inputs[0]);
   }
 
+  void setLocation(mlir::Location location) { loc = location; }
+
 private:
   llvm::SmallVector<mlir::Type> visitedTypes;
+  mlir::Location loc;
 };
 
 /// A `boxproc` is an abstraction for a Fortran procedure reference. Typically,
 /// Fortran procedures can be referenced directly through a function pointer.
 /// However, Fortran has one-level dynamic scoping between a host procedure and
 /// its internal procedures. This allows internal procedures to directly access
 /// and modify the state of the host procedure's variables.
 ///
public:
   BoxedProcedurePass(bool useThunks) { options = {useThunks}; }
 
   inline mlir::ModuleOp getModule() { return getOperation(); }
 
   void runOnOperation() override final {
     if (options.useThunks) {
       auto *context = &getContext();
       mlir::IRRewriter rewriter(context);
-      BoxprocTypeRewriter typeConverter;
+      BoxprocTypeRewriter typeConverter(mlir::UnknownLoc::get(context));
       mlir::Dialect *firDialect = context->getLoadedDialect("fir");
       getModule().walk([&](mlir::Operation *op) {
+        typeConverter.setLocation(op->getLoc());
         if (auto addr = mlir::dyn_cast<BoxAddrOp>(op)) {
           auto ty = addr.getVal().getType();
           if (typeConverter.needsConversion(ty) ||
               ty.isa<mlir::FunctionType>()) {
             // Rewrite all `fir.box_addr` ops on values of type `!fir.boxproc`
             // or function type to be `fir.convert` ops.
             rewriter.setInsertionPoint(addr);
             rewriter.replaceOpWithNewOp<ConvertOp>(

flang/test/Fir/boxproc-2.fir

+// Test fir.boxproc type conversion in the boxed-procedure pass.
+// RUN: fir-opt --boxed-procedure %s | FileCheck %s
+
+//CHECK-LABEL:  func.func private @test1(!fir.type<a{x:i32,y:f64}>, () -> ()) -> none
+func.func private @test1(!fir.type<a{x:i32, y:f64}>, !fir.boxproc<() -> ()>) -> none
+
+//CHECK-LABEL:  func.func private @test2((!fir.type<a{x:i32,y:f64}>) -> ()) -> none
+func.func private @test2(!fir.boxproc<(!fir.type<a{x:i32, y:f64}>) -> ()>) -> none
+
+//CHECK-LABEL:  func.func private @test3(() -> !fir.type<a{x:i32,y:f64}>) -> none
+func.func private @test3(!fir.boxproc<() -> (!fir.type<a{x:i32, y:f64}>)>) -> none
+
+//CHECK-LABEL:  func.func private @test5(((i32) -> f32) -> ())
+func.func private @test5(!fir.boxproc<(!fir.boxproc<(i32) -> (f32)>) -> ()>)