[IDE] Allow resolving of unresolved decl refs when computing shorthand shadows

Author: ahoppen

include/swift/Sema/IDETypeChecking.h

@@ -164,6 +164,11 @@ namespace swift {
       constraints::SolutionApplicationTarget &target, bool needsPrecheck,
       llvm::function_ref<void(const constraints::Solution &)> callback);
 
+  /// Thunk around \c TypeChecker::resolveDeclRefExpr to make it available to
+  /// \c swift::ide
+  Expr *resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE, DeclContext *Context,
+                         bool replaceInvalidRefsWithErrors);
+
   LookupResult
   lookupSemanticMember(DeclContext *DC, Type ty, DeclName name);
 
@@ -339,12 +344,17 @@ namespace swift {
   /// these shorthand shadows.
   /// The first element in the pair is the implicitly declared variable and the
   /// second variable is the shadowed one.
+  /// If a \c DeclContext is passed, it is used to resolve any
+  /// \c UnresolvedDeclRef that a shorthand shadow may refer to.
   SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1>
-  getShorthandShadows(CaptureListExpr *CaptureList);
+  getShorthandShadows(CaptureListExpr *CaptureList, DeclContext *DC = nullptr);
 
   /// Same as above but for shorthand `if let foo {` syntax.
+  /// If a \c DeclContext is passed, it is used to resolve any
+  /// \c UnresolvedDeclRef that a shorthand shadow may refer to.
   SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1>
-  getShorthandShadows(LabeledConditionalStmt *CondStmt);
+  getShorthandShadows(LabeledConditionalStmt *CondStmt,
+                      DeclContext *DC = nullptr);
 }
 
 #endif

lib/IDE/IDETypeChecking.cpp

@@ -915,13 +915,22 @@ Type swift::getResultTypeOfKeypathDynamicMember(SubscriptDecl *SD) {
 }
 
 SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1>
-swift::getShorthandShadows(CaptureListExpr *CaptureList) {
+swift::getShorthandShadows(CaptureListExpr *CaptureList, DeclContext *DC) {
   SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1> Result;
   for (auto Capture : CaptureList->getCaptureList()) {
     if (Capture.PBD->getPatternList().size() != 1) {
       continue;
     }
-    auto *DRE = dyn_cast_or_null<DeclRefExpr>(Capture.PBD->getInit(0));
+    Expr *Init = Capture.PBD->getInit(0);
+    if (!Init) {
+      continue;
+    }
+    if (auto UDRE = dyn_cast<UnresolvedDeclRefExpr>(Init)) {
+      if (DC) {
+        Init = resolveDeclRefExpr(UDRE, DC, /*replaceInvalidRefsWithErrors=*/false);
+      }
+    }
+    auto *DRE = dyn_cast_or_null<DeclRefExpr>(Init);
     if (!DRE) {
       continue;
     }
@@ -946,17 +955,23 @@ swift::getShorthandShadows(CaptureListExpr *CaptureList) {
 }
 
 SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1>
-swift::getShorthandShadows(LabeledConditionalStmt *CondStmt) {
+swift::getShorthandShadows(LabeledConditionalStmt *CondStmt, DeclContext *DC) {
   SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1> Result;
   for (const StmtConditionElement &Cond : CondStmt->getCond()) {
     if (Cond.getKind() != StmtConditionElement::CK_PatternBinding) {
       continue;
     }
-    auto Init = dyn_cast<DeclRefExpr>(Cond.getInitializer());
-    if (!Init) {
+    Expr *Init = Cond.getInitializer();
+    if (auto UDRE = dyn_cast<UnresolvedDeclRefExpr>(Init)) {
+      if (DC) {
+        Init = resolveDeclRefExpr(UDRE, DC, /*replaceInvalidRefsWithErrors=*/false);
+      }
+    }
+    auto InitDeclRef = dyn_cast<DeclRefExpr>(Init);
+    if (!InitDeclRef) {
       continue;
     }
-    auto ReferencedVar = dyn_cast_or_null<VarDecl>(Init->getDecl());
+    auto ReferencedVar = dyn_cast_or_null<VarDecl>(InitDeclRef->getDecl());
     if (!ReferencedVar) {
       continue;
     }

lib/Sema/TypeChecker.cpp

@@ -551,6 +551,11 @@ bool swift::typeCheckForCodeCompletion(
                                                  callback);
 }
 
+Expr *swift::resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE, DeclContext *Context,
+                                bool replaceInvalidRefsWithErrors) {
+  return TypeChecker::resolveDeclRefExpr(UDRE, Context, replaceInvalidRefsWithErrors);
+}
+
 void TypeChecker::checkForForbiddenPrefix(ASTContext &C, DeclBaseName Name) {
   if (C.TypeCheckerOpts.DebugForbidTypecheckPrefix.empty())
     return;