[IDE] Propagate the Stmt visit failure while walking TopLevelCodeDecls

`visitTopLevelCodeDecl` ignored the `Stmt` visit returning a nullptr.
This caused the `walkToDeclPost` to run for the `TopLevelCodeDecl` and
thus an imbalance in the `RangeResolver` pre and posts (since none of
the children would have their `walkTo*Post` called).

This was originally incorrectly fixed while assuming that the
`walkTo*Post` are called regardless of whether the children were visited
or not. Those changes have been reverted - fixing an imbalance in
`ExtDecls` in `SemaAnnotator`.

Added a test case for the `ExtDecls` imbalance which can occur while an
extension is being added.

Resolves rdar://74820040

Author: bnbarham

lib/AST/ASTWalker.cpp

@@ -217,9 +217,11 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   bool visitTopLevelCodeDecl(TopLevelCodeDecl *TLCD) {
-    if (BraceStmt *S = cast_or_null<BraceStmt>(doIt(TLCD->getBody())))
+    if (BraceStmt *S = cast_or_null<BraceStmt>(doIt(TLCD->getBody()))) {
       TLCD->setBody(S);
-    return false;
+      return false;
+    }
+    return true;
   }
 
   bool visitIfConfigDecl(IfConfigDecl *ICD) {

lib/IDE/IDERequests.cpp

@@ -995,24 +995,16 @@ bool RangeResolver::walkToDeclPre(Decl *D, CharSourceRange Range) {
 }
 
 bool RangeResolver::walkToExprPost(Expr *E) {
-  if (!Impl->shouldEnter(E))
-    return true;
   Impl->leave(E);
   return !Impl->hasResult();
 }
 
 bool RangeResolver::walkToStmtPost(Stmt *S) {
-  if (!Impl->shouldEnter(S))
-    return true;
   Impl->leave(S);
   return !Impl->hasResult();
 };
 
 bool RangeResolver::walkToDeclPost(Decl *D) {
-  if (D->isImplicit())
-    return true;
-  if (!Impl->shouldEnter(D))
-    return true;
   Impl->leave(D);
   return !Impl->hasResult();
 }

lib/IDE/SourceEntityWalker.cpp

@@ -181,7 +181,9 @@ bool SemaAnnotator::walkToDeclPre(Decl *D) {
   CharSourceRange Range = (Loc.isValid()) ? CharSourceRange(Loc, NameLen)
                                           : CharSourceRange();
   bool ShouldVisitChildren = SEWalker.walkToDeclPre(D, Range);
-  if (IsExtension) {
+  // walkToDeclPost is only called when visiting children, so make sure to only
+  // push the extension decl in that case (otherwise it won't be popped)
+  if (IsExtension && ShouldVisitChildren) {
     ExtDecls.push_back(static_cast<ExtensionDecl*>(D));
   }
   return ShouldVisitChildren;
@@ -194,13 +196,6 @@ bool SemaAnnotator::walkToDeclPost(Decl *D) {
   if (shouldIgnore(D))
     return true;
 
-  // Note: This should match any early (non-cancelling) returns in
-  // `walkToDeclPre` above
-  if (auto *ICD = dyn_cast<IfConfigDecl>(D)) {
-    if (SEWalker.shouldWalkInactiveConfigRegion())
-      return true;
-  }
-
   if (isa<ExtensionDecl>(D)) {
     assert(ExtDecls.back() == D);
     ExtDecls.pop_back();

test/SourceKit/CursorInfo/cursor_nested_extension.swift

@@ -0,0 +1,18 @@
+struct Foo {}
+struct Bar {}
+
+extension Foo {
+
+extension Bar {
+}
+
+func someFunc() {}
+
+// Closing brace for extension Foo intentionally left off, user is in the
+// process of adding the extension
+
+// This check is mostly to make sure we don't crash with an assertion -
+// extensions are pushed and popped in an AST walk pre/post and they need to
+// match up.
+// RUN: %sourcekitd-test -req=cursor -pos=9:6 %s -- %s | %FileCheck %s
+// CHECK: source.lang.swift.decl.function.method.instance