Merge pull request swiftlang#23255 from nkcsgexi/proco-filter

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

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>
 
@@ -180,12 +181,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: }