Fix findAstAncestry when position is at eof (#490)

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);
 }

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();