[Sema] Implement type erasure for dynamic replacement.

include/swift/AST/Attr.def

@@ -520,7 +520,7 @@ SIMPLE_DECL_ATTR(_inheritsConvenienceInitializers,
   93)
 
 DECL_ATTR(_typeEraser, TypeEraser,
-  OnProtocol | UserInaccessible | NotSerialized |
+  OnProtocol | UserInaccessible |
   ABIStableToAdd | ABIBreakingToRemove | APIStableToAdd | APIBreakingToRemove,
   94)
 

lib/AST/Attr.cpp

@@ -985,6 +985,20 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     break;
   }
 
+  case DAK_TypeEraser: {
+    Printer.printAttrName("@_typeEraser");
+    Printer << "(";
+    Printer.callPrintNamePre(PrintNameContext::Attribute);
+    auto typeLoc = cast<TypeEraserAttr>(this)->getTypeEraserLoc();
+    if (auto type = typeLoc.getType())
+      type->print(Printer, Options);
+    else
+      typeLoc.getTypeRepr()->print(Printer, Options);
+    Printer.printNamePost(PrintNameContext::Attribute);
+    Printer << ")";
+    break;
+  }
+
   case DAK_Custom: {
     Printer.callPrintNamePre(PrintNameContext::Attribute);
     Printer << "@";

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

@@ -3919,6 +3919,10 @@ bool ConstraintSystem::generateConstraints(
       target.setExprConversionType(TypeChecker::getOptionalType(expr->getLoc(), var));
     }
 
+    expr = buildTypeErasedExpr(expr, target.getDeclContext(),
+                               target.getExprContextualType(),
+                               target.getExprContextualTypePurpose());
+
     // Generate constraints for the main system.
     expr = generateConstraints(expr, target.getDeclContext());
     if (!expr)

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.

lib/Serialization/Deserialization.cpp

@@ -4208,6 +4208,19 @@ llvm::Error DeclDeserializer::deserializeDeclAttributes() {
         break;
       }
 
+      case decls_block::TypeEraser_DECL_ATTR: {
+        bool isImplicit;
+        TypeID typeEraserID;
+        serialization::decls_block::TypeEraserDeclAttrLayout::readRecord(
+            scratch, isImplicit, typeEraserID);
+
+        auto typeEraser = MF.getType(typeEraserID);
+        assert(!isImplicit);
+        Attr = new (ctx) TypeEraserAttr(SourceLoc(), SourceRange(),
+                                        TypeLoc::withoutLoc(typeEraser));
+        break;
+      }
+
       case decls_block::Custom_DECL_ATTR: {
         bool isImplicit;
         TypeID typeID;

lib/Serialization/Serialization.cpp

@@ -2173,7 +2173,6 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
     case DAK_RestatedObjCConformance:
     case DAK_ClangImporterSynthesizedType:
     case DAK_PrivateImport:
-    case DAK_TypeEraser:
       llvm_unreachable("cannot serialize attribute");
 
     case DAK_Count:
@@ -2364,6 +2363,16 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
       return;
     }
 
+    case DAK_TypeEraser: {
+      auto abbrCode = S.DeclTypeAbbrCodes[TypeEraserDeclAttrLayout::Code];
+      auto attr = cast<TypeEraserAttr>(DA);
+      auto typeEraser = attr->getTypeEraserLoc().getType();
+      TypeEraserDeclAttrLayout::emitRecord(S.Out, S.ScratchRecord, abbrCode,
+                                           attr->isImplicit(),
+                                           S.addTypeRef(typeEraser));
+      return;
+    }
+
     case DAK_Custom: {
       auto abbrCode = S.DeclTypeAbbrCodes[CustomDeclAttrLayout::Code];
       auto theAttr = cast<CustomAttr>(DA);

test/Sema/type_eraser.swift

@@ -0,0 +1,109 @@
+// RUN: %target-swift-frontend -typecheck -disable-availability-checking -dump-ast %s | %FileCheck %s
+
+class AnyP: P {
+  init<T: P>(erasing: T) {}
+}
+
+@_typeEraser(AnyP)
+protocol P {}
+
+struct ConcreteP: P, Hashable {}
+
+// CHECK-LABEL: testBasic
+dynamic func testBasic() -> some P {
+  // CHECK: underlying_to_opaque_expr{{.*}}'some P'
+  // CHECK-NEXT: call_expr implicit type='AnyP'{{.*}}arg_labels=erasing:
+  // CHECK: call_expr type='ConcreteP'
+  ConcreteP()
+}
+
+// CHECK-LABEL: testTypeAlias
+typealias AliasForP = P
+dynamic func testTypeAlias() -> some AliasForP {
+  // CHECK: underlying_to_opaque_expr{{.*}}'some P'
+  // CHECK-NEXT: call_expr implicit type='AnyP'{{.*}}arg_labels=erasing:
+  // CHECK: call_expr type='ConcreteP'
+  ConcreteP()
+}
+
+// CHECK-LABEL: testNoDynamic
+func testNoDynamic() -> some P {
+  // CHECK: underlying_to_opaque_expr{{.*}}'some P'
+  // CHECK-NEXT: call_expr type='ConcreteP'
+  ConcreteP()
+}
+
+// CHECK-LABEL: testNoOpaque
+dynamic func testNoOpaque() -> P {
+  // CHECK: erasure_expr implicit type='P'
+  // CHECK-NEXT: normal_conformance type=ConcreteP protocol=P
+  // CHECK-NEXT: call_expr type='ConcreteP'
+  ConcreteP()
+}
+
+// CHECK-LABEL: testComposition
+typealias Composition = P & Hashable
+dynamic func testComposition() -> some Composition {
+  // CHECK: underlying_to_opaque_expr{{.*}}'some Hashable & P'
+  // CHECK-NEXT: call_expr type='ConcreteP'
+  ConcreteP()
+}
+
+// CHECK-LABEL: struct_decl{{.*}}Builder
+@_functionBuilder
+struct Builder {
+  static func buildBlock(_ params: P...) -> ConcreteP {
+    return ConcreteP()
+  }
+}
+
+// CHECK-LABEL: TestFunctionBuilder
+class TestFunctionBuilder {
+  // CHECK-LABEL: testTransformFnBody
+  @Builder dynamic var testTransformFnBody: some P {
+    // CHECK: return_stmt
+    // CHECK-NEXT: underlying_to_opaque_expr implicit type='some P'
+    // CHECK-NEXT: call_expr implicit type='AnyP'{{.*}}arg_labels=erasing:
+    // CHECK: declref_expr implicit type='@lvalue ConcreteP'
+    ConcreteP()
+  }
+
+  // CHECK-LABEL: func_decl{{.*}}takesBuilder
+  func takesBuilder(@Builder closure: () -> ConcreteP) -> ConcreteP { closure() }
+
+  // CHECK-LABEL: testClosureBuilder
+  dynamic var testClosureBuilder: some P {
+    // CHECK: underlying_to_opaque_expr implicit type='some P'
+    // CHECK-NEXT: call_expr implicit type='AnyP'{{.*}}arg_labels=erasing:
+    // CHECK: closure_expr type='() -> ConcreteP'
+    takesBuilder {
+      // CHECK: return_stmt
+      // CHECK-NEXT: load_expr implicit type='ConcreteP'
+      ConcreteP()
+    }
+  }
+}
+
+// CHECK-LABEL: class_decl{{.*}}DynamicReplacement
+class DynamicReplacement {
+  dynamic func testDynamicReplaceable() -> some P {
+    // CHECK: underlying_to_opaque_expr implicit type='some P'
+    // CHECK-NEXT: call_expr implicit type='AnyP'{{.*}}arg_labels=erasing:
+    // CHECK: call_expr type='ConcreteP'
+    ConcreteP()
+  }
+}
+
+// CHECK-LABEL: extension_decl{{.*}}DynamicReplacement
+extension DynamicReplacement {
+  // CHECK-LABEL: testDynamicReplacement
+  @_dynamicReplacement(for: testDynamicReplaceable)
+  func testDynamicReplacement() -> some P {
+    print("not single expr return")
+    // CHECK: return_stmt
+    // CHECK-NEXT: underlying_to_opaque_expr implicit type='some P'
+    // CHECK-NEXT: call_expr implicit type='AnyP'{{.*}}arg_labels=erasing:
+    // CHECK: call_expr type='ConcreteP'
+    return ConcreteP()
+  }
+}

test/Serialization/serialize_attr.swift

@@ -82,3 +82,15 @@ public class CC<T : PP> {
 // CHECK-DAG: sil [serialized] [_specialize exported: false, kind: full, where T == Int, U == Float] [canonical] [ossa] @$s14serialize_attr14specializeThis_1uyx_q_tr0_lF : $@convention(thin) <T, U> (@in_guaranteed T, @in_guaranteed U) -> () {
 
 // CHECK-DAG: sil [serialized] [noinline] [_specialize exported: false, kind: full, where T == RR, U == SS] [canonical] [ossa] @$s14serialize_attr2CCC3foo_1gqd___AA2GGVyxGtqd___AHtAA2QQRd__lF : $@convention(method) <T where T : PP><U where U : QQ> (@in_guaranteed U, GG<T>, @guaranteed CC<T>) -> (@out U, GG<T>) {
+
+
+ // @_typeEraser
+ // -----------------------------------------------------------------------------
+
+public class AnyP: P {
+  public init<T: P>(erasing: T) {}
+}
+
+// CHECK-DAG: @_typeEraser(AnyP) protocol P
+@_typeEraser(AnyP)
+public protocol P {}