[Sema] Support type erasure for dynamic replacement in function

builders.

Author: hborla

lib/Sema/BuilderTransform.cpp

@@ -210,6 +210,10 @@ class BuilderClosureVisitor
       return None;
 
     applied.returnExpr = buildVarRef(bodyVar, stmt->getEndLoc());
+    applied.returnExpr = cs->buildTypeErasedExpr(applied.returnExpr,
+                                                 dc, applied.bodyResultType,
+                                                 CTP_ReturnStmt);
+
     applied.returnExpr = cs->generateConstraints(applied.returnExpr, dc);
     if (!applied.returnExpr) {
       hadError = true;

lib/Sema/CSApply.cpp

@@ -7417,6 +7417,7 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
       return None;
 
     result.setFunctionBody(newBody);
+    fn.getAbstractFunctionDecl()->setHasSingleExpressionBody(false);
   }
 
   // Follow-up tasks.

lib/Sema/CSGen.cpp

@@ -3903,38 +3903,6 @@ static bool generateInitPatternConstraints(
   return false;
 }
 
-static Expr *eraseTypeForDynamicReplacement(SolutionApplicationTarget target,
-                                            Expr *expr) {
-  auto purpose = target.getExprContextualTypePurpose();
-  auto *dc = target.getDeclContext();
-  auto *decl = dyn_cast_or_null<ValueDecl>(dc->getAsDecl());
-  if (!decl)
-    return expr;
-
-  if (!(purpose == CTP_ReturnStmt || purpose == CTP_ReturnSingleExpr) ||
-      !(decl->isDynamic() || decl->getDynamicallyReplacedDecl()))
-    return expr;
-
-  auto *opaque =
-      target.getExprContextualType()->getAs<OpaqueTypeArchetypeType>();
-  if (!opaque)
-    return expr;
-
-  auto protocols = opaque->getConformsTo();
-  if (protocols.size() != 1)
-    return expr;
-
-  auto *attr = protocols.front()->getAttrs().getAttribute<TypeEraserAttr>();
-  if (!attr)
-    return expr;
-
-  auto typeEraser = attr->getTypeEraserLoc().getType();
-  auto &ctx = dc->getASTContext();
-  return CallExpr::createImplicit(ctx,
-                                  TypeExpr::createImplicit(typeEraser, ctx),
-                                  {expr}, {ctx.Id_erasing});
-}
-
 bool ConstraintSystem::generateConstraints(
     SolutionApplicationTarget &target,
     FreeTypeVariableBinding allowFreeTypeVariables) {
@@ -3951,7 +3919,9 @@ bool ConstraintSystem::generateConstraints(
       target.setExprConversionType(TypeChecker::getOptionalType(expr->getLoc(), var));
     }
 
-    expr = eraseTypeForDynamicReplacement(target, expr);
+    expr = buildTypeErasedExpr(expr, target.getDeclContext(),
+                               target.getExprContextualType(),
+                               target.getExprContextualTypePurpose());
 
     // Generate constraints for the main system.
     expr = generateConstraints(expr, target.getDeclContext());

lib/Sema/ConstraintSystem.cpp

@@ -4062,6 +4062,36 @@ Expr *ConstraintSystem::buildAutoClosureExpr(Expr *expr,
   return result;
 }
 
+Expr *ConstraintSystem::buildTypeErasedExpr(Expr *expr, DeclContext *dc,
+                                            Type contextualType,
+                                            ContextualTypePurpose purpose) {
+  if (!(purpose == CTP_ReturnStmt || purpose == CTP_ReturnSingleExpr))
+    return expr;
+
+  auto *decl = dyn_cast_or_null<ValueDecl>(dc->getAsDecl());
+  if (!decl ||
+      !(decl->isDynamic() || decl->getDynamicallyReplacedDecl()))
+    return expr;
+
+  auto *opaque = contextualType->getAs<OpaqueTypeArchetypeType>();
+  if (!opaque)
+    return expr;
+
+  auto protocols = opaque->getConformsTo();
+  if (protocols.size() != 1)
+    return expr;
+
+  auto *attr = protocols.front()->getAttrs().getAttribute<TypeEraserAttr>();
+  if (!attr)
+    return expr;
+
+  auto typeEraser = attr->getTypeEraserLoc().getType();
+  auto &ctx = dc->getASTContext();
+  return CallExpr::createImplicit(ctx,
+                                  TypeExpr::createImplicit(typeEraser, ctx),
+                                  {expr}, {ctx.Id_erasing});
+}
+
 /// If an UnresolvedDotExpr, SubscriptMember, etc has been resolved by the
 /// constraint system, return the decl that it references.
 ValueDecl *ConstraintSystem::findResolvedMemberRef(ConstraintLocator *locator) {

lib/Sema/ConstraintSystem.h

@@ -3564,6 +3564,21 @@ class ConstraintSystem {
   /// Given expression represents computed result of the closure.
   Expr *buildAutoClosureExpr(Expr *expr, FunctionType *closureType);
 
+  /// Builds a type-erased return expression that can be used in dynamic
+  /// replacement.
+  ///
+  /// An expression needs type erasure if:
+  ///  1. The expression is a return value.
+  ///  2. The enclosing function is dynamic or a dynamic replacement.
+  ///  3. The enclosing function returns an opaque type.
+  ///  4. The opaque type conforms to (exactly) one protocol, and the protocol
+  ///     has a declared type eraser.
+  ///
+  /// \returns the transformed return expression, or the original expression if
+  /// no type erasure is needed.
+  Expr *buildTypeErasedExpr(Expr *expr, DeclContext *dc, Type contextualType,
+                            ContextualTypePurpose purpose);
+
 private:
   /// Determines whether or not a given conversion at a given locator requires
   /// the creation of a temporary value that's only valid for a limited scope.

test/Sema/type_eraser.swift

@@ -16,3 +16,27 @@ dynamic func opaque() -> some P {
   ConcreteP()
 }
 
+class AnyQ: Q {
+  init<T: Q>(erasing: T) {}
+}
+@_typeEraser(AnyQ)
+protocol Q {}
+
+struct ConcreteQ: Q {}
+
+@_functionBuilder
+struct Builder {
+  static func buildBlock(_ params: Q...) -> ConcreteQ {
+    return ConcreteQ()
+  }
+}
+
+// CHECK-LABEL: transformFnBody
+@Builder
+dynamic var transformFnBody: some Q {
+  // CHECK: return_stmt
+  // CHECK-NEXT: underlying_to_opaque_expr implicit type='some Q'
+  // CHECK-NEXT: call_expr implicit type='AnyQ'{{.*}}arg_labels=erasing:
+  // CHECK: declref_expr implicit type='@lvalue ConcreteQ'
+  ConcreteQ()
+}