Merge pull request #77550 from tshortli/opaque-return-type-printing

Author: tshortli

include/swift/AST/PrintOptions.h

@@ -366,6 +366,10 @@ struct PrintOptions {
   OpaqueReturnTypePrintingMode OpaqueReturnTypePrinting =
       OpaqueReturnTypePrintingMode::WithOpaqueKeyword;
 
+  /// If non-null, opaque types that have this naming decl should be printed as
+  /// `some P1` instead of as a stable reference.
+  const ValueDecl *OpaqueReturnTypeNamingDecl = nullptr;
+
   /// Whether to print decl attributes that are only used internally,
   /// such as _silgen_name, transparent, etc.
   bool PrintUserInaccessibleAttrs = true;

lib/AST/ASTPrinter.cpp

@@ -194,24 +194,6 @@ bool PrintOptions::excludeAttr(const DeclAttribute *DA) const {
   return false;
 }
 
-/// Forces printing types with the `some` keyword, instead of the full stable
-/// reference.
-struct PrintWithOpaqueResultTypeKeywordRAII {
-  PrintWithOpaqueResultTypeKeywordRAII(PrintOptions &Options)
-      : Options(Options) {
-    SavedMode = Options.OpaqueReturnTypePrinting;
-    Options.OpaqueReturnTypePrinting =
-        PrintOptions::OpaqueReturnTypePrintingMode::WithOpaqueKeyword;
-  }
-  ~PrintWithOpaqueResultTypeKeywordRAII() {
-    Options.OpaqueReturnTypePrinting = SavedMode;
-  }
-
-private:
-  PrintOptions &Options;
-  PrintOptions::OpaqueReturnTypePrintingMode SavedMode;
-};
-
 PrintOptions PrintOptions::printSwiftInterfaceFile(ModuleDecl *ModuleToPrint,
                                                    bool preferTypeRepr,
                                                    bool printFullConvention,
@@ -969,10 +951,19 @@ class PrintAST : public ASTVisitor<PrintAST> {
     printTypeLocWithOptions(TL, Options, printBeforeType);
   }
 
-  void printTypeLocForImplicitlyUnwrappedOptional(TypeLoc TL, bool IUO) {
-    PrintOptions options = Options;
-    options.PrintOptionalAsImplicitlyUnwrapped = IUO;
-    printTypeLocWithOptions(TL, options);
+  void printTypeLocForImplicitlyUnwrappedOptional(
+      TypeLoc TL, bool IUO, const ValueDecl *opaqueTypeNamingDecl) {
+    auto savedIOU = Options.PrintOptionalAsImplicitlyUnwrapped;
+    Options.PrintOptionalAsImplicitlyUnwrapped = IUO;
+
+    auto savedOpaqueTypeNamingDecl = Options.OpaqueReturnTypeNamingDecl;
+    if (opaqueTypeNamingDecl)
+      Options.OpaqueReturnTypeNamingDecl = opaqueTypeNamingDecl;
+
+    printTypeLocWithOptions(TL, Options);
+
+    Options.PrintOptionalAsImplicitlyUnwrapped = savedIOU;
+    Options.OpaqueReturnTypeNamingDecl = savedOpaqueTypeNamingDecl;
   }
 
   void printContextIfNeeded(const Decl *decl) {
@@ -1376,18 +1367,17 @@ void PrintAST::printTypedPattern(const TypedPattern *TP) {
   printPattern(TP->getSubPattern());
   Printer << ": ";
 
-  PrintWithOpaqueResultTypeKeywordRAII x(Options);
-
-  // Make sure to check if the underlying var decl is an implicitly unwrapped
-  // optional.
-  bool isIUO = false;
+  VarDecl *varDecl = nullptr;
   if (auto *named = dyn_cast<NamedPattern>(TP->getSubPattern()))
     if (auto decl = named->getDecl())
-      isIUO = decl->isImplicitlyUnwrappedOptional();
+      varDecl = decl;
 
   const auto TyLoc = TypeLoc(TP->getTypeRepr(),
                              TP->hasType() ? TP->getType() : Type());
-  printTypeLocForImplicitlyUnwrappedOptional(TyLoc, isIUO);
+
+  printTypeLocForImplicitlyUnwrappedOptional(
+      TyLoc, varDecl ? varDecl->isImplicitlyUnwrappedOptional() : false,
+      varDecl);
 }
 
 /// Determines if we are required to print the name of a property declaration,
@@ -3882,9 +3872,8 @@ void PrintAST::visitVarDecl(VarDecl *decl) {
     }
     Printer.printDeclResultTypePre(decl, tyLoc);
 
-    PrintWithOpaqueResultTypeKeywordRAII x(Options);
     printTypeLocForImplicitlyUnwrappedOptional(
-      tyLoc, decl->isImplicitlyUnwrappedOptional());
+        tyLoc, decl->isImplicitlyUnwrappedOptional(), decl);
   }
 
   printAccessors(decl);
@@ -3966,7 +3955,7 @@ void PrintAST::printOneParameter(const ParamDecl *param,
     }
 
     printTypeLocForImplicitlyUnwrappedOptional(
-      TheTypeLoc, param->isImplicitlyUnwrappedOptional());
+        TheTypeLoc, param->isImplicitlyUnwrappedOptional(), nullptr);
   }
 
   if (param->isDefaultArgument() && Options.PrintDefaultArgumentValue) {
@@ -4256,8 +4245,6 @@ void PrintAST::visitFuncDecl(FuncDecl *decl) {
         }
       }
 
-      PrintWithOpaqueResultTypeKeywordRAII x(Options);
-
       // Check if we would go down the type repr path... in such a case, see if
       // we can find a type repr and if that type has a sending type repr. In
       // such a case, look through the sending type repr since we handle it here
@@ -4284,7 +4271,7 @@ void PrintAST::visitFuncDecl(FuncDecl *decl) {
       // If we printed using type repr printing, do not print again.
       if (!usedTypeReprPrinting) {
         printTypeLocForImplicitlyUnwrappedOptional(
-            ResultTyLoc, decl->isImplicitlyUnwrappedOptional());
+            ResultTyLoc, decl->isImplicitlyUnwrappedOptional(), decl);
       }
       Printer.printStructurePost(PrintStructureKind::FunctionReturnType);
     }
@@ -4433,9 +4420,8 @@ void PrintAST::visitSubscriptDecl(SubscriptDecl *decl) {
     Printer.printDeclResultTypePre(decl, elementTy);
     Printer.callPrintStructurePre(PrintStructureKind::FunctionReturnType);
 
-    PrintWithOpaqueResultTypeKeywordRAII x(Options);
     printTypeLocForImplicitlyUnwrappedOptional(
-      elementTy, decl->isImplicitlyUnwrappedOptional());
+        elementTy, decl->isImplicitlyUnwrappedOptional(), decl);
     Printer.printStructurePost(PrintStructureKind::FunctionReturnType);
   }
 
@@ -7245,7 +7231,11 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
     auto genericSig = namingDecl->getInnermostDeclContext()
           ->getGenericSignatureOfContext();
 
-    switch (Options.OpaqueReturnTypePrinting) {
+    auto mode = Options.OpaqueReturnTypePrinting;
+    if (Options.OpaqueReturnTypeNamingDecl == T->getDecl()->getNamingDecl())
+      mode = PrintOptions::OpaqueReturnTypePrintingMode::WithOpaqueKeyword;
+
+    switch (mode) {
     case PrintOptions::OpaqueReturnTypePrintingMode::WithOpaqueKeyword:
       if (printNamedOpaque())
         return;

test/ModuleInterface/opaque-result-types.swift

@@ -50,6 +50,7 @@ public protocol AssocTypeInference {
   subscript() -> AssocSubscript { get }
 }
 
+// CHECK-LABEL: public struct Bar<T> : OpaqueResultTypes.AssocTypeInference
 @available(SwiftStdlib 5.1, *)
 public struct Bar<T>: AssocTypeInference {
   public init() {}
@@ -259,3 +260,28 @@ public struct Zim: AssocTypeInference {
     return 123
   }
 }
+
+public protocol PrimaryAssociatedTypeInference<Assoc> {
+  associatedtype Assoc
+
+  func foo(_: Int) -> Assoc
+}
+
+// CHECK-LABEL: public struct Baz : OpaqueResultTypes.PrimaryAssociatedTypeInference
+
+public struct Baz: PrimaryAssociatedTypeInference {
+  // CHECK-LABEL: public func foo(_: Swift.Int) -> some OpaqueResultTypes.Foo
+  public func foo(_: Int) -> some Foo {
+    return 123
+  }
+
+  // CHECK-LABEL: public func callsFoo() -> @_opaqueReturnTypeOf("$s17OpaqueResultTypes3BazV3fooyQrSiF", 0) __
+  public func callsFoo() -> Assoc {
+    return foo(123)
+  }
+
+  // CHECK-LABEL: public func identity() -> some OpaqueResultTypes.PrimaryAssociatedTypeInference<@_opaqueReturnTypeOf("$s17OpaqueResultTypes3BazV3fooyQrSiF", 0) __>
+  public func identity() -> some PrimaryAssociatedTypeInference<Assoc> {
+    return self
+  }
+}