SourceKit: allow expression type request to specify a list of protocol USRs for filtering

The client usually cares about a subset of all expressions. A way to differentiate
them is by the protocols these expressions' types conform to. This patch allows
the request to add a list of protocol USRs so that the response only includes those
interested expressions that conform to any of the input protocols.

We also add a field to the response for each expression type to indicate the
conforming protocols names that were originally in the input list.

When an empty list of protocol USRs are given, we report all expressions' types
in the file like the old behavior.

rdar://35199889

Author: nkcsgexi

include/swift/Sema/IDETypeChecking.h

@@ -19,6 +19,7 @@
 #ifndef SWIFT_SEMA_IDETYPECHECKING_H
 #define SWIFT_SEMA_IDETYPECHECKING_H
 
+#include "llvm/ADT/MapVector.h"
 #include "swift/Basic/SourceLoc.h"
 #include <memory>
 
@@ -177,12 +178,21 @@ namespace swift {
 
     /// The length of the printed type
     uint32_t typeLength;
+
+    /// The offsets and lengths of all protocols the type conforms to
+    std::vector<std::pair<uint32_t, uint32_t>> protocols;
   };
 
   /// Collect type information for every expression in \c SF; all types will
   /// be printed to \c OS.
   ArrayRef<ExpressionTypeInfo> collectExpressionType(SourceFile &SF,
+    ArrayRef<const char *> ExpectedProtocols,
     std::vector<ExpressionTypeInfo> &scratch, llvm::raw_ostream &OS);
+
+  /// Resolve a list of mangled names to accessible protocol decls from
+  /// the decl context.
+  bool resolveProtocolNames(DeclContext *DC, ArrayRef<const char *> names,
+                            llvm::MapVector<ProtocolDecl*, StringRef> &result);
 }
 
 #endif

lib/IDE/ConformingMethodList.cpp

