[flang] Enhance location information (#95862)

Add inclusion location information by using FusedLocation with
attribute.

More context here:
https://discourse.llvm.org/t/rfc-enhancing-location-information/79650

flang/include/flang/Optimizer/Dialect/FIRAttr.td

@@ -111,4 +111,17 @@ def fir_LowerBoundModifierAttribute : I32EnumAttr<
   let cppNamespace = "::fir";
 }
 
+def fir_LocationKind : I32EnumAttr<"LocationKind", "Flang location kind",
+  [
+    I32EnumAttrCase<"Base", 0, "base">,
+    I32EnumAttrCase<"Inclusion", 1, "inclusion">,
+  ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::fir";
+}
+def fir_LocationKindAttr : EnumAttr<FIROpsDialect, fir_LocationKind, "loc_kind">;
+
+def LocationKindArrayAttr : ArrayOfAttr<FIROpsDialect, "LocationKindArray",
+    "loc_kind_array", "LocationKindAttr">;
+
 #endif // FIR_DIALECT_FIR_ATTRS

flang/include/flang/Parser/provenance.h

@@ -172,6 +172,8 @@ public:
   }
   void setShowColors(bool showColors) { showColors_ = showColors; }
   bool getShowColors() const { return showColors_; }
+  std::optional<ProvenanceRange> GetInclusionInfo(
+      const std::optional<ProvenanceRange> &) const;
   void EmitMessage(llvm::raw_ostream &, const std::optional<ProvenanceRange> &,
       const std::string &message, const std::string &prefix,
       llvm::raw_ostream::Colors color, bool echoSourceLine = false) const;

flang/lib/Lower/Bridge.cpp

@@ -939,24 +939,59 @@ public:
     return mlir::UnknownLoc::get(&getMLIRContext());
   }
 
+  static mlir::Location genLocation(Fortran::parser::SourcePosition pos,
+                                    mlir::MLIRContext &ctx) {
+    llvm::SmallString<256> path(*pos.path);
+    llvm::sys::fs::make_absolute(path);
+    llvm::sys::path::remove_dots(path);
+    return mlir::FileLineColLoc::get(&ctx, path.str(), pos.line, pos.column);
+  }
+
   /// Generate a `Location` from the `CharBlock`.
   mlir::Location
   genLocation(const Fortran::parser::CharBlock &block) override final {
+    mlir::Location mainLocation = genUnknownLocation();
     if (const Fortran::parser::AllCookedSources *cooked =
             bridge.getCookedSource()) {
       if (std::optional<Fortran::parser::ProvenanceRange> provenance =
               cooked->GetProvenanceRange(block)) {
         if (std::optional<Fortran::parser::SourcePosition> filePos =
-                cooked->allSources().GetSourcePosition(provenance->start())) {
-          llvm::SmallString<256> filePath(*filePos->path);
-          llvm::sys::fs::make_absolute(filePath);
-          llvm::sys::path::remove_dots(filePath);
-          return mlir::FileLineColLoc::get(&getMLIRContext(), filePath.str(),
-                                           filePos->line, filePos->column);
+                cooked->allSources().GetSourcePosition(provenance->start()))
+          mainLocation = genLocation(*filePos, getMLIRContext());
+
+        llvm::SmallVector<mlir::Location> locs;
+        locs.push_back(mainLocation);
+
+        llvm::SmallVector<fir::LocationKindAttr> locAttrs;
+        locAttrs.push_back(fir::LocationKindAttr::get(&getMLIRContext(),
+                                                      fir::LocationKind::Base));
+
+        // Gather include location information if any.
+        Fortran::parser::ProvenanceRange *prov = &*provenance;
+        while (prov) {
+          if (std::optional<Fortran::parser::ProvenanceRange> include =
+                  cooked->allSources().GetInclusionInfo(*prov)) {
+            if (std::optional<Fortran::parser::SourcePosition> incPos =
+                    cooked->allSources().GetSourcePosition(include->start())) {
+              locs.push_back(genLocation(*incPos, getMLIRContext()));
+              locAttrs.push_back(fir::LocationKindAttr::get(
+                  &getMLIRContext(), fir::LocationKind::Inclusion));
+            }
+            prov = &*include;
+          } else {
+            prov = nullptr;
+          }
+        }
+        if (locs.size() > 1) {
+          assert(locs.size() == locAttrs.size() &&
+                 "expect as many attributes as locations");
+          return mlir::FusedLocWith<fir::LocationKindArrayAttr>::get(
+              &getMLIRContext(), locs,
+              fir::LocationKindArrayAttr::get(&getMLIRContext(), locAttrs));
         }
       }
     }
-    return genUnknownLocation();
+    return mainLocation;
   }
 
   const Fortran::semantics::Scope &getCurrentScope() override final {

flang/lib/Optimizer/Dialect/FIRAttr.cpp

@@ -298,5 +298,6 @@ void fir::printFirAttribute(FIROpsDialect *dialect, mlir::Attribute attr,
 void FIROpsDialect::registerAttributes() {
   addAttributes<ClosedIntervalAttr, ExactTypeAttr, FortranVariableFlagsAttr,
                 LowerBoundAttr, PointIntervalAttr, RealAttr, ReduceAttr,
-                SubclassAttr, UpperBoundAttr>();
+                SubclassAttr, UpperBoundAttr, LocationKindAttr,
+                LocationKindArrayAttr>();
 }

flang/lib/Optimizer/Transforms/AddDebugInfo.cpp

@@ -76,8 +76,11 @@ static uint32_t getLineFromLoc(mlir::Location loc) {
 }
 
 bool debugInfoIsAlreadySet(mlir::Location loc) {
-  if (mlir::isa<mlir::FusedLoc>(loc))
+  if (mlir::isa<mlir::FusedLoc>(loc)) {
+    if (loc->findInstanceOf<mlir::FusedLocWith<fir::LocationKindAttr>>())
+      return false;
     return true;
+  }
   return false;
 }
 

flang/lib/Parser/provenance.cpp

@@ -246,6 +246,27 @@ static void EmitPrefix(llvm::raw_ostream &o, llvm::raw_ostream::Colors color,
   }
 }
 
+std::optional<ProvenanceRange> AllSources::GetInclusionInfo(
+    const std::optional<ProvenanceRange> &range) const {
+  if (!range)
+    return std::nullopt;
+  const Origin &origin{MapToOrigin(range->start())};
+
+  return common::visit(
+      common::visitors{
+          [&](const Inclusion &inc) -> std::optional<ProvenanceRange> {
+            if (IsValid(origin.replaces) &&
+                range_.Contains(origin.replaces.start()))
+              return origin.replaces;
+            return std::nullopt;
+          },
+          [&](const auto &) -> std::optional<ProvenanceRange> {
+            return std::nullopt;
+          },
+      },
+      origin.u);
+}
+
 void AllSources::EmitMessage(llvm::raw_ostream &o,
     const std::optional<ProvenanceRange> &range, const std::string &message,
     const std::string &prefix, llvm::raw_ostream::Colors color,

flang/test/Lower/location.f90

@@ -0,0 +1,13 @@
+! RUN: bbc -emit-hlfir --mlir-print-debuginfo %s -o - | FileCheck %s
+
+program test
+include 'location0.inc'
+
+end 
+
+! CHECK-LABEL: func.func @_QQmain() attributes {fir.bindc_name = "test"} {
+! CHECK: fir.call @_FortranAioOutputAscii(%{{.*}}, %{{.*}}, %{{.*}}) fastmath<contract> : (!fir.ref<i8>, !fir.ref<i8>, i64) -> i1 loc(fused<#fir<loc_kind_array[ base,  inclusion,  inclusion]>>["{{.*}}location1.inc":1:10, "{{.*}}location0.inc":1:1, "{{.*}}location.f90":4:1])
+! CHECK: return loc("{{.*}}location.f90":6:1)
+! CHECK: } loc("{{.*}}location.f90":3:1)
+
+

flang/test/Lower/location0.inc

@@ -0,0 +1 @@
+include 'location1.inc'

flang/test/Lower/location1.inc

@@ -0,0 +1 @@
+print *, 'from location.inc'