Add getOperatorDecls to FileUnit and ModuleDecl

Query the SourceLookupCache for the operator decls,
and use ModuleDecl::getOperatorDecls for both
frontend stats and to clean up some code
completion logic.

In addition, this commit switches getPrecedenceGroups
over to querying SourceLookupCache.

Author: hamishknight

include/swift/AST/FileUnit.h

@@ -196,6 +196,13 @@ class FileUnit : public DeclContext {
               SmallVectorImpl<Decl*> &Results,
               llvm::function_ref<bool(DeclAttributes)> matchAttributes) const;
 
+  /// Finds all operator decls in this file.
+  ///
+  /// This does a simple local lookup, not recursively looking through imports.
+  /// The order of the results is not guaranteed to be meaningful.
+  virtual void
+  getOperatorDecls(SmallVectorImpl<OperatorDecl *> &results) const {}
+
   /// Finds all precedence group decls in this file.
   ///
   /// This does a simple local lookup, not recursively looking through imports.

include/swift/AST/Module.h

@@ -620,6 +620,12 @@ class ModuleDecl : public DeclContext, public TypeDecl {
   /// The order of the results is not guaranteed to be meaningful.
   void getLocalTypeDecls(SmallVectorImpl<TypeDecl*> &Results) const;
 
+  /// Finds all operator decls of this module.
+  ///
+  /// This does a simple local lookup, not recursively looking through imports.
+  /// The order of the results is not guaranteed to be meaningful.
+  void getOperatorDecls(SmallVectorImpl<OperatorDecl *> &results) const;
+
   /// Finds all precedence group decls of this module.
   ///
   /// This does a simple local lookup, not recursively looking through imports.

include/swift/AST/SourceFile.h

@@ -448,6 +448,9 @@ class SourceFile final : public FileUnit {
 public:
   virtual void getTopLevelDecls(SmallVectorImpl<Decl*> &results) const override;
 
+  virtual void
+  getOperatorDecls(SmallVectorImpl<OperatorDecl *> &results) const override;
+
   virtual void
   getPrecedenceGroups(SmallVectorImpl<PrecedenceGroupDecl*> &results) const override;
 

include/swift/Basic/Statistics.def

@@ -130,10 +130,8 @@ FRONTEND_STATISTIC(AST, NumLocalTypeDecls)
 /// Number of Objective-C declarations in the AST context.
 FRONTEND_STATISTIC(AST, NumObjCMethods)
 
-/// Number of infix, postfix, and prefix operators in the AST context.
-FRONTEND_STATISTIC(AST, NumInfixOperators)
-FRONTEND_STATISTIC(AST, NumPostfixOperators)
-FRONTEND_STATISTIC(AST, NumPrefixOperators)
+/// Number of operators in the AST context.
+FRONTEND_STATISTIC(AST, NumOperators)
 
 /// Number of precedence groups in the AST context.
 FRONTEND_STATISTIC(AST, NumPrecedenceGroups)

include/swift/Serialization/SerializedModuleLoader.h

@@ -384,6 +384,9 @@ class SerializedASTFile final : public LoadedFile {
       SmallVectorImpl<Decl*> &Results,
       llvm::function_ref<bool(DeclAttributes)> matchAttributes) const override;
 
+  virtual void
+  getOperatorDecls(SmallVectorImpl<OperatorDecl *> &results) const override;
+
   virtual void
   getPrecedenceGroups(SmallVectorImpl<PrecedenceGroupDecl*> &Results) const override;
 

lib/AST/Module.cpp

@@ -745,18 +745,36 @@ void SourceFile::getTopLevelDecls(SmallVectorImpl<Decl*> &Results) const {
   Results.append(decls.begin(), decls.end());
 }
 
+void ModuleDecl::getOperatorDecls(
+    SmallVectorImpl<OperatorDecl *> &results) const {
+  // For a parsed module, we can check the source cache on the module rather
+  // than doing an O(N) search over the source files.
+  if (isParsedModule(this)) {
+    getSourceLookupCache().getOperatorDecls(results);
+    return;
+  }
+  FORWARD(getOperatorDecls, (results));
+}
+
+void SourceFile::getOperatorDecls(
+       SmallVectorImpl<OperatorDecl*> &results) const {
+  getCache().getOperatorDecls(results);
+}
+
 void ModuleDecl::getPrecedenceGroups(
-       SmallVectorImpl<PrecedenceGroupDecl*> &Results) const {
-  FORWARD(getPrecedenceGroups, (Results));
+       SmallVectorImpl<PrecedenceGroupDecl*> &results) const {
+  // For a parsed module, we can check the source cache on the module rather
+  // than doing an O(N) search over the source files.
+  if (isParsedModule(this)) {
+    getSourceLookupCache().getPrecedenceGroups(results);
+    return;
+  }
+  FORWARD(getPrecedenceGroups, (results));
 }
 
 void SourceFile::getPrecedenceGroups(
-       SmallVectorImpl<PrecedenceGroupDecl*> &Results) const {
-  for (auto pair : PrecedenceGroups) {
-    if (pair.second.getPointer() && pair.second.getInt()) {
-      Results.push_back(pair.second.getPointer());
-    }
-  }
+       SmallVectorImpl<PrecedenceGroupDecl*> &results) const {
+  getCache().getPrecedenceGroups(results);
 }
 
 void SourceFile::getLocalTypeDecls(SmallVectorImpl<TypeDecl*> &Results) const {

lib/FrontendTool/FrontendTool.cpp

@@ -655,10 +655,14 @@ static void countStatsOfSourceFile(UnifiedStatsReporter &Stats,
   C.NumDecls += SF->getTopLevelDecls().size();
   C.NumLocalTypeDecls += SF->LocalTypeDecls.size();
   C.NumObjCMethods += SF->ObjCMethods.size();
-  C.NumInfixOperators += SF->InfixOperators.size();
-  C.NumPostfixOperators += SF->PostfixOperators.size();
-  C.NumPrefixOperators += SF->PrefixOperators.size();
-  C.NumPrecedenceGroups += SF->PrecedenceGroups.size();
+
+  SmallVector<OperatorDecl *, 2> operators;
+  SF->getOperatorDecls(operators);
+  C.NumOperators += operators.size();
+
+  SmallVector<PrecedenceGroupDecl *, 2> groups;
+  SF->getPrecedenceGroups(groups);
+  C.NumPrecedenceGroups += groups.size();
 
   auto bufID = SF->getBufferID();
   if (bufID.hasValue()) {

lib/IDE/CodeCompletion.cpp

@@ -3421,56 +3421,10 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
                              IncludeInstanceMembers);
   }
 
-  template <typename T>
-  void collectOperatorsFromMap(SourceFile::OperatorMap<T> &map,
-                               bool includePrivate,
-                               std::vector<OperatorDecl *> &results) {
-    for (auto &pair : map) {
-      if (pair.second.getPointer() &&
-          (pair.second.getInt() || includePrivate)) {
-        results.push_back(pair.second.getPointer());
-      }
-    }
-  }
-
-  void collectOperatorsFrom(SourceFile *SF,
-                            std::vector<OperatorDecl *> &results) {
-    bool includePrivate = CurrDeclContext->getParentSourceFile() == SF;
-    collectOperatorsFromMap(SF->PrefixOperators, includePrivate, results);
-    collectOperatorsFromMap(SF->PostfixOperators, includePrivate, results);
-    collectOperatorsFromMap(SF->InfixOperators, includePrivate, results);
-  }
-
-  void collectOperatorsFrom(LoadedFile *F,
-                            std::vector<OperatorDecl *> &results) {
-    SmallVector<Decl *, 64> topLevelDecls;
-    F->getTopLevelDecls(topLevelDecls);
-    for (auto D : topLevelDecls) {
-      if (auto op = dyn_cast<OperatorDecl>(D))
-        results.push_back(op);
-    }
-  }
-
-  std::vector<OperatorDecl *> collectOperators() {
-    std::vector<OperatorDecl *> results;
+  void collectOperators(SmallVectorImpl<OperatorDecl *> &results) {
     assert(CurrDeclContext);
-    for (auto import : namelookup::getAllImports(CurrDeclContext)) {
-      for (auto fileUnit : import.second->getFiles()) {
-        switch (fileUnit->getKind()) {
-        case FileUnitKind::Builtin:
-        case FileUnitKind::ClangModule:
-        case FileUnitKind::DWARFModule:
-          continue;
-        case FileUnitKind::Source:
-          collectOperatorsFrom(cast<SourceFile>(fileUnit), results);
-          break;
-        case FileUnitKind::SerializedAST:
-          collectOperatorsFrom(cast<LoadedFile>(fileUnit), results);
-          break;
-        }
-      }
-    }
-    return results;
+    for (auto import : namelookup::getAllImports(CurrDeclContext))
+      import.second->getOperatorDecls(results);
   }
 
   void addPostfixBang(Type resultType) {
@@ -3618,7 +3572,8 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
 
     Expr *foldedExpr = typeCheckLeadingSequence(LHS, leadingSequence);
 
-    std::vector<OperatorDecl *> operators = collectOperators();
+    SmallVector<OperatorDecl *, 16> operators;
+    collectOperators(operators);
     // FIXME: this always chooses the first operator with the given name.
     llvm::DenseSet<Identifier> seenPostfixOperators;
     llvm::DenseSet<Identifier> seenInfixOperators;

lib/Serialization/ModuleFile.cpp

@@ -2694,6 +2694,17 @@ void ModuleFile::getTopLevelDecls(
   }
 }
 
+void ModuleFile::getOperatorDecls(SmallVectorImpl<OperatorDecl *> &results) {
+  PrettyStackTraceModuleFile stackEntry(*this);
+  if (!OperatorDecls)
+    return;
+
+  for (auto entry : OperatorDecls->data()) {
+    for (auto item : entry)
+      results.push_back(cast<OperatorDecl>(getDecl(item.second)));
+  }
+}
+
 void ModuleFile::getPrecedenceGroups(
        SmallVectorImpl<PrecedenceGroupDecl*> &results) {
   PrettyStackTraceModuleFile stackEntry(*this);

lib/Serialization/ModuleFile.h

@@ -829,6 +829,9 @@ class ModuleFile
          SmallVectorImpl<Decl*> &Results,
          llvm::function_ref<bool(DeclAttributes)> matchAttributes = nullptr);
 
+  /// Adds all operators to the given vector.
+  void getOperatorDecls(SmallVectorImpl<OperatorDecl *> &Results);
+
   /// Adds all precedence groups to the given vector.
   void getPrecedenceGroups(SmallVectorImpl<PrecedenceGroupDecl*> &Results);