[CodeCompletion] Report type relations when completing inside result builders

This requires navigating the constraint system solution to retrieve the argument type of the `buildBlock` call. The comments in the code should describe what I’m doing well enough.

rdar://83846531

Author: ahoppen

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);
 

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^#
+    }
+  }
 }