[5.7-04182022][CodeCompletion] Report type relations when completing inside result builders

include/swift/Sema/ConstraintSystem.h

@@ -5296,18 +5296,21 @@ class ConstraintSystem {
                                  = FreeTypeVariableBinding::Disallow,
                                  bool allowFixes = false);
 
-  /// Construct and solve a system of constraints based on the given expression
-  /// and its contextual information.
+  /// Assuming that constraints have already been generated, solve the
+  /// constraint system for code completion, writing all solutions to
+  /// \p solutions.
   ///
   /// This method is designed to be used for code completion which means that
   /// it doesn't mutate given expression, even if there is a single valid
   /// solution, and constraint solver is allowed to produce partially correct
   /// solutions. Such solutions can have any number of holes in them.
   ///
-  /// \param target The expression involved in code completion.
-  ///
   /// \param solutions The solutions produced for the given target without
   /// filtering.
+  void solveForCodeCompletion(SmallVectorImpl<Solution> &solutions);
+
+  /// Generate constraints for \p target and solve the resulting constraint
+  /// system for code completion (see overload above).
   ///
   /// \returns `false` if this call fails (e.g. pre-check or constraint
   /// generation fails), `true` otherwise.

lib/Sema/BuilderTransform.cpp

@@ -1714,10 +1714,10 @@ Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
   }
 
   // Solve the constraint system.
-  SmallVector<Solution, 4> solutions;
-  bool solvingFailed = cs.solve(solutions);
-
   if (cs.getASTContext().CompletionCallback) {
+    SmallVector<Solution, 4> solutions;
+    cs.solveForCodeCompletion(solutions);
+
     CompletionContextFinder analyzer(func, func->getDeclContext());
     filterSolutionsForCodeCompletion(solutions, analyzer);
     for (const auto &solution : solutions) {
@@ -1726,6 +1726,9 @@ Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
     return nullptr;
   }
 
+  SmallVector<Solution, 4> solutions;
+  bool solvingFailed = cs.solve(solutions);
+
   if (solvingFailed || solutions.size() != 1) {
     // Try to fix the system or provide a decent diagnostic.
     auto salvagedResult = cs.salvage();

lib/Sema/CSSolver.cpp

@@ -1516,26 +1516,8 @@ void ConstraintSystem::solveImpl(SmallVectorImpl<Solution> &solutions) {
   }
 }
 
-bool ConstraintSystem::solveForCodeCompletion(
-    SolutionApplicationTarget &target, SmallVectorImpl<Solution> &solutions) {
-  auto *expr = target.getAsExpr();
-  // Tell the constraint system what the contextual type is.
-  setContextualType(expr, target.getExprContextualTypeLoc(),
-                    target.getExprContextualTypePurpose());
-
-  // Set up the expression type checker timer.
-  Timer.emplace(expr, *this);
-
-  shrink(expr);
-
-  if (isDebugMode()) {
-    auto &log = llvm::errs();
-    log << "--- Code Completion ---\n";
-  }
-
-  if (generateConstraints(target, FreeTypeVariableBinding::Disallow))
-    return false;
-
+void ConstraintSystem::solveForCodeCompletion(
+    SmallVectorImpl<Solution> &solutions) {
   {
     SolverState state(*this, FreeTypeVariableBinding::Disallow);
 
@@ -1556,6 +1538,30 @@ bool ConstraintSystem::solveForCodeCompletion(
     }
   }
 
+  return;
+}
+
+bool ConstraintSystem::solveForCodeCompletion(
+    SolutionApplicationTarget &target, SmallVectorImpl<Solution> &solutions) {
+  auto *expr = target.getAsExpr();
+  // Tell the constraint system what the contextual type is.
+  setContextualType(expr, target.getExprContextualTypeLoc(),
+                    target.getExprContextualTypePurpose());
+
+  // Set up the expression type checker timer.
+  Timer.emplace(expr, *this);
+
+  shrink(expr);
+
+  if (isDebugMode()) {
+    auto &log = llvm::errs();
+    log << "--- Code Completion ---\n";
+  }
+
+  if (generateConstraints(target, FreeTypeVariableBinding::Disallow))
+    return false;
+
+  solveForCodeCompletion(solutions);
   return true;
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -3558,18 +3558,79 @@ Type Solution::simplifyTypeForCodeCompletion(Type Ty) const {
 
   // Replace all type variables (which must come from placeholders) by their
   // generic parameters. Because we call into simplifyTypeImpl
-  Ty = CS.simplifyTypeImpl(Ty, [](TypeVariableType *typeVar) -> Type {
-    if (auto *GP = typeVar->getImpl().getGenericParameter()) {
-      // Code completion depends on generic parameter type being
-      // represented in terms of `ArchetypeType` since it's easy
-      // to extract protocol requirements from it.
-      if (auto *GPD = GP->getDecl()) {
-        return GPD->getInnermostDeclContext()->mapTypeIntoContext(GP);
+  Ty = CS.simplifyTypeImpl(Ty, [&CS](TypeVariableType *typeVar) -> Type {
+    // Code completion depends on generic parameter type being represented in
+    // terms of `ArchetypeType` since it's easy to extract protocol requirements
+    // from it.
+    auto getTypeVarAsArchetype = [](TypeVariableType *typeVar) -> Type {
+      if (auto *GP = typeVar->getImpl().getGenericParameter()) {
+        if (auto *GPD = GP->getDecl()) {
+          return GPD->getInnermostDeclContext()->mapTypeIntoContext(GP);
+        }
+      }
+      return Type();
+    };
+
+    if (auto archetype = getTypeVarAsArchetype(typeVar)) {
+      return archetype;
+    }
+
+    // When applying the logic below to get contextual types inside result
+    // builders, the code completion type variable is connected by a one-way
+    // constraint to a type variable in the buildBlock call, but that is not the
+    // type variable that represents the argument type. We need to find the type
+    // variable representing the argument to retrieve protocol requirements from
+    // it. Look for a ArgumentConversion constraint that allows us to retrieve
+    // the argument type var.
+    for (auto argConstraint :
+         CS.getConstraintGraph()[typeVar].getConstraints()) {
+      if (argConstraint->getKind() == ConstraintKind::ArgumentConversion &&
+          argConstraint->getFirstType()->getRValueType()->isEqual(typeVar)) {
+        if (auto argTV =
+                argConstraint->getSecondType()->getAs<TypeVariableType>()) {
+          if (auto archetype = getTypeVarAsArchetype(argTV)) {
+            return archetype;
+          }
+        }
       }
     }
+
     return typeVar;
   });
 
+  // Logic to determine the contextual type inside buildBlock result builders:
+  //
+  // When completing inside a result builder, the result builder
+  //   @ViewBuilder var body: some View {
+  //     Text("Foo")
+  //     #^COMPLETE^#
+  //   }
+  // gets rewritten to
+  //   @ViewBuilder var body: some View {
+  //     let $__builder2: Text
+  //     let $__builder0 = Text("Foo")
+  //     let $__builder1 = #^COMPLETE^#
+  //     $__builder2 = ViewBuilder.buildBlock($__builder0, $__builder1)
+  //     return $__builder2
+  //   }
+  // Inside the constraint system
+  //     let $__builder1 = #^COMPLETE^#
+  // gets type checked without context, so we can't know the contexutal type for
+  // the code completion token. But we know that $__builder1 (and thus the type
+  // of #^COMPLETE^#) is used as the second argument to ViewBuilder.buildBlock,
+  // so we can extract the contextual type from that call. To do this, figure
+  // out the type variable that is used for $__builder1 in the buildBlock call.
+  // This type variable is connected to the type variable of $__builder1's
+  // definition by a one-way constraint.
+  if (auto TV = Ty->getAs<TypeVariableType>()) {
+    for (auto constraint : CS.getConstraintGraph()[TV].getConstraints()) {
+      if (constraint->getKind() == ConstraintKind::OneWayEqual &&
+          constraint->getSecondType()->isEqual(TV)) {
+        return simplifyTypeForCodeCompletion(constraint->getFirstType());
+      }
+    }
+  }
+
   // Remove any remaining type variables and placeholders
   Ty = simplifyType(Ty);
 

lib/Sema/TypeCheckConstraints.cpp

@@ -1406,6 +1406,8 @@ void ConstraintSystem::print(raw_ostream &out) const {
       out << " [inout allowed]";
     if (tv->getImpl().canBindToNoEscape())
       out << " [noescape allowed]";
+    if (tv->getImpl().canBindToHole())
+      out << " [hole allowed]";
     auto rep = getRepresentative(tv);
     if (rep == tv) {
       if (auto fixed = getFixedType(tv)) {

test/IDE/complete_in_result_builder.swift

@@ -36,25 +36,28 @@ func testGlobalLookup() {
   @TupleBuilder<String> var x1 {
     #^GLOBAL_LOOKUP^#
     // GLOBAL_LOOKUP: Begin completions
-    // GLOBAL_LOOKUP: Decl[GlobalVar]/CurrModule:         MyConstantString[#String#];
+    // GLOBAL_LOOKUP: Decl[GlobalVar]/CurrModule/TypeRelation[Identical]:         MyConstantString[#String#];
     // GLOBAL_LOOKUP: End completions
   }
 
   @TupleBuilder<String> var x2 {
     if true {
-      #^GLOBAL_LOOKUP_IN_IF_BODY?check=GLOBAL_LOOKUP^#
+      #^GLOBAL_LOOKUP_IN_IF_BODY?check=GLOBAL_LOOKUP_NO_TYPE_RELATION^#
+// GLOBAL_LOOKUP_NO_TYPE_RELATION: Begin completions
+// GLOBAL_LOOKUP_NO_TYPE_RELATION: Decl[GlobalVar]/CurrModule:         MyConstantString[#String#];
+// GLOBAL_LOOKUP_NO_TYPE_RELATION: End completions
     }
   }
 
   @TupleBuilder<String> var x3 {
     if {
-      #^GLOBAL_LOOKUP_IN_IF_BODY_WITHOUT_CONDITION?check=GLOBAL_LOOKUP^#
+      #^GLOBAL_LOOKUP_IN_IF_BODY_WITHOUT_CONDITION?check=GLOBAL_LOOKUP_NO_TYPE_RELATION^#
     }
   }
 
   @TupleBuilder<String> var x4 {
     guard else {
-      #^GLOBAL_LOOKUP_IN_GUARD_BODY_WITHOUT_CONDITION?check=GLOBAL_LOOKUP^#
+      #^GLOBAL_LOOKUP_IN_GUARD_BODY_WITHOUT_CONDITION?check=GLOBAL_LOOKUP_NO_TYPE_RELATION^#
     }
   }
 
@@ -78,13 +81,16 @@ func testStaticMemberLookup() {
   @TupleBuilder<String> var x1 {
     StringFactory.#^COMPLETE_STATIC_MEMBER^#
     // COMPLETE_STATIC_MEMBER: Begin completions
-    // COMPLETE_STATIC_MEMBER: Decl[StaticMethod]/CurrNominal:     makeString({#x: String#})[#String#];
+    // COMPLETE_STATIC_MEMBER: Decl[StaticMethod]/CurrNominal/TypeRelation[Identical]:     makeString({#x: String#})[#String#];
     // COMPLETE_STATIC_MEMBER: End completions
   }
 
   @TupleBuilder<String> var x2 {
     if true {
-      StringFactory.#^COMPLETE_STATIC_MEMBER_IN_IF_BODY?check=COMPLETE_STATIC_MEMBER^#
+      StringFactory.#^COMPLETE_STATIC_MEMBER_IN_IF_BODY^#
+// COMPLETE_STATIC_MEMBER_IN_IF_BODY: Begin completions
+// COMPLETE_STATIC_MEMBER_IN_IF_BODY: Decl[StaticMethod]/CurrNominal:     makeString({#x: String#})[#String#];
+// COMPLETE_STATIC_MEMBER_IN_IF_BODY: End completions
     }
   }
 
@@ -208,13 +214,37 @@ func testCompleteInStringLiteral() {
 // STRING_LITERAL_VAR-DAG: Keyword[self]/CurrNominal:          self[#Island#]; name=self
 // STRING_LITERAL_VAR-DAG: Decl[InstanceVar]/CurrNominal/TypeRelation[Convertible]: turnipPrice[#String#]; name=turnipPrice
 // STRING_LITERAL_VAR: End completions
+  }
 
-
-    func bar(island: Island) {
+  func bar(island: Island) {
     BStack {
       Text("\(island.#^STRING_LITERAL_AS_ARGUMENT?check=STRING_LITERAL_VAR^#turnipPrice)")
       takeTrailingClosure {}
-    }
+    } 
+  }
+}
+
+func testTypeRelationInResultBuilder() {
+  protocol View2 {}
+
+  @resultBuilder public struct ViewBuilder2 {
+    static func buildBlock<Content>(_ content: Content) -> Content where Content : View2 { fatalError() }
+    static func buildBlock<C0, C1>(_ c0: C0, _ c1: C1) -> C0 where C0 : View2, C1: View2 { fatalError() }
   }
 
+  struct MyText: View2 {}
+
+  struct MyView {
+    @ViewBuilder2 var body: some View2 {
+      #^SINGLE_ELEMENT^#
+    }
+    // SINGLE_ELEMENT: Begin completions
+    // SINGLE_ELEMENT-DAG: Decl[Struct]/Local/TypeRelation[Convertible]: MyText[#MyText#];
+    // SINGLE_ELEMENT: End completions
+
+    @ViewBuilder2 var body2: some View2 {
+      MyText()
+      #^SECOND_ELEMENT?check=SINGLE_ELEMENT^#
+    }
+  }
 }