[codecompletion] Add @escaping to override completions

Flush out the ASTPrinter's ability to exclude and include specific
attributes to cover TypeAttrKinds and have code-completion use this to
print @escaping in override completions.  Incidentally fix a case where
we weren't forwarding important options after type transformation, which
prevented printing @escaping in transformed parameter types.

rdar://problem/27772722

Author: benlangmuir

include/swift/AST/PrintOptions.h

@@ -110,6 +110,36 @@ class BracketOptions {
   }
 };
 
+/// A union of DeclAttrKind and TypeAttrKind.
+class AnyAttrKind {
+  unsigned kind : 31;
+  unsigned isType : 1;
+
+public:
+  AnyAttrKind(TypeAttrKind K) : kind(static_cast<unsigned>(K)), isType(1) {
+    static_assert(TAK_Count < UINT_MAX, "TypeAttrKind is > 31 bits");
+  }
+  AnyAttrKind(DeclAttrKind K) : kind(static_cast<unsigned>(K)), isType(0) {
+    static_assert(DAK_Count < UINT_MAX, "DeclAttrKind is > 31 bits");
+  }
+  AnyAttrKind() : kind(TAK_Count), isType(1) {}
+  AnyAttrKind(const AnyAttrKind &) = default;
+
+  /// Returns the TypeAttrKind, or TAK_Count if this is not a type attribute.
+  TypeAttrKind type() const {
+    return isType ? static_cast<TypeAttrKind>(kind) : TAK_Count;
+  }
+  /// Returns the DeclAttrKind, or DAK_Count if this is not a decl attribute.
+  DeclAttrKind decl() const {
+    return isType ? DAK_Count : static_cast<DeclAttrKind>(kind);
+  }
+
+  bool operator==(AnyAttrKind K) const {
+    return kind == K.kind && isType == K.isType;
+  }
+  bool operator!=(AnyAttrKind K) const { return !(*this == K); }
+};
+
 /// Options for printing AST nodes.
 ///
 /// A default-constructed PrintOptions is suitable for printing to users;
@@ -225,12 +255,12 @@ struct PrintOptions {
   bool PrintUserInaccessibleAttrs = true;
 
   /// List of attribute kinds that should not be printed.
-  std::vector<DeclAttrKind> ExcludeAttrList =
+  std::vector<AnyAttrKind> ExcludeAttrList =
       { DAK_Transparent, DAK_Effects, DAK_FixedLayout };
 
   /// List of attribute kinds that should be printed exclusively.
   /// Empty means allow all.
-  std::vector<DeclAttrKind> ExclusiveAttrList;
+  std::vector<AnyAttrKind> ExclusiveAttrList;
 
   /// Whether to print function @convention attribute on function types.
   bool PrintFunctionRepresentationAttrs = true;
@@ -326,6 +356,16 @@ struct PrintOptions {
 
   BracketOptions BracketOptions;
 
+  bool excludeAttrKind(AnyAttrKind K) const {
+    if (std::any_of(ExcludeAttrList.begin(), ExcludeAttrList.end(),
+                    [K](AnyAttrKind other) { return other == K; }))
+      return true;
+    if (!ExclusiveAttrList.empty())
+      return std::none_of(ExclusiveAttrList.begin(), ExclusiveAttrList.end(),
+                          [K](AnyAttrKind other) { return other == K; });
+    return false;
+  }
+
   /// Retrieve the set of options for verbose printing to users.
   static PrintOptions printVerbose() {
     PrintOptions result;

include/swift/AST/TypeRepr.h

@@ -166,7 +166,7 @@ class AttributedTypeRepr : public TypeRepr {
   TypeRepr *getTypeRepr() const { return Ty; }
 
   void printAttrs(llvm::raw_ostream &OS) const;
-  void printAttrs(ASTPrinter &Printer) const;
+  void printAttrs(ASTPrinter &Printer, const PrintOptions &Options) const;
 
   static bool classof(const TypeRepr *T) {
     return T->getKind() == TypeReprKind::Attributed;

lib/AST/ASTPrinter.cpp

@@ -1172,7 +1172,12 @@ class PrintAST : public ASTVisitor<PrintAST> {
   void printTypeLoc(const TypeLoc &TL) {
     if (Options.TransformContext && TL.getType()) {
       if (auto RT = Options.TransformContext->transform(TL.getType())) {
+        // FIXME: it's not clear exactly what we want to keep from the existing
+        // options, and what we want to discard.
         PrintOptions FreshOptions;
+        FreshOptions.PrintAsInParamType = Options.PrintAsInParamType;
+        FreshOptions.ExcludeAttrList = Options.ExcludeAttrList;
+        FreshOptions.ExclusiveAttrList = Options.ExclusiveAttrList;
         RT.print(Printer, FreshOptions);
         return;
       }
@@ -2527,10 +2532,8 @@ void PrintAST::visitVarDecl(VarDecl *decl) {
 
 void PrintAST::visitParamDecl(ParamDecl *decl) {
   // Set and restore in-parameter-position printing of types
-  auto prior = Options.PrintAsInParamType;
-  Options.PrintAsInParamType = true;
+  llvm::SaveAndRestore<bool> savePrintParam(Options.PrintAsInParamType, true);
   visitVarDecl(decl);
-  Options.PrintAsInParamType = prior;
 }
 
 void PrintAST::printOneParameter(const ParamDecl *param, bool Curried,
@@ -2585,10 +2588,8 @@ void PrintAST::printOneParameter(const ParamDecl *param, bool Curried,
   }
 
   // Set and restore in-parameter-position printing of types
-  auto prior = Options.PrintAsInParamType;
-  Options.PrintAsInParamType = true;
+  llvm::SaveAndRestore<bool> savePrintParam(Options.PrintAsInParamType, true);
   printTypeLoc(TheTypeLoc);
-  Options.PrintAsInParamType = prior;
 
   if (param->isVariadic())
     Printer << "...";
@@ -3763,16 +3764,18 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
     if (Options.SkipAttributes)
       return;
 
-    if (info.isAutoClosure()) {
+    if (info.isAutoClosure() && !Options.excludeAttrKind(TAK_autoclosure)) {
       Printer.printAttrName("@autoclosure");
       Printer << " ";
     }
-    if (inParameterPrinting && !info.isNoEscape()) {
+    if (inParameterPrinting && !info.isNoEscape() &&
+        !Options.excludeAttrKind(TAK_escaping)) {
       Printer.printAttrName("@escaping");
       Printer << " ";
     }
 
-    if (Options.PrintFunctionRepresentationAttrs) {
+    if (Options.PrintFunctionRepresentationAttrs &&
+        !Options.excludeAttrKind(TAK_convention)) {
       // TODO: coalesce into a single convention attribute.
       switch (info.getSILRepresentation()) {
       case SILFunctionType::Representation::Thick:

lib/AST/Attr.cpp

@@ -249,16 +249,8 @@ void DeclAttributes::print(ASTPrinter &Printer,
     if (!Options.PrintUserInaccessibleAttrs &&
         DeclAttribute::isUserInaccessible(DA->getKind()))
       continue;
-    if (std::find(Options.ExcludeAttrList.begin(),
-                  Options.ExcludeAttrList.end(),
-                  DA->getKind()) != Options.ExcludeAttrList.end())
+    if (Options.excludeAttrKind(DA->getKind()))
       continue;
-    if (!Options.ExclusiveAttrList.empty()) {
-      if (std::find(Options.ExclusiveAttrList.begin(),
-                    Options.ExclusiveAttrList.end(),
-                    DA->getKind()) == Options.ExclusiveAttrList.end())
-        continue;
-    }
 
     AttributeVector &which = DA->isDeclModifier() ? modifiers :
                              isShortAvailable(DA) ? shortAvailableAttributes :

lib/AST/TypeRepr.cpp

@@ -264,36 +264,43 @@ void ErrorTypeRepr::printImpl(ASTPrinter &Printer,
 
 void AttributedTypeRepr::printImpl(ASTPrinter &Printer,
                                    const PrintOptions &Opts) const {
-  printAttrs(Printer);
+  printAttrs(Printer, Opts);
   printTypeRepr(Ty, Printer, Opts);
 }
 
 void AttributedTypeRepr::printAttrs(llvm::raw_ostream &OS) const {
   StreamPrinter Printer(OS);
-  printAttrs(Printer);
+  printAttrs(Printer, PrintOptions());
 }
 
-void AttributedTypeRepr::printAttrs(ASTPrinter &Printer) const {
+void AttributedTypeRepr::printAttrs(ASTPrinter &Printer,
+                                    const PrintOptions &Options) const {
   const TypeAttributes &Attrs = getAttrs();
 
-  if (Attrs.has(TAK_autoclosure)) {
+  auto hasAttr = [&](TypeAttrKind K) -> bool {
+    if (Options.excludeAttrKind(K))
+      return false;
+    return Attrs.has(K);
+  };
+
+  if (hasAttr(TAK_autoclosure)) {
     Printer.printAttrName("@autoclosure");
     Printer << " ";
   }
-  if (Attrs.has(TAK_escaping)) {
+  if (hasAttr(TAK_escaping)) {
     Printer.printAttrName("@escaping");
     Printer << " ";
   }
 
-  if (Attrs.has(TAK_thin)) {
+  if (hasAttr(TAK_thin)) {
     Printer.printAttrName("@thin");
     Printer << " ";
   }
-  if (Attrs.has(TAK_thick)) {
+  if (hasAttr(TAK_thick)) {
     Printer.printAttrName("@thick");
     Printer << " ";
   }
-  if (Attrs.convention.hasValue()) {
+  if (hasAttr(TAK_convention) && Attrs.convention.hasValue()) {
     Printer.printAttrName("@convention");
     Printer << "(" << Attrs.convention.getValue() << ") ";
   }

lib/IDE/CodeCompletion.cpp

@@ -4117,7 +4117,7 @@ class CompletionOverrideLookup : public swift::VisibleDeclConsumer {
         Options.setArchetypeSelfTransform(transformType, VD->getDeclContext());
       Options.PrintDefaultParameterPlaceholder = false;
       Options.PrintImplicitAttrs = false;
-      Options.SkipAttributes = true;
+      Options.ExclusiveAttrList.push_back(TAK_escaping);
       Options.PrintOverrideKeyword = false;
       Options.PrintPropertyAccessors = false;
       VD->print(Printer, Options);

test/IDE/complete_override.swift

@@ -105,6 +105,7 @@
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=ASSOC_TYPE1 -code-completion-keywords=false | %FileCheck %s -check-prefix=ASSOC_TYPE1
 
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=DEPRECATED_1 -code-completion-keywords=false | %FileCheck %s -check-prefix=DEPRECATED_1
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=ESCAPING_1 -code-completion-keywords=false | %FileCheck %s -check-prefix=ESCAPING_1
 
 @objc
 class TagPA {}
@@ -463,8 +464,7 @@ class TestClassWithThrows : HasThrowing, HasThrowingProtocol {
 // HAS_THROWING: Begin completions
 // HAS_THROWING-DAG: Decl[InstanceMethod]/Super:         func foo() throws {|}; name=foo() throws
 // HAS_THROWING-DAG: Decl[InstanceMethod]/Super:         override func bar() throws {|}; name=bar() throws
-// FIXME: SR-2214 make the below require printing @escaping 
-// HAS_THROWING-DAG: Decl[InstanceMethod]/Super:         override func baz(x: {{(@escaping )?}}() throws -> ()) rethrows {|}; name=baz(x: {{(@escaping )?}}() throws -> ()) rethrows
+// HAS_THROWING-DAG: Decl[InstanceMethod]/Super:         override func baz(x: @escaping () throws -> ()) rethrows {|}; name=baz(x: {{(@escaping )?}}() throws -> ()) rethrows
 // HAS_THROWING-DAG: Decl[Constructor]/Super:            init() throws {|}; name=init() throws
 // HAS_THROWING: End completions
 
@@ -490,3 +490,11 @@ class Deprecated2 : Deprecated1 {
   override func #^DEPRECATED_1^#
 }
 // DEPRECATED_1: Decl[InstanceMethod]/Super/NotRecommended: deprecated() {|};
+
+class EscapingBase {
+  func method(_ x: @escaping (@escaping ()->()) -> (@escaping ()->())) -> (@escaping (@escaping ()->() )->()) { }
+}
+class Escaping : EscapingBase {
+  override func #^ESCAPING_1^#
+}
+// ESCAPING_1: Decl[InstanceMethod]/Super:         method(_ x: @escaping (@escaping () -> ()) -> (@escaping () -> ())) -> ((@escaping () -> ()) -> ()) {|};