Merge pull request #62637 from ahoppen/ahoppen/stmtcondition-at-loc

[Sema] Teach typeCheckASTNodeAtLoc to check StmtConditionElement

Author: ahoppen

lib/Sema/TypeCheckStmt.cpp

@@ -1836,6 +1836,12 @@ void StmtChecker::typeCheckASTNode(ASTNode &node) {
     return;
   }
 
+  if (auto *Cond = node.dyn_cast<StmtConditionElement *>()) {
+    bool IsFalsable; // ignored
+    TypeChecker::typeCheckStmtConditionElement(*Cond, IsFalsable, DC);
+    return;
+  }
+
   llvm_unreachable("Type checking null ASTNode");
 }
 
@@ -2163,6 +2169,19 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
   class ASTNodeFinder : public ASTWalker {
     SourceManager &SM;
     SourceLoc Loc;
+
+    /// When the \c ASTNode that we want to check was found inside a brace
+    /// statement, we need to store a *reference* to the element in the
+    /// \c BraceStmt. When the brace statement gets type checked for result
+    /// builders its elements will be updated in-place, which makes
+    /// \c FoundNodeRef now point to the type-checked replacement node. We need
+    /// this behavior.
+    ///
+    /// But for all other cases, we just want to store a plain \c ASTNode. To
+    /// make sure we free the \c ASTNode again, we store it in
+    /// \c FoundNodeStorage and set \c FoundNodeRef to point to
+    /// \c FoundNodeStorage.
+    ASTNode FoundNodeStorage;
     ASTNode *FoundNode = nullptr;
 
     /// The innermost DeclContext that contains \c FoundNode.
@@ -2230,6 +2249,20 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
         }
         // Already walked into.
         return Action::Stop();
+      } else if (auto Conditional = dyn_cast<LabeledConditionalStmt>(S)) {
+        for (StmtConditionElement &Cond : Conditional->getCond()) {
+          if (SM.isBeforeInBuffer(Loc, Cond.getStartLoc())) {
+            break;
+          }
+          SourceLoc endLoc = Lexer::getLocForEndOfToken(SM, Cond.getEndLoc());
+          if (SM.isBeforeInBuffer(endLoc, Loc) || endLoc == Loc) {
+            continue;
+          }
+
+          FoundNodeStorage = ASTNode(&Cond);
+          FoundNode = &FoundNodeStorage;
+          return Action::Stop();
+        }
       }
 
       return Action::Continue(S);
@@ -2270,7 +2303,8 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
         SourceLoc endLoc = Lexer::getLocForEndOfToken(SM, D->getEndLoc());
         if (!(SM.isBeforeInBuffer(endLoc, Loc) || endLoc == Loc)) {
           if (!isa<TopLevelCodeDecl>(D)) {
-            FoundNode = new ASTNode(D);
+            FoundNodeStorage = ASTNode(D);
+            FoundNode = &FoundNodeStorage;
           }
         }
       }

test/SourceKit/CursorInfo/cursor_on_if_let_binding.swift

@@ -0,0 +1,12 @@
+func foo() {
+  let myVar: Int? = 2
+  // RUN: %sourcekitd-test -req=cursor -pos=%(line + 1):10 %s -- %s | %FileCheck %s
+  if let myVar = myVar {
+  }
+}
+
+// CHECK: source.lang.swift.decl.var.local (4:10-4:15)
+// CHECK-NEXT: myVar
+// CHECK-NEXT: s:24cursor_on_if_let_binding3fooyyF5myVarL0_Sivp
+// CHECK-NEXT: source.lang.swift
+// CHECK-NEXT: Int