@@ -32,9 +32,8 @@ class ConformingMethodListCallbacks : public CodeCompletionCallbacks {
   Expr *ParsedExpr = nullptr;
   DeclContext *CurDeclContext = nullptr;
 
-  void resolveExpectedTypes(ArrayRef<const char *> names, SourceLoc loc,
-                            SmallVectorImpl<ProtocolDecl *> &result);
-  void getMatchingMethods(Type T, ArrayRef<ProtocolDecl *> expectedTypes,
+  void getMatchingMethods(Type T,
+                          llvm::MapVector<ProtocolDecl*, StringRef> &expectedTypes,
                           SmallVectorImpl<ValueDecl *> &result);
 
 public:
@@ -89,9 +88,8 @@ void ConformingMethodListCallbacks::doneParsing() {
   if (!T || T->is<ErrorType>() || T->is<UnresolvedType>())
     return;
 
-  SmallVector<ProtocolDecl *, 4> expectedProtocols;
-  resolveExpectedTypes(ExpectedTypeNames, ParsedExpr->getLoc(),
-                       expectedProtocols);
+  llvm::MapVector<ProtocolDecl*, StringRef> expectedProtocols;
+  resolveProtocolNames(CurDeclContext, ExpectedTypeNames, expectedProtocols);
 
   // Collect the matching methods.
   ConformingMethodListResult result(CurDeclContext, T);
@@ -100,21 +98,8 @@ void ConformingMethodListCallbacks::doneParsing() {
   Consumer.handleResult(result);
 }
 
-void ConformingMethodListCallbacks::resolveExpectedTypes(
-    ArrayRef<const char *> names, SourceLoc loc,
-    SmallVectorImpl<ProtocolDecl *> &result) {
-  auto &ctx = CurDeclContext->getASTContext();
-
-  for (auto name : names) {
-    if (auto ty = Demangle::getTypeForMangling(ctx, name)) {
-      if (auto Proto = dyn_cast_or_null<ProtocolDecl>(ty->getAnyGeneric()))
-        result.push_back(Proto);
-    }
-  }
-}
-
 void ConformingMethodListCallbacks::getMatchingMethods(
-    Type T, ArrayRef<ProtocolDecl *> expectedTypes,
+    Type T, llvm::MapVector<ProtocolDecl*, StringRef> &expectedTypes,
     SmallVectorImpl<ValueDecl *> &result) {
   if (!T->mayHaveMembers())
     return;
@@ -126,7 +111,7 @@ void ConformingMethodListCallbacks::getMatchingMethods(
     Type T;
 
     /// The list of expected types.
-    ArrayRef<ProtocolDecl *> ExpectedTypes;
+    llvm::MapVector<ProtocolDecl*, StringRef> &ExpectedTypes;
 
     /// Result sink to populate.
     SmallVectorImpl<ValueDecl *> &Result;
@@ -149,7 +134,7 @@ void ConformingMethodListCallbacks::getMatchingMethods(
 
       // The return type conforms to any of the requested protocols.
       for (auto Proto : ExpectedTypes) {
-        if (CurModule->conformsToProtocol(declTy, Proto))
+        if (CurModule->conformsToProtocol(declTy, Proto.first))
           return true;
       }
 
@@ -158,7 +143,7 @@ void ConformingMethodListCallbacks::getMatchingMethods(
 
   public:
     LocalConsumer(DeclContext *DC, Type T,
-                  ArrayRef<ProtocolDecl *> expectedTypes,
+                  llvm::MapVector<ProtocolDecl*, StringRef> &expectedTypes,
                   SmallVectorImpl<ValueDecl *> &result)
         : CurModule(DC->getParentModule()), T(T), ExpectedTypes(expectedTypes),
           Result(result) {}

lib/IDE/IDETypeChecking.cpp

@@ -21,6 +21,7 @@
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/NameLookup.h"
+#include "swift/AST/ASTDemangler.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/Sema/IDETypeChecking.h"
 #include "swift/IDE/SourceEntityWalker.h"
@@ -582,6 +583,7 @@ collectDefaultImplementationForProtocolMembers(ProtocolDecl *PD,
 
 /// This walker will traverse the AST and report types for every expression.
 class ExpressionTypeCollector: public SourceEntityWalker {
+  ModuleDecl &Module;
   SourceManager &SM;
   unsigned int BufferId;
   std::vector<ExpressionTypeInfo> &Results;
@@ -596,7 +598,13 @@ class ExpressionTypeCollector: public SourceEntityWalker {
   // [offset, length].
   llvm::DenseMap<unsigned, llvm::DenseSet<unsigned>> AllPrintedTypes;
 
-  bool shouldReport(unsigned Offset, unsigned Length, Expr *E) {
+  // When non empty, we only print expression types that conform to any of
+  // these protocols.
+  llvm::MapVector<ProtocolDecl*, StringRef> &InterestedProtocols;
+
+  bool shouldReport(unsigned Offset, unsigned Length, Expr *E,
+                    std::vector<StringRef> &Conformances) {
+    assert(Conformances.empty());
     // We shouldn't report null types.
     if (E->getType().isNull())
       return false;
@@ -605,58 +613,116 @@ class ExpressionTypeCollector: public SourceEntityWalker {
     // report again. This makes sure we always report the outtermost type of
     // several overlapping expressions.
     auto &Bucket = AllPrintedTypes[Offset];
-    return Bucket.find(Length) == Bucket.end();
+    if (Bucket.find(Length) != Bucket.end())
+      return false;
+
+    // We print every expression if the interested protocols are empty.
+    if (InterestedProtocols.empty())
+      return true;
+
+    // Collecting protocols conformed by this expressions that are in the list.
+    for (auto Proto: InterestedProtocols) {
+      if (Module.conformsToProtocol(E->getType(), Proto.first)) {
+        Conformances.push_back(Proto.second);
+      }
+    }
+
+    // We only print the type of the expression if it conforms to any of the
+    // interested protocols.
+    return !Conformances.empty();
   }
 
   // Find an existing offset in the type buffer otherwise print the type to
   // the buffer.
-  uint32_t getTypeOffsets(StringRef PrintedType) {
+  std::pair<uint32_t, uint32_t> getTypeOffsets(StringRef PrintedType) {
     auto It = TypeOffsets.find(PrintedType);
     if (It == TypeOffsets.end()) {
       TypeOffsets[PrintedType] = OS.tell();
-      OS << PrintedType;
+      OS << PrintedType << '\0';
     }
-    return TypeOffsets[PrintedType];
+    return {TypeOffsets[PrintedType], PrintedType.size()};
   }
 
+
 public:
-  ExpressionTypeCollector(SourceFile &SF, std::vector<ExpressionTypeInfo> &Results,
-    llvm::raw_ostream &OS): SM(SF.getASTContext().SourceMgr),
+  ExpressionTypeCollector(SourceFile &SF,
+              llvm::MapVector<ProtocolDecl*, StringRef> &InterestedProtocols,
+                          std::vector<ExpressionTypeInfo> &Results,
+                          llvm::raw_ostream &OS): Module(*SF.getParentModule()),
+                            SM(SF.getASTContext().SourceMgr),
                             BufferId(*SF.getBufferID()),
-                            Results(Results), OS(OS) {}
+                            Results(Results), OS(OS),
+                            InterestedProtocols(InterestedProtocols) {}
   bool walkToExprPre(Expr *E) override {
     if (E->getSourceRange().isInvalid())
       return true;
     CharSourceRange Range =
       Lexer::getCharSourceRangeFromSourceRange(SM, E->getSourceRange());
     unsigned Offset = SM.getLocOffsetInBuffer(Range.getStart(), BufferId);
     unsigned Length = Range.getByteLength();
-    if (!shouldReport(Offset, Length, E))
+    std::vector<StringRef> Conformances;
+    if (!shouldReport(Offset, Length, E, Conformances))
       return true;
     // Print the type to a temporary buffer.
     SmallString<64> Buffer;
     {
       llvm::raw_svector_ostream OS(Buffer);
       E->getType()->getRValueType()->reconstituteSugar(true)->print(OS);
-      // Ensure the end user can directly use the char*
-      OS << '\0';
     }
-
+    auto Ty = getTypeOffsets(Buffer.str());
     // Add the type information to the result list.
-    Results.push_back({Offset, Length, getTypeOffsets(Buffer.str()),
-      static_cast<uint32_t>(Buffer.size()) - 1});
+    Results.push_back({Offset, Length, Ty.first, Ty.second, {}});
+
+    // Adding all protocol names to the result.
+    for(auto Con: Conformances) {
+      auto Ty = getTypeOffsets(Con);
+      Results.back().protocols.push_back({Ty.first, Ty.second});
+    }
 
     // Keep track of that we have a type reported for this range.
     AllPrintedTypes[Offset].insert(Length);
     return true;
   }
 };
 
+bool swift::resolveProtocolNames(DeclContext *DC,
+                                 ArrayRef<const char *> names,
+                          llvm::MapVector<ProtocolDecl*, StringRef> &result) {
+  assert(result.empty());
+  auto &ctx = DC->getASTContext();
+  for (auto name : names) {
+    // First try to solve by usr
+    ProtocolDecl *pd = dyn_cast_or_null<ProtocolDecl>(Demangle::
+      getTypeDeclForUSR(ctx, name));
+    if (!pd) {
+      // Second try to solve by mangled symbol name
+      pd = dyn_cast_or_null<ProtocolDecl>(Demangle::getTypeDeclForMangling(ctx, name));
+    }
+    if (!pd) {
+      // Thirdly try to solve by mangled type name
+      if (auto ty = Demangle::getTypeForMangling(ctx, name)) {
+        pd = dyn_cast_or_null<ProtocolDecl>(ty->getAnyGeneric());
+      }
+    }
+    if (pd) {
+      result.insert({pd, name});
+    }
+  }
+  if (names.size() == result.size())
+    return false;
+  // If we resolved none but the given names are not empty, return true for failure.
+  return result.size() == 0;
+}
+
 ArrayRef<ExpressionTypeInfo>
 swift::collectExpressionType(SourceFile &SF,
+                             ArrayRef<const char *> ExpectedProtocols,
                              std::vector<ExpressionTypeInfo> &Scratch,
                              llvm::raw_ostream &OS) {
-  ExpressionTypeCollector Walker(SF, Scratch, OS);
+  llvm::MapVector<ProtocolDecl*, StringRef> InterestedProtocols;
+  if (resolveProtocolNames(&SF, ExpectedProtocols, InterestedProtocols))
+    return {};
+  ExpressionTypeCollector Walker(SF, InterestedProtocols, Scratch, OS);
   Walker.walk(SF);
   return Scratch;
 }

test/IDE/Inputs/ExprTypeFiltered.swift

@@ -0,0 +1,31 @@
+protocol ProtEmpty {}
+
+protocol Prot {}
+
+protocol Prot1 {}
+
+class Clas: Prot {
+  var value: Clas { return self }
+  func getValue() -> Clas { return self }
+}
+
+struct Stru: Prot, Prot1 {
+  var value: Stru { return self }
+  func getValue() -> Stru { return self }
+}
+
+class C {}
+
+func ArrayC(_ a: [C]) {
+	_ = a.count
+	_ = a.description.count.advanced(by: 1).description
+	_ = a[0]
+}
+
+func ArrayClas(_ a: [Clas]) {
+	_ = a[0].value.getValue().value
+}
+
+func ArrayClas(_ a: [Stru]) {
+	_ = a[0].value.getValue().value
+}

test/IDE/expr_type_filtered.swift

@@ -0,0 +1,30 @@
+// RUN: %target-swift-ide-test -print-expr-type -source-filename %S/Inputs/ExprTypeFiltered.swift -swift-version 5 -module-name filtered -usr-filter 's:8filtered9ProtEmptyP' | %FileCheck %s -check-prefix=EMPTY
+// RUN: %target-swift-ide-test -print-expr-type -source-filename %S/Inputs/ExprTypeFiltered.swift -swift-version 5 -module-name filtered -usr-filter 's:8filtered4ProtP' | %FileCheck %s -check-prefix=PROTO
+// RUN: %target-swift-ide-test -print-expr-type -source-filename %S/Inputs/ExprTypeFiltered.swift -swift-version 5 -module-name filtered -usr-filter 's:8filtered5Prot1P' | %FileCheck %s -check-prefix=PROTO1
+
+// EMPTY: class Clas: Prot {
+// EMPTY:   var value: Clas { return self }
+// EMPTY:   func getValue() -> Clas { return self }
+// EMPTY: }
+// EMPTY: struct Stru: Prot, Prot1 {
+// EMPTY:   var value: Stru { return self }
+// EMPTY:   func getValue() -> Stru { return self }
+// EMPTY: }
+
+// PROTO: class Clas: Prot {
+// PROTO:   var value: Clas { return <expr type:"Clas">self</expr> }
+// PROTO:   func getValue() -> Clas { return <expr type:"Clas">self</expr> }
+// PROTO: }
+// PROTO: struct Stru: Prot, Prot1 {
+// PROTO:   var value: Stru { return <expr type:"Stru">self</expr> }
+// PROTO:   func getValue() -> Stru { return <expr type:"Stru">self</expr> }
+// PROTO: }
+
+// PROTO1: class Clas: Prot {
+// PROTO1:   var value: Clas { return self }
+// PROTO1:   func getValue() -> Clas { return self }
+// PROTO1: }
+// PROTO1: struct Stru: Prot, Prot1 {
+// PROTO1:   var value: Stru { return <expr type:"Stru">self</expr> }
+// PROTO1:   func getValue() -> Stru { return <expr type:"Stru">self</expr> }
+// PROTO1: }