[TypeChecker] Replace invalid refs with ErrorExpr only when explicitly allowed

Add a flag to `ConstraintSystem::preCheckExpression` and subsequently
to `TypeChecker::resolveDeclRefExpr` to indicate whether it's allowed
to replace invalid member refs with `ErrorExpr`.

It is useful for diagnostics and code completion to preserve AST
in it's original state otherwise it's impossible to diagnose errors
post factum or extract `CodeCompletionExpr` when it's a child of an
invalid reference.

Author: xedin

lib/Sema/BuilderTransform.cpp

@@ -1749,7 +1749,8 @@ class PreCheckFunctionBuilderApplication : public ASTWalker {
       // Suppress any diangostics which could be produced by this expression.
       DiagnosticTransaction transaction(diagEngine);
 
-      HasError |= ConstraintSystem::preCheckExpression(E, DC);
+      HasError |= ConstraintSystem::preCheckExpression(
+          E, DC, /*replaceInvalidRefsWithErrors=*/false);
       HasError |= transaction.hasDiagnostics();
 
       transaction.abort();

lib/Sema/CSFix.cpp

@@ -1563,7 +1563,8 @@ bool IgnoreInvalidFunctionBuilderBody::diagnose(const Solution &solution,
     PreCheckWalker(DeclContext *dc) : DC(dc) {}
 
     std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
-      auto hasError = ConstraintSystem::preCheckExpression(E, DC);
+      auto hasError = ConstraintSystem::preCheckExpression(
+          E, DC, /*replaceInvalidRefsWithErrors=*/true);
       return std::make_pair(false, hasError ? nullptr : E);
     }
 

lib/Sema/CSSolver.cpp

@@ -1411,7 +1411,8 @@ void ConstraintSystem::solveForCodeCompletion(
     llvm::function_ref<void(const Solution &)> callback) {
   // First, pre-check the expression, validating any types that occur in the
   // expression and folding sequence expressions.
-  if (ConstraintSystem::preCheckExpression(expr, DC))
+  if (ConstraintSystem::preCheckExpression(
+          expr, DC, /*replaceInvalidRefsWithErrors=*/true))
     return;
 
   ConstraintSystemOptions options;

lib/Sema/ConstraintSystem.h

@@ -4955,8 +4955,12 @@ class ConstraintSystem {
 public:
   /// Pre-check the expression, validating any types that occur in the
   /// expression and folding sequence expressions.
-  static bool preCheckExpression(Expr *&expr, DeclContext *dc);
-        
+  ///
+  /// \param replaceInvalidRefsWithErrors Indicates whether it's allowed
+  /// to replace any discovered invalid member references with `ErrorExpr`.
+  static bool preCheckExpression(Expr *&expr, DeclContext *dc,
+                                 bool replaceInvalidRefsWithErrors);
+
   /// Solve the system of constraints generated from provided target.
   ///
   /// \param target The target that we'll generate constraints from, which

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -553,10 +553,10 @@ TypeChecker::getTypeOfCompletionOperator(DeclContext *DC, Expr *LHS,
   // We allocate these expressions on the stack because we know they can't
   // escape and there isn't a better way to allocate scratch Expr nodes.
   UnresolvedDeclRefExpr UDRE(DeclNameRef(opName), refKind, DeclNameLoc(Loc));
-  auto *opExpr = TypeChecker::resolveDeclRefExpr(&UDRE, DC);
+  auto *opExpr = TypeChecker::resolveDeclRefExpr(
+      &UDRE, DC, /*replaceInvalidRefsWithErrors=*/true);
 
   switch (refKind) {
-
   case DeclRefKind::PostfixOperator: {
     // (postfix_unary_expr
     //   (declref_expr name=<opName>)
@@ -611,7 +611,8 @@ static Optional<Type> getTypeOfCompletionContextExpr(
                         CompletionTypeCheckKind kind,
                         Expr *&parsedExpr,
                         ConcreteDeclRef &referencedDecl) {
-  if (constraints::ConstraintSystem::preCheckExpression(parsedExpr, DC))
+  if (constraints::ConstraintSystem::preCheckExpression(
+          parsedExpr, DC, /*replaceInvalidRefsWithErrors=*/true))
     return None;
 
   switch (kind) {

lib/Sema/TypeCheckConstraints.cpp

@@ -481,11 +481,18 @@ static bool findNonMembers(ArrayRef<LookupResultEntry> lookupResults,
 /// returning the resultant expression. Context is the DeclContext used
 /// for the lookup.
 Expr *TypeChecker::resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE,
-                                      DeclContext *DC) {
+                                      DeclContext *DC,
+                                      bool replaceInvalidRefsWithErrors) {
   // Process UnresolvedDeclRefExpr by doing an unqualified lookup.
   DeclNameRef Name = UDRE->getName();
   SourceLoc Loc = UDRE->getLoc();
 
+  auto errorResult = [&]() -> Expr * {
+    if (replaceInvalidRefsWithErrors)
+      return new (DC->getASTContext()) ErrorExpr(UDRE->getSourceRange());
+    return UDRE;
+  };
+
   // Perform standard value name lookup.
   NameLookupOptions lookupOptions = defaultUnqualifiedLookupOptions;
   if (isa<AbstractFunctionDecl>(DC))
@@ -506,7 +513,7 @@ Expr *TypeChecker::resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE,
     if (diagnoseRangeOperatorMisspell(Context.Diags, UDRE) ||
         diagnoseIncDecOperator(Context.Diags, UDRE) ||
         diagnoseOperatorJuxtaposition(UDRE, DC)) {
-      return new (Context) ErrorExpr(UDRE->getSourceRange());
+      return errorResult();
     }
 
     // Try ignoring access control.
@@ -531,7 +538,7 @@ Expr *TypeChecker::resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE,
 
       // Don't try to recover here; we'll get more access-related diagnostics
       // downstream if the type of the inaccessible decl is also inaccessible.
-      return new (Context) ErrorExpr(UDRE->getSourceRange());
+      return errorResult();
     }
 
     // TODO: Name will be a compound name if it was written explicitly as
@@ -625,7 +632,7 @@ Expr *TypeChecker::resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE,
     // TODO: consider recovering from here.  We may want some way to suppress
     // downstream diagnostics, though.
 
-    return new (Context) ErrorExpr(UDRE->getSourceRange());
+    return errorResult();
   }
 
   // FIXME: Need to refactor the way we build an AST node from a lookup result!
@@ -916,6 +923,10 @@ namespace {
 
     Expr *ParentExpr;
 
+    /// Indicates whether pre-check is allowed to insert
+    /// implicit `ErrorExpr` in place of invalid references.
+    bool UseErrorExprs;
+
     /// A stack of expressions being walked, used to determine where to
     /// insert RebindSelfInConstructorExpr nodes.
     llvm::SmallVector<Expr *, 8> ExprStack;
@@ -1055,8 +1066,10 @@ namespace {
     }
 
   public:
-    PreCheckExpression(DeclContext *dc, Expr *parent)
-        : Ctx(dc->getASTContext()), DC(dc), ParentExpr(parent) {}
+    PreCheckExpression(DeclContext *dc, Expr *parent,
+                       bool replaceInvalidRefsWithErrors)
+        : Ctx(dc->getASTContext()), DC(dc), ParentExpr(parent),
+          UseErrorExprs(replaceInvalidRefsWithErrors) {}
 
     ASTContext &getASTContext() const { return Ctx; }
 
@@ -1122,7 +1135,8 @@ namespace {
       if (auto unresolved = dyn_cast<UnresolvedDeclRefExpr>(expr)) {
         TypeChecker::checkForForbiddenPrefix(
             getASTContext(), unresolved->getName().getBaseName());
-        return finish(true, TypeChecker::resolveDeclRefExpr(unresolved, DC));
+        return finish(true, TypeChecker::resolveDeclRefExpr(unresolved, DC,
+                                                            UseErrorExprs));
       }
 
       // Let's try to figure out if `InOutExpr` is out of place early
@@ -1979,8 +1993,9 @@ Expr *PreCheckExpression::simplifyTypeConstructionWithLiteralArg(Expr *E) {
 
 /// Pre-check the expression, validating any types that occur in the
 /// expression and folding sequence expressions.
-bool ConstraintSystem::preCheckExpression(Expr *&expr, DeclContext *dc) {
-  PreCheckExpression preCheck(dc, expr);
+bool ConstraintSystem::preCheckExpression(Expr *&expr, DeclContext *dc,
+                                          bool replaceInvalidRefsWithErrors) {
+  PreCheckExpression preCheck(dc, expr, replaceInvalidRefsWithErrors);
   // Perform the pre-check.
   if (auto result = expr->walk(preCheck)) {
     expr = result;
@@ -2107,7 +2122,8 @@ TypeChecker::typeCheckExpression(
 
   // First, pre-check the expression, validating any types that occur in the
   // expression and folding sequence expressions.
-  if (ConstraintSystem::preCheckExpression(expr, dc)) {
+  if (ConstraintSystem::preCheckExpression(
+          expr, dc, /*replaceInvalidRefsWithErrors=*/true)) {
     target.setExpr(expr);
     return None;
   }
@@ -2336,13 +2352,15 @@ bool TypeChecker::typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt) {
 
   // Precheck the sequence.
   Expr *sequence = stmt->getSequence();
-  if (ConstraintSystem::preCheckExpression(sequence, dc))
+  if (ConstraintSystem::preCheckExpression(
+          sequence, dc, /*replaceInvalidRefsWithErrors=*/true))
     return failed();
   stmt->setSequence(sequence);
 
   // Precheck the filtering condition.
   if (Expr *whereExpr = stmt->getWhere()) {
-    if (ConstraintSystem::preCheckExpression(whereExpr, dc))
+    if (ConstraintSystem::preCheckExpression(
+            whereExpr, dc, /*replaceInvalidRefsWithErrors=*/true))
       return failed();
 
     stmt->setWhere(whereExpr);

lib/Sema/TypeCheckPattern.cpp

@@ -530,9 +530,10 @@ class ResolvePattern : public ASTVisitor<ResolvePattern,
 
 
     // Perform unqualified name lookup to find out what the UDRE is.
-    return getSubExprPattern(TypeChecker::resolveDeclRefExpr(ude, DC));
+    return getSubExprPattern(TypeChecker::resolveDeclRefExpr(
+        ude, DC, /*replaceInvalidRefsWithErrors=*/true));
   }
-  
+
   // Call syntax forms a pattern if:
   // - the callee in 'Element(x...)' or '.Element(x...)'
   //   references an enum element. The arguments then form a tuple

lib/Sema/TypeChecker.h

@@ -315,7 +315,11 @@ Type getOptionalType(SourceLoc loc, Type elementType);
 /// Bind an UnresolvedDeclRefExpr by performing name lookup and
 /// returning the resultant expression.  Context is the DeclContext used
 /// for the lookup.
-Expr *resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE, DeclContext *Context);
+///
+/// \param replaceInvalidRefsWithErrors Indicates whether it's allowed
+/// to replace any discovered invalid member references with `ErrorExpr`.
+Expr *resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE, DeclContext *Context,
+                         bool replaceInvalidRefsWithErrors);
 
 /// Check for unsupported protocol types in the given declaration.
 void checkUnsupportedProtocolType(Decl *decl);