Support nested examples like @Array.Builder

Author: calda

include/swift/AST/DiagnosticsSema.def

@@ -7798,7 +7798,7 @@ NOTE(result_builder_infer_pick_specific, none,
      "apply result builder %0 (inferred from %select{protocol|dynamic replacement of}1 %2)",
      (Type, unsigned, DeclName))
 ERROR(result_builder_generic_inference_failed, none,
-      "unable to infer generic parameters for result builder %0", (DeclName))
+      "unable to infer generic parameters for result builder %0", (DeclAttribute))
 GROUPED_WARNING(result_builder_missing_limited_availability,
         ResultBuilderMethods, none,
         "result builder %0 does not implement 'buildLimitedAvailability'; "

lib/Sema/TypeCheckType.cpp

@@ -6984,47 +6984,33 @@ struct ResultBuilderUnboundTypeOpener {
   CustomAttr *attr;
 
   Type operator()(UnboundGenericType *unboundTy) const {
-    auto resultBuilderDecl =
-        dyn_cast_or_null<NominalTypeDecl>(unboundTy->getDecl());
-    if (!resultBuilderDecl ||
-        !resultBuilderDecl->getAttrs().hasAttribute<ResultBuilderAttr>()) {
-      return invalidResultBuilderType(unboundTy);
+    auto resultBuilderDecl = attr->getNominalDecl();
+    if (!resultBuilderDecl) {
+      return invalidResultBuilderType();
     }
 
-    Decl *owningDecl = attr->getOwner().getAsDecl();
-    if (!owningDecl) {
-      return invalidResultBuilderType(unboundTy);
-    }
-
-    // Retrieve the return type of the owning declaration that
-    // provides type inference for the result builder type.
-    Type owningDeclResultType;
-    if (auto varDecl = dyn_cast<VarDecl>(owningDecl)) {
-      if (auto closureResultType =
-              varDecl->getInterfaceType()->getAs<FunctionType>()) {
-        owningDeclResultType = closureResultType->getResult();
-      } else {
-        owningDeclResultType = varDecl->getInterfaceType();
-      }
-    } else if (auto funcDecl = dyn_cast<FuncDecl>(owningDecl)) {
-      owningDeclResultType = funcDecl->getResultInterfaceType();
+    auto unboundTyDecl =
+        dyn_cast_or_null<NominalTypeDecl>(unboundTy->getDecl());
+    if (!resultBuilderDecl) {
+      return invalidResultBuilderType();
     }
 
-    if (!owningDeclResultType) {
-      return invalidResultBuilderType(unboundTy);
+    auto owningDeclReturnType = getOwningDeclReturnType();
+    if (!owningDeclReturnType) {
+      return invalidResultBuilderType();
     }
 
     // Retrieve the supported result types of the result builder.
-    auto resultTypes = retrieveResultBuilderResultTypes(resultBuilderDecl);
+    auto resultTypes = getResultBuilderResultTypes(resultBuilderDecl);
     if (resultTypes.empty()) {
-      return invalidResultBuilderType(unboundTy);
+      return invalidResultBuilderType();
     }
 
-    auto genericSig = resultBuilderDecl->getGenericSignature();
+    auto genericSig = unboundTyDecl->getGenericSignature();
 
     // Try each result type and return the first one that
     // produces a valid solution.
-    for (auto componentType : resultTypes) {
+    for (auto resultType : resultTypes) {
       using namespace constraints;
 
       // Avoid recording failed constraints, which are not used here.
@@ -7038,7 +7024,7 @@ struct ResultBuilderUnboundTypeOpener {
       llvm::SmallVector<TypeVariableType *, 8> typeVars;
       for (unsigned i = 0; i < genericSig.getGenericParams().size(); ++i) {
         auto locator = cs.getConstraintLocator(
-            resultBuilderDecl, {ConstraintLocator::GenericArgument, i});
+            unboundTyDecl, {ConstraintLocator::GenericArgument, i});
         auto typeVar = cs.createTypeVariable(locator, TVO_CanBindToHole);
         typeVarReplacements.push_back(typeVar);
         typeVars.push_back(typeVar);
@@ -7049,12 +7035,12 @@ struct ResultBuilderUnboundTypeOpener {
       auto subMap = SubstitutionMap::get(genericSig, typeVarReplacements,
                                          LookUpConformanceInModule());
 
-      auto componentTypeWithTypeVars = componentType.subst(subMap);
+      auto resultTypeWithTypeVars = resultType.subst(subMap);
 
       // The result builder result type should be equal to the return type of
       // the attached declaration.
-      cs.addConstraint(ConstraintKind::Equal, owningDeclResultType,
-                       componentTypeWithTypeVars,
+      cs.addConstraint(ConstraintKind::Equal, owningDeclReturnType,
+                       resultTypeWithTypeVars,
                        /*preparedOverload:*/ nullptr);
 
       auto solution = cs.solveSingle();
@@ -7067,30 +7053,52 @@ struct ResultBuilderUnboundTypeOpener {
           solvedReplacements.push_back(solution->typeBindings[typeVar]);
         }
 
-        return BoundGenericType::get(resultBuilderDecl, unboundTy->getParent(),
+        return BoundGenericType::get(unboundTyDecl, unboundTy->getParent(),
                                      solvedReplacements);
       }
     }
 
     // No result type produced a valid solution
-    return invalidResultBuilderType(unboundTy);
+    return invalidResultBuilderType();
   }
 
 private:
-  Type invalidResultBuilderType(UnboundGenericType *unboundTy) const {
+  Type invalidResultBuilderType() const {
     dc->getASTContext().Diags.diagnose(
         attr->getTypeExpr()->getLoc(),
-        diag::result_builder_generic_inference_failed,
-        unboundTy->getDecl()->getName());
+        diag::result_builder_generic_inference_failed, attr);
 
     return ErrorType::get(dc->getASTContext());
   }
 
+  /// Retrieve the return type of the owning declaration that
+  /// provides type inference for the result builder type.
+  Type getOwningDeclReturnType() const {
+    Decl *owningDecl = attr->getOwner().getAsDecl();
+    if (!owningDecl) {
+      return nullptr;
+    }
+
+    Type owningDeclReturnType;
+    if (auto varDecl = dyn_cast<VarDecl>(owningDecl)) {
+      if (auto closureResultType =
+              varDecl->getInterfaceType()->getAs<FunctionType>()) {
+        owningDeclReturnType = closureResultType->getResult();
+      } else {
+        owningDeclReturnType = varDecl->getInterfaceType();
+      }
+    } else if (auto funcDecl = dyn_cast<FuncDecl>(owningDecl)) {
+      owningDeclReturnType = funcDecl->getResultInterfaceType();
+    }
+
+    return owningDeclReturnType;
+  }
+
   /// Retrieves the valid result types of the result builder
   /// (the return types of `buildFinalResult` / `buildBlock` /
   /// `buildPartialBlock`).
   llvm::SmallVector<Type, 4>
-  retrieveResultBuilderResultTypes(NominalTypeDecl *builder) const {
+  getResultBuilderResultTypes(NominalTypeDecl *builder) const {
     ASTContext &ctx = builder->getASTContext();
     llvm::SmallVector<Type, 4> resultTypes;
 

test/Constraints/result_builder_diags.swift

@@ -1119,7 +1119,7 @@ func testInferResultBuilderGenerics() {
     print(makeArray())
   }
 
-  @SimpleArrayBuilder // expected-error{{unable to infer generic parameters for result builder 'SimpleArrayBuilder'}}
+  @SimpleArrayBuilder // expected-error{{unable to infer generic parameters for result builder @SimpleArrayBuilder}}
   var string: String {
     "foo"
   }
@@ -1167,7 +1167,7 @@ func testInferResultBuilderGenerics() {
     }
   }
 
-  @TooManyArgsBuilder // expected-error {{unable to infer generic parameters for result builder 'TooManyArgsBuilder'}}
+  @TooManyArgsBuilder // expected-error {{unable to infer generic parameters for result builder @TooManyArgsBuilder}}
   var invalidTupleArray: [(String, Int)] {
     ("foo", 1) // expected-warning {{expression of type '(String, Int)' is unused}}
     ("bar", 2) // expected-warning {{expression of type '(String, Int)' is unused}}
@@ -1235,7 +1235,7 @@ func testInferResultBuilderGenerics() {
       "d"
   }
 
-  @CollectionBuilder // expected-error {{unable to infer generic parameters for result builder 'CollectionBuilder'}}
+  @CollectionBuilder // expected-error {{unable to infer generic parameters for result builder @CollectionBuilder}}
   var contiguousArray: ContiguousArray<String> {
       "c" // expected-warning {{string literal is unused}}
       "d" // expected-warning {{string literal is unused}}
@@ -1255,7 +1255,7 @@ extension Array {
       self = build3()
   }
 
-  init(@SimpleArrayBuilder build3: () -> ContiguousArray<Element>) { // expected-error {{unable to infer generic parameters for result builder 'SimpleArrayBuilder'}}
+  init(@SimpleArrayBuilder build3: () -> ContiguousArray<Element>) { // expected-error {{unable to infer generic parameters for result builder @SimpleArrayBuilder}}
       self = Array(build3())
   }
 }
@@ -1269,7 +1269,51 @@ extension Set {
     self = Set(build2())
   }
 
-  init(@SimpleArrayBuilder build3: () -> Self) { // expected-error {{unable to infer generic parameters for result builder 'SimpleArrayBuilder'}}
+  init(@SimpleArrayBuilder build3: () -> Self) { // expected-error {{unable to infer generic parameters for result builder @SimpleArrayBuilder}}
     self = build3()
   }
 }
+
+extension Array {
+  @resultBuilder
+  struct Builder {
+    static func buildBlock(_ elements: Element...) -> [Element] {
+      elements
+    }
+  }
+}
+
+extension Dictionary {
+  @resultBuilder
+  struct Builder {
+    static func buildBlock(_ elements: (Key, Value)...) -> Dictionary {
+      Dictionary(elements, uniquingKeysWith: { $1 })
+    }
+  }
+}
+
+func testNonGenericResultBuildersInGenericTypeExtensions() {
+  @Array.Builder
+  var elements: [String] {
+    "foo"
+    "bar"
+  }
+  
+  @Array.Builder // expected-error {{unable to infer generic parameters for result builder @Array.Builder}}
+  var set: Set<String> {
+    "foo" // expected-warning {{string literal is unused}}
+    "bar" // expected-warning {{string literal is unused}}
+  }
+  
+  @Dictionary.Builder
+  var dictionary: [String: String] {
+    ("foo", "bar")
+    ("baaz", "quux")
+  }
+  
+  @Dictionary.Builder // expected-error {{unable to infer generic parameters for result builder @Dictionary.Builder}}
+  var dictionary2: [(String, String)] {
+    ("foo", "bar") // expected-warning {{expression of type '(String, String)' is unused}}
+    ("baaz", "quux") // expected-warning {{expression of type '(String, String)' is unused}}
+  }
+}

test/decl/var/result_builders.swift

@@ -72,13 +72,13 @@ struct GenericContainer<T> {
   struct Maker {} // expected-error {{result builder must provide at least one static 'buildBlock' method}}
 }
 
-func makeParamUnbound(@GenericMaker // expected-error {{unable to infer generic parameters for result builder 'GenericMaker'}}
+func makeParamUnbound(@GenericMaker // expected-error {{unable to infer generic parameters for result builder @GenericMaker}}
                       fn: () -> ()) {}
 
 func makeParamBound(@GenericMaker<Int>
                     fn: () -> ()) {}
 
-func makeParamNestedUnbound(@GenericContainer.Maker // expected-error {{unable to infer generic parameters for result builder 'GenericContainer'}}
+func makeParamNestedUnbound(@GenericContainer.Maker // expected-error {{unable to infer generic parameters for result builder @GenericContainer.Maker}}
                             fn: () -> ()) {}
 
 func makeParamNestedBound(@GenericContainer<Int>.Maker