From ee7b6fd4c1c11641f9b1b54441ed038fd761de13 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin@nvidia.com>
Date: Fri, 17 Mar 2023 16:52:19 -0700
Subject: [PATCH] [flang] Set proper source location for the main function.

Take the source position for the anonymous program from its scope.
If the first evaluation is a construct or directive, then it has
null source position.

Author: vdonaldson
Differential Revision: https://reviews.llvm.org/D146445
---
 flang/lib/Lower/PFTBuilder.cpp     |  4 +---
 flang/test/Lower/main_location.f90 | 26 ++++++++++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)
 create mode 100644 flang/test/Lower/main_location.f90

diff --git a/flang/lib/Lower/PFTBuilder.cpp b/flang/lib/Lower/PFTBuilder.cpp
index 6261a90..2c63852 100644
--- a/flang/lib/Lower/PFTBuilder.cpp
+++ b/flang/lib/Lower/PFTBuilder.cpp
@@ -1671,9 +1671,7 @@ parser::CharBlock
 Fortran::lower::pft::FunctionLikeUnit::getStartingSourceLoc() const {
   if (beginStmt)
     return stmtSourceLoc(*beginStmt);
-  if (!evaluationList.empty())
-    return evaluationList.front().position;
-  return stmtSourceLoc(endStmt);
+  return scope->sourceRange();
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/flang/test/Lower/main_location.f90 b/flang/test/Lower/main_location.f90
new file mode 100644
index 0000000..db63339
--- /dev/null
+++ b/flang/test/Lower/main_location.f90
@@ -0,0 +1,26 @@
+! RUN: split-file %s %t
+! RUN: bbc %t/test1.f90 -o - --emit-fir --mlir-print-debuginfo | FileCheck %s --check-prefix=TEST1
+! RUN: bbc %t/test2.f90 -o - --emit-fir --mlir-print-debuginfo | FileCheck %s --check-prefix=TEST2
+
+! Check that the missing optional program-stmt (R1401)
+! does not result in unknown source location of the corresponding
+! function.
+
+!--- test1.f90
+if (.false.) then
+endif
+end
+
+! TEST1: func.func @_QQmain() {
+! TEST1-NEXT: return loc("{{.*}}test1.f90":3:1)
+! TEST1-NEXT: } loc("{{.*}}test1.f90":1:1)
+
+!--- test2.f90
+!!! keep me here
+if (.true.) then
+endif
+end program
+
+! TEST2: func.func @_QQmain() {
+! TEST2-NEXT: return loc("{{.*}}test2.f90":4:1)
+! TEST2-NEXT: } loc("{{.*}}test2.f90":2:1)
-- 
2.7.4