Clean up

Author: calda

lib/Sema/TypeCheckType.cpp

@@ -6976,36 +6976,52 @@ void TypeChecker::checkExistentialTypes(
   checker.checkRequirements(genericParams->getRequirements());
 }
 
-// copied from elsewhere, to unify later
-Type inferResultBuilderComponentType(NominalTypeDecl *builder) {
-  Type componentType;
-
-  SmallVector<ValueDecl *, 4> potentialMatches;
+/// Retrieves the valid result types of the result builder
+/// (the return types of `buildFinalResult` / `buildBlock` / `buildPartialResult`).
+llvm::SmallVector<Type, 4> retrieveResultBuilderResultTypes(NominalTypeDecl *builder) {
   ASTContext &ctx = builder->getASTContext();
-  bool supportsBuildBlock = TypeChecker::typeSupportsBuilderOp(
-      builder->getDeclaredInterfaceType(), builder, ctx.Id_buildBlock,
-      /*argLabels=*/{}, &potentialMatches);
-  if (supportsBuildBlock) {
+  llvm::SmallVector<Type, 4> resultTypes;
+
+  Identifier methodIds[] = {
+    ctx.Id_buildFinalResult,
+    ctx.Id_buildBlock,
+    ctx.Id_buildPartialBlock
+  };
+
+  for (auto methodId : methodIds) {
+    SmallVector<ValueDecl *, 4> potentialMatches;
+    bool supportsMethod = TypeChecker::typeSupportsBuilderOp(
+        builder->getDeclaredInterfaceType(), builder, methodId,
+        /*argLabels=*/{}, &potentialMatches);
+
+    if (!supportsMethod)
+      continue;
+
     for (auto decl : potentialMatches) {
       auto func = dyn_cast<FuncDecl>(decl);
       if (!func || !func->isStatic())
         continue;
 
-      // If we haven't seen a component type before, gather it.
-      if (!componentType) {
-        componentType = func->getResultInterfaceType();
+      auto resultType = func->getResultInterfaceType();
+      if (!resultType || resultType->hasError())
         continue;
+
+      // Add the result type if we haven't seen it before
+      bool isDuplicate = false;
+      for (auto existingType : resultTypes) {
+        if (existingType->isEqual(resultType)) {
+          isDuplicate = true;
+          break;
+        }
       }
 
-      // If there are inconsistent component types, bail out.
-      if (!componentType->isEqual(func->getResultInterfaceType())) {
-        componentType = Type();
-        break;
+      if (!isDuplicate) {
+        resultTypes.push_back(resultType);
       }
     }
   }
-
-  return componentType;
+  
+  return resultTypes;
 }
 
 Type invalidResultBuilderType(UnboundGenericType* unboundTy,
@@ -7018,10 +7034,13 @@ Type invalidResultBuilderType(UnboundGenericType* unboundTy,
   return ErrorType::get(dc->getASTContext());
 }
 
+/// Opens a result builder `UnboundGenericType` into a `BoundGenericType`
+/// by creating and solving a constraint between the result builder's
+/// result type and the return type of the attached declaration.
 Type openUnboundResultBuilderType(UnboundGenericType* unboundTy,
                                   CustomAttr *attr, DeclContext *dc) {
-  auto builder = dyn_cast_or_null<NominalTypeDecl>(unboundTy->getDecl());
-  if (!builder) {
+  auto resultBuilderDecl = dyn_cast_or_null<NominalTypeDecl>(unboundTy->getDecl());
+  if (!resultBuilderDecl) {
     return invalidResultBuilderType(unboundTy, attr, dc);
   }
 
@@ -7043,57 +7062,61 @@ Type openUnboundResultBuilderType(UnboundGenericType* unboundTy,
   if (!owningDeclResultType) {
     return invalidResultBuilderType(unboundTy, attr, dc);
   }
-  
-  // Retrieve the result type of the result builder itself
-  auto componentType = inferResultBuilderComponentType(builder);
-  if (!componentType) {
+
+  // Retrieve the supported result types of the result builder.
+  auto resultTypes = retrieveResultBuilderResultTypes(resultBuilderDecl);
+  if (resultTypes.empty()) {
     return invalidResultBuilderType(unboundTy, attr, dc);
   }
-  
-  using namespace constraints;
-  ConstraintSystem cs(dc, std::nullopt);
 
-  // Create a type variable for each of the result builder's generic params
-  auto genericSig = builder->getGenericSignature();
+  auto genericSig = resultBuilderDecl->getGenericSignature();
 
-  llvm::SmallVector<Type, 8> typeVarReplacements;
-  llvm::SmallVector<TypeVariableType*, 8> typeVars;
-  for (unsigned i = 0; i < genericSig.getGenericParams().size(); ++i) {
-    auto locator = cs.getConstraintLocator(builder);
-    auto typeVar = cs.createTypeVariable(locator, TVO_CanBindToHole);
-    typeVarReplacements.push_back(typeVar);
-    typeVars.push_back(typeVar);
-  }
+  for (auto componentType : resultTypes) {
+    using namespace constraints;
+    ConstraintSystem cs(dc, std::nullopt);
 
-  // Retrieve the generic types of the ResultBuilder component and return type
-  // of the attached function, replacing any reference to the ResultBuilder's
-  // generic arguments with the corresponding type vars.
-  auto subMap = SubstitutionMap::get(
-    genericSig,
-    typeVarReplacements,
-    LookUpConformanceInModule());
+    // Create a type variable for each of the result builder's generic params
+    llvm::SmallVector<Type, 8> typeVarReplacements;
+    llvm::SmallVector<TypeVariableType*, 8> typeVars;
+    for (unsigned i = 0; i < genericSig.getGenericParams().size(); ++i) {
+      auto locator = cs.getConstraintLocator(
+        resultBuilderDecl,
+        {ConstraintLocator::GenericArgument, i});
+      auto typeVar = cs.createTypeVariable(locator, TVO_CanBindToHole);
+      typeVarReplacements.push_back(typeVar);
+      typeVars.push_back(typeVar);
+    }
 
-  auto componentTypeWithTypeVars = componentType.subst(subMap);
+    // Replace any references to the result builder's generic params
+    // in the result type with the corresponding type variables.
+    auto subMap = SubstitutionMap::get(
+      genericSig,
+      typeVarReplacements,
+      LookUpConformanceInModule());
 
-  // The result builder result type should be equal to the return type of the attached declaration
-  cs.addConstraint(ConstraintKind::Equal,
-                   owningDeclResultType,
-                   componentTypeWithTypeVars,
-                   /*preparedOverload:*/ nullptr);
+    auto componentTypeWithTypeVars = componentType.subst(subMap);
 
-  auto solution = cs.solveSingle();
+    // The result builder result type should be equal to the return type of the attached declaration.
+    cs.addConstraint(ConstraintKind::Equal,
+                     owningDeclResultType,
+                     componentTypeWithTypeVars,
+                     /*preparedOverload:*/ nullptr);
 
-  if (!solution) {
-    return invalidResultBuilderType(unboundTy, attr, dc);
-  }
+    auto solution = cs.solveSingle();
 
-  // Bind the result builder type to the solved type parameters
-  llvm::SmallVector<Type, 8> solvedReplacements;
-  for (auto typeVar : typeVars) {
-    solvedReplacements.push_back(solution->typeBindings[typeVar]);
+    // If a solution exists, bind the result builder's generic params to the solved types.
+    if (solution) {
+      llvm::SmallVector<Type, 8> solvedReplacements;
+      for (auto typeVar : typeVars) {
+        solvedReplacements.push_back(solution->typeBindings[typeVar]);
+      }
+
+      return BoundGenericType::get(resultBuilderDecl, unboundTy->getParent(), solvedReplacements);
+    }
   }
 
-  return BoundGenericType::get(builder, unboundTy->getParent(), solvedReplacements);
+  // No result type produced a valid solution
+  return invalidResultBuilderType(unboundTy, attr, dc);
 }
 
 Type CustomAttrTypeRequest::evaluate(Evaluator &eval, CustomAttr *attr,

test/Constraints/result_builder_diags.swift

@@ -1157,5 +1157,54 @@ func testInferResultBuilderGenerics() {
     ("foo", 1) // expected-warning {{expression of type '(String, Int)' is unused}}
     ("bar", 2) // expected-warning {{expression of type '(String, Int)' is unused}}
   }
+  
+  @resultBuilder
+  enum ComplexListBuilder<Element> {
+    static func buildBlock(_ elements: Element...) -> [Element] {
+      elements
+    }
+    
+    static func buildFinalResult(_ component: [Element]) -> Set<Element> where Element: Hashable {
+      Set(component)
+    }
+    
+    static func buildFinalResult(_ component: [Element]) -> ContiguousArray<Element> {
+      ContiguousArray(component)
+    }
+  }
+
+  @ComplexListBuilder
+  var stringSetFromBuildFinalResult: Set<String> {
+    "foo"
+    "bar"
+  }
+
+  @ComplexListBuilder
+  var contiguousStringsFromBuildFinalResult: ContiguousArray<String> {
+    "foo"
+    "bar"
+  }
+
+  @ComplexListBuilder
+  var stringArrayFromBuildFinalResult: [String] { // expected-error {{cannot convert return expression of type 'Set<String>' to return type '[String]'}}
+    "foo"
+    "bar"
+  }
+
+  @resultBuilder
+  enum PartialArrayBuilder<Element> {
+    static func buildPartialBlock(first: Element) -> [Element] {
+      [first]
+    }
+
+    static func buildPartialBlock(accumulated: [Element], next: Element) -> [Element] {
+      accumulated + [next]
+    }
+  }
 
+  @PartialArrayBuilder
+  var stringArrayFromBuildPartial: [String] {
+    "foo"
+    "bar"
+  }
 }