[CodeCompletion] Use opaque type for override completion if preferable

rdar://problem/49354106

Author: rintaro

include/swift/AST/ASTPrinter.h

@@ -118,6 +118,10 @@ class ASTPrinter {
   /// Callers should use callPrintDeclPost().
   virtual void printDeclPost(const Decl *D, Optional<BracketOptions> Bracket) {}
 
+  /// Called before printing the result type of the declaration. Printer can
+  /// replace \p TL to customize the input.
+  virtual void printDeclResultTypePre(ValueDecl *VD, TypeLoc &TL) {}
+
   /// Called before printing a type.
   virtual void printTypePre(const TypeLoc &TL) {}
   /// Called after printing a type.

include/swift/AST/PrintOptions.h

@@ -135,6 +135,9 @@ struct PrintOptions {
   /// Whether to print *any* accessors on properties.
   bool PrintPropertyAccessors = true;
 
+  /// Whether to print *any* accessors on subscript.
+  bool PrintSubscriptAccessors = true;
+
   /// Whether to print the accessors of a property abstractly,
   /// i.e. always as:
   /// ```

include/swift/Parse/Parser.h

@@ -1028,6 +1028,8 @@ class Parser {
   ParserResult<PrecedenceGroupDecl>
   parseDeclPrecedenceGroup(ParseDeclOptions flags, DeclAttributes &attributes);
 
+  ParserResult<TypeRepr> parseDeclResultType(Diag<> MessageID);
+
   //===--------------------------------------------------------------------===//
   // Type Parsing
 

lib/AST/ASTPrinter.cpp

@@ -1714,6 +1714,8 @@ void PrintAST::printMutatingModifiersIfNeeded(const AccessorDecl *accessor) {
 void PrintAST::printAccessors(const AbstractStorageDecl *ASD) {
   if (isa<VarDecl>(ASD) && !Options.PrintPropertyAccessors)
     return;
+  if (isa<SubscriptDecl>(ASD) && !Options.PrintSubscriptAccessors)
+    return;
 
   auto impl = ASD->getImplInfo();
 
@@ -2534,6 +2536,7 @@ void PrintAST::visitVarDecl(VarDecl *decl) {
     if (!tyLoc.getTypeRepr())
       tyLoc = TypeLoc::withoutLoc(decl->getInterfaceType());
 
+    Printer.printDeclResultTypePre(decl, tyLoc);
     if (decl->getAttrs().hasAttribute<ImplicitlyUnwrappedOptionalAttr>())
       printTypeLocForImplicitlyUnwrappedOptional(tyLoc);
     else
@@ -2804,6 +2807,8 @@ void PrintAST::visitFuncDecl(FuncDecl *decl) {
           ResultTyLoc = TypeLoc::withoutLoc(ResultTy);
       }
       Printer << " -> ";
+
+      Printer.printDeclResultTypePre(decl, ResultTyLoc);
       Printer.callPrintStructurePre(PrintStructureKind::FunctionReturnType);
       if (decl->getAttrs().hasAttribute<ImplicitlyUnwrappedOptionalAttr>())
         printTypeLocForImplicitlyUnwrappedOptional(ResultTyLoc);
@@ -2931,8 +2936,9 @@ void PrintAST::visitSubscriptDecl(SubscriptDecl *decl) {
   });
   Printer << " -> ";
 
-  Printer.callPrintStructurePre(PrintStructureKind::FunctionReturnType);
   TypeLoc elementTy = decl->getElementTypeLoc();
+  Printer.printDeclResultTypePre(decl, elementTy);
+  Printer.callPrintStructurePre(PrintStructureKind::FunctionReturnType);
   if (!elementTy.getTypeRepr())
     elementTy = TypeLoc::withoutLoc(decl->getElementInterfaceType());
   if (decl->getAttrs().hasAttribute<ImplicitlyUnwrappedOptionalAttr>())

lib/IDE/CodeCompletion.cpp

@@ -4038,27 +4038,82 @@ class CompletionOverrideLookup : public swift::VisibleDeclConsumer {
       Builder.addAccessControlKeyword(Access);
   }
 
+  /// Return type if the result type if \p VD should be represented as opaque
+  /// result type.
+  TypeLoc getOpaqueResultTypeLoc(const ValueDecl *VD, DeclVisibilityKind Reason) {
+    if (Reason !=
+        DeclVisibilityKind::MemberOfProtocolImplementedByCurrentNominal)
+      return nullptr;
+
+    auto currTy = CurrDeclContext->getDeclaredTypeInContext();
+    if (!currTy)
+      return nullptr;
+
+    Type ResultT;
+    if (auto *FD = dyn_cast<FuncDecl>(VD))
+      ResultT = FD->getResultInterfaceType();
+    else if (auto *SD = dyn_cast<SubscriptDecl>(VD))
+      ResultT = SD->getElementInterfaceType();
+    else if (auto *VarD = dyn_cast<VarDecl>(VD))
+      ResultT = VarD->getInterfaceType();
+    else
+      return nullptr;
+
+    if (!ResultT->is<DependentMemberType>())
+      // The result is not associatedtype.
+      return nullptr;
+
+    // If associatedtype doesn't have conformance/superclass constraint, we
+    // can't use opaque type.
+    auto assocTyD = ResultT->castTo<DependentMemberType>()->getAssocType();
+    if (!assocTyD->getInherited().size())
+      return nullptr;
+
+    // Try substitution to see if the associated type is resolved to concrete
+    // type.
+    auto substMap = currTy->getMemberSubstitutionMap(
+        CurrDeclContext->getParentModule(), VD);
+    ResultT = ResultT.subst(substMap, SubstFlags::UseErrorType);
+    if (!ResultT || !ResultT->is<DependentMemberType>())
+      // If resolved print it.
+      return nullptr;
+
+    return assocTyD->getInherited()[0];
+  }
+
   void addValueOverride(const ValueDecl *VD, DeclVisibilityKind Reason,
                         CodeCompletionResultBuilder &Builder,
                         bool hasDeclIntroducer) {
+    class DeclPrinter : public StreamPrinter {
+      TypeLoc OpaqueBaseTy;
 
-    class DeclNameOffsetLocatorPrinter : public StreamPrinter {
     public:
       using StreamPrinter::StreamPrinter;
 
       Optional<unsigned> NameOffset;
 
+      DeclPrinter(raw_ostream &OS, TypeLoc OpaqueBaseTy)
+          : StreamPrinter(OS), OpaqueBaseTy(OpaqueBaseTy) {}
+
       void printDeclLoc(const Decl *D) override {
         if (!NameOffset.hasValue())
           NameOffset = OS.tell();
       }
+
+      // As for FuncDecl, SubscriptDecl, and VarDecl,
+      void printDeclResultTypePre(ValueDecl *VD, TypeLoc &TL) override {
+        if (!OpaqueBaseTy.isNull()) {
+          OS << "some ";
+          TL = OpaqueBaseTy;
+        }
+      }
     };
 
     llvm::SmallString<256> DeclStr;
     unsigned NameOffset = 0;
     {
       llvm::raw_svector_ostream OS(DeclStr);
-      DeclNameOffsetLocatorPrinter Printer(OS);
+      DeclPrinter Printer(OS, getOpaqueResultTypeLoc(VD, Reason));
       PrintOptions Options;
       if (auto transformType = CurrDeclContext->getDeclaredTypeInContext())
         Options.setBaseType(transformType);
@@ -4067,6 +4122,7 @@ class CompletionOverrideLookup : public swift::VisibleDeclConsumer {
       Options.ExclusiveAttrList.push_back(TAK_autoclosure);
       Options.PrintOverrideKeyword = false;
       Options.PrintPropertyAccessors = false;
+      Options.PrintSubscriptAccessors = false;
       Options.PrintStaticKeyword = !hasStaticOrClass;
       VD->print(Printer, Options);
       NameOffset = Printer.NameOffset.getValue();
@@ -4080,7 +4136,7 @@ class CompletionOverrideLookup : public swift::VisibleDeclConsumer {
         Builder.addOverrideKeyword();
       else {
         auto dist = Ctx.SourceMgr.getByteDistance(
-                      introducerLoc, Ctx.SourceMgr.getCodeCompletionLoc());
+            introducerLoc, Ctx.SourceMgr.getCodeCompletionLoc());
         Builder.setNumBytesToErase(dist);
         Builder.addOverrideKeyword();
         Builder.addDeclIntroducer(DeclStr.str().substr(0, NameOffset));
@@ -4118,6 +4174,15 @@ class CompletionOverrideLookup : public swift::VisibleDeclConsumer {
     addValueOverride(VD, Reason, Builder, hasVarIntroducer);
   }
 
+  void addSubscriptOverride(const SubscriptDecl *SD, DeclVisibilityKind Reason) {
+    CodeCompletionResultBuilder Builder(
+        Sink, CodeCompletionResult::ResultKind::Declaration,
+        SemanticContextKind::Super, {});
+    Builder.setAssociatedDecl(SD);
+    addValueOverride(SD, Reason, Builder, false);
+    Builder.addBraceStmtWithCursor();
+  }
+
   void addTypeAlias(const AssociatedTypeDecl *ATD, DeclVisibilityKind Reason) {
     CodeCompletionResultBuilder Builder(Sink,
       CodeCompletionResult::ResultKind::Declaration,
@@ -4222,6 +4287,11 @@ class CompletionOverrideLookup : public swift::VisibleDeclConsumer {
       return;
     }
 
+    if (auto *SD = dyn_cast<SubscriptDecl>(D)) {
+      if (!hasIntroducer && !hasInitializerModifier)
+        addSubscriptOverride(SD, Reason);
+    }
+
     if (auto *CD = dyn_cast<ConstructorDecl>(D)) {
       if (!isa<ProtocolDecl>(CD->getDeclContext()))
         return;

test/IDE/complete_opaque_result.swift

@@ -15,10 +15,21 @@
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=STRUCT_SUBSCRIPT | %FileCheck %s -check-prefix=BEGINNING_WITH_SOME
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=STRUCT_TYPEALIAS_RHS | %FileCheck %s -check-prefix=BEGINNING_WITHOUT_SOME
 
-protocol MyProtocol {}
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=OVERRIDE_TestClass | %FileCheck %s -check-prefix=OVERRIDE
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=OVERRIDE_TestStruct | %FileCheck %s -check-prefix=OVERRIDE
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=OVERRIDE_HasTypealias | %FileCheck %s -check-prefix=OVERRIDE_HasTypealias
+
+protocol MyProtocol {
+  associatedtype Mistery
+}
 struct MyStruct {}
 enum MyEnum {}
 class MyClass {}
+struct ConcreteMyProtocol : MyProtocol {
+  typealias Mistery = MyStruct
+}
+
+// MARK: 'some' keyword.
 
 // BEGINNING_WITH_SOME: Begin completions
 // BEGINNING_WITH_SOME-DAG: Keyword/None:                       some[#some#]; name=some
@@ -43,8 +54,6 @@ func gloabalFunc() -> #^GLOBAL_FUNC^#
 var globalVar: #^GLOBAL_VAR^#
 
 protocol SomeProto {
-  associatedtype Assoc: MyProtocol
-
   associatedtype protoAssocTy: #^PROTOCOL_ASSOCIATEDTYPE^#
   func protoMethodReq() -> #^PROTOCOL_METHOD_REQUIREMENT^#
   var protoVarReq: #^PROTOCOL_VAR_REQUIREMENT^#
@@ -63,3 +72,67 @@ struct SomeStruct {
   var structVarExt: #^STRUCT_VAR^#
   subscript(struct: Int) -> #^STRUCT_SUBSCRIPT^#
 }
+
+// MARK: Conformance.
+
+protocol HasAssocPlain {
+  associatedtype AssocPlain
+  func returnAssocPlain() -> AssocPlain
+}
+protocol HasAssocWithConformanceConstraint {
+  associatedtype AssocWithConformanceConstraint: MyProtocol
+  func returnAssocWithConformanceConstraint(fn: (Int) -> Int) -> AssocWithConformanceConstraint
+}
+protocol HasAssocWithSuperClassConstraint {
+  associatedtype AssocWithSuperClassConstraint: MyClass
+  var valAssocWithSuperClassConstraint: AssocWithSuperClassConstraint { get }
+}
+protocol HasAssocWithCompositionConstraint {
+  associatedtype AssocWithCompositionConstraint: MyClass & MyProtocol
+  subscript<T>(idx: T) -> AssocWithCompositionConstraint where T: Comparable { get }
+}
+protocol HasAssocWithDefault {
+  associatedtype AssocWithDefault = MyEnum
+  func returnAssocWithDefault() -> AssocWithDefault
+}
+protocol HasAssocWithConstraintAndDefault {
+  associatedtype AssocWithConstraintAndDefault: MyProtocol = ConcreteMyProtocol
+  func returnAssocWithConstraintAndDefault() -> AssocWithConstraintAndDefault
+}
+
+class TestClass :
+    HasAssocPlain,
+    HasAssocWithConformanceConstraint,
+    HasAssocWithSuperClassConstraint,
+    HasAssocWithCompositionConstraint,
+    HasAssocWithDefault,
+    HasAssocWithConstraintAndDefault {
+  #^OVERRIDE_TestClass^#
+// OVERRIDE: found code completion token OVERRIDE_[[BASETYPE:[A-Za-z0-9]+]] at
+// OVERRIDE: Begin completions
+// OVERRIDE-DAG: Decl[InstanceMethod]/Super:         func returnAssocPlain() -> [[BASETYPE]].AssocPlain {|};
+// OVERRIDE-DAG: Decl[InstanceMethod]/Super:         func returnAssocWithConformanceConstraint(fn: (Int) -> Int) -> some MyProtocol {|};
+// OVERRIDE-DAG: Decl[InstanceVar]/Super:            var valAssocWithSuperClassConstraint: some MyClass;
+// OVERRIDE-DAG: Decl[Subscript]/Super:              subscript<T>(idx: T) -> some MyClass & MyProtocol where T : Comparable {|};
+// OVERRIDE-DAG: Decl[InstanceMethod]/Super:         func returnAssocWithDefault() -> MyEnum {|};
+// OVERRIDE-DAG: Decl[InstanceMethod]/Super:         func returnAssocWithConstraintAndDefault() -> ConcreteMyProtocol {|};
+// OVERRIDE: End completions
+}
+
+struct TestStruct :
+    HasAssocPlain,
+    HasAssocWithConformanceConstraint,
+    HasAssocWithSuperClassConstraint,
+    HasAssocWithCompositionConstraint,
+    HasAssocWithDefault,
+    HasAssocWithConstraintAndDefault {
+  #^OVERRIDE_TestStruct^#
+}
+
+class HasTypealias : HasAssocWithConformanceConstraint {
+  typealias AssocWithConformanceConstraint = ConcreteMyProtocol
+  #^OVERRIDE_HasTypealias^#
+// OVERRIDE_HasTypealias: Begin completions
+// OVERRIDE_HasTypealias-DAG: Decl[InstanceMethod]/Super:         func returnAssocWithConformanceConstraint(fn: (Int) -> Int) -> ConcreteMyProtocol {|};
+// OVERRIDE_HasTypealias: End completions
+}