[flang][NFC] Cache derived type translation in lowering (#80179)

Derived type translation is proving expensive in modern fortran apps
with many big derived types with dozens of components and parents.

Extending the cache that prevent recursion is proving to have little
cost on apps with small derived types and significant gain (can divide
compile time by 2) on modern fortran apps.

It is legal since the cache lifetime is less than the MLIRContext
lifetime that owns the cached mlir::Type.

Doing so also exposed that the current caching was incorrect, the type
symbol is the same for kind parametrized derived types regardless of the
kind parameters. Instances with different kinds should lower to
different MLIR types. See added test.
Using the type scopes fixes the problem.

3 files changed, 25 insertions, 18 deletions

diff --git a/flang/include/flang/Lower/AbstractConverter.h b/flang/include/flang/Lower/AbstractConverter.h
index c19dcbdcdb39..796933a4eb5f 100644
--- a/flang/include/flang/Lower/AbstractConverter.h
+++ b/flang/include/flang/Lower/AbstractConverter.h
@@ -47,6 +47,7 @@ class CharBlock;
 }
 namespace semantics {
 class Symbol;
+class Scope;
 class DerivedTypeSpec;
 } // namespace semantics
 
@@ -59,7 +60,7 @@ struct Variable;
 using SomeExpr = Fortran::evaluate::Expr<Fortran::evaluate::SomeType>;
 using SymbolRef = Fortran::common::Reference<const Fortran::semantics::Symbol>;
 using TypeConstructionStack =
-    llvm::SmallVector<std::pair<const Fortran::lower::SymbolRef, mlir::Type>>;
+    llvm::DenseMap<const Fortran::semantics::Scope *, mlir::Type>;
 class StatementContext;
 
 using ExprToValueMap = llvm::DenseMap<const SomeExpr *, mlir::Value>;
diff --git a/flang/lib/Lower/ConvertType.cpp b/flang/lib/Lower/ConvertType.cpp
index 8caafb72e472..21564e8b81d7 100644
--- a/flang/lib/Lower/ConvertType.cpp
+++ b/flang/lib/Lower/ConvertType.cpp
@@ -374,19 +374,20 @@ struct TypeBuilderImpl {
   mlir::Type genDerivedType(const Fortran::semantics::DerivedTypeSpec &tySpec) {
     std::vector<std::pair<std::string, mlir::Type>> ps;
     std::vector<std::pair<std::string, mlir::Type>> cs;
-    const Fortran::semantics::Symbol &typeSymbol = tySpec.typeSymbol();
-    if (mlir::Type ty = getTypeIfDerivedAlreadyInConstruction(typeSymbol))
-      return ty;
-
     if (tySpec.IsVectorType()) {
       return genVectorType(tySpec);
     }
 
+    const Fortran::semantics::Symbol &typeSymbol = tySpec.typeSymbol();
     const Fortran::semantics::Scope &derivedScope = DEREF(tySpec.GetScope());
+    if (mlir::Type ty = getTypeIfDerivedAlreadyInConstruction(derivedScope))
+      return ty;
 
     auto rec = fir::RecordType::get(context, converter.mangleName(tySpec));
-    // Maintain the stack of types for recursive references.
-    derivedTypeInConstruction.emplace_back(typeSymbol, rec);
+    // Maintain the stack of types for recursive references and to speed-up
+    // the derived type constructions that can be expensive for derived type
+    // with dozens of components/parents (modern Fortran).
+    derivedTypeInConstruction.try_emplace(&derivedScope, rec);
 
     // Gather the record type fields.
     // (1) The data components.
@@ -446,7 +447,6 @@ struct TypeBuilderImpl {
       }
 
     rec.finalize(ps, cs);
-    popDerivedTypeInConstruction();
 
     if (!ps.empty()) {
       // TODO: this type is a PDT (parametric derived type) with length
@@ -552,16 +552,8 @@ struct TypeBuilderImpl {
   /// type `t` have type `t`. This helper returns `t` if it is already being
   /// lowered to avoid infinite loops.
   mlir::Type getTypeIfDerivedAlreadyInConstruction(
-      const Fortran::lower::SymbolRef derivedSym) const {
-    for (const auto &[sym, type] : derivedTypeInConstruction)
-      if (sym == derivedSym)
-        return type;
-    return {};
-  }
-
-  void popDerivedTypeInConstruction() {
-    assert(!derivedTypeInConstruction.empty());
-    derivedTypeInConstruction.pop_back();
+      const Fortran::semantics::Scope &derivedScope) const {
+    return derivedTypeInConstruction.lookup(&derivedScope);
   }
 
   /// Stack derived type being processed to avoid infinite loops in case of
diff --git a/flang/test/Lower/derived-types-kind-params-2.f90 b/flang/test/Lower/derived-types-kind-params-2.f90
new file mode 100644
index 000000000000..5833079901b9
--- /dev/null
+++ b/flang/test/Lower/derived-types-kind-params-2.f90
@@ -0,0 +1,14 @@
+! This is a crazy program, recursive derived types with recursive kind
+! parameters are a terrible idea if they do not converge quickly.
+
+! RUN: bbc -emit-hlfir -o - -I nw %s | FileCheck %s
+
+subroutine foo(x)
+  type t(k)
+    integer, kind :: k
+    type(t(modulo(k+1,2))), pointer :: p
+  end type
+  type(t(1)) :: x
+end subroutine
+! CHECK-LABEL: func.func @_QPfoo(
+! CHECK-SAME: !fir.ref<!fir.type<_QFfooTtK1{p:!fir.box<!fir.ptr<!fir.type<_QFfooTtK0{p:!fir.box<!fir.ptr<!fir.type<_QFfooTtK1>>>}>>>}>>