From 8b4c6aabc271c8bee2910c5b752759b21a6bcca9 Mon Sep 17 00:00:00 2001
From: JohnnyMorganz <johnnymorganz@outlook.com>
Date: Thu, 19 May 2022 00:26:05 +0100
Subject: [PATCH] Fix findAstAncestry when position is at eof (#490)

---
 Analysis/src/AstQuery.cpp | 20 ++++++++++++++++++--
 tests/AstQuery.test.cpp   | 13 +++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/Analysis/src/AstQuery.cpp b/Analysis/src/AstQuery.cpp
index 0aed34c0..0522b1fa 100644
--- a/Analysis/src/AstQuery.cpp
+++ b/Analysis/src/AstQuery.cpp
@@ -71,9 +71,11 @@ struct FindFullAncestry final : public AstVisitor
 {
     std::vector<AstNode*> nodes;
     Position pos;
+    Position documentEnd;
 
-    explicit FindFullAncestry(Position pos)
+    explicit FindFullAncestry(Position pos, Position documentEnd)
         : pos(pos)
+        , documentEnd(documentEnd)
     {
     }
 
@@ -84,6 +86,16 @@ struct FindFullAncestry final : public AstVisitor
             nodes.push_back(node);
             return true;
         }
+
+        // Edge case: If we ask for the node at the position that is the very end of the document
+        // return the innermost AST element that ends at that position.
+
+        if (node->location.end == documentEnd && pos >= documentEnd)
+        {
+            nodes.push_back(node);
+            return true;
+        }
+
         return false;
     }
 };
@@ -92,7 +104,11 @@ struct FindFullAncestry final : public AstVisitor
 
 std::vector<AstNode*> findAstAncestryOfPosition(const SourceModule& source, Position pos)
 {
-    FindFullAncestry finder(pos);
+    const Position end = source.root->location.end;
+    if (pos > end)
+        pos = end;
+
+    FindFullAncestry finder(pos, end);
     source.root->visit(&finder);
     return std::move(finder.nodes);
 }
diff --git a/tests/AstQuery.test.cpp b/tests/AstQuery.test.cpp
index 12c68450..f0017509 100644
--- a/tests/AstQuery.test.cpp
+++ b/tests/AstQuery.test.cpp
@@ -92,4 +92,17 @@ bar(foo())
     CHECK_EQ("number", toString(*expectedOty));
 }
 
+TEST_CASE_FIXTURE(Fixture, "ast_ancestry_at_eof")
+{
+    check(R"(
+if true then
+    )");
+
+    std::vector<AstNode*> ancestry = findAstAncestryOfPosition(*getMainSourceModule(), Position(2, 4));
+    REQUIRE_GE(ancestry.size(), 2);
+    AstStat* parentStat = ancestry[ancestry.size() - 2]->asStat();
+    REQUIRE(bool(parentStat));
+    REQUIRE(parentStat->is<AstStatIf>());
+}
+
 TEST_SUITE_END();