Allow OperatorLookupDescriptor to hold a file or module

Author: hamishknight

include/swift/AST/NameLookupRequests.h

@@ -18,6 +18,7 @@
 
 #include "swift/AST/SimpleRequest.h"
 #include "swift/AST/ASTTypeIDs.h"
+#include "swift/AST/FileUnit.h"
 #include "swift/AST/Identifier.h"
 #include "swift/Basic/Statistic.h"
 #include "llvm/ADT/Hashing.h"
@@ -518,29 +519,53 @@ class DirectLookupRequest
 
 class OperatorLookupDescriptor final {
 public:
-  SourceFile *SF;
+  using Storage = llvm::PointerUnion<FileUnit *, ModuleDecl *>;
+  Storage fileOrModule;
   Identifier name;
   bool isCascading;
   SourceLoc diagLoc;
 
-  OperatorLookupDescriptor(SourceFile *SF, Identifier name, bool isCascading,
-                           SourceLoc diagLoc)
-      : SF(SF), name(name), isCascading(isCascading), diagLoc(diagLoc) {}
+private:
+  OperatorLookupDescriptor(Storage fileOrModule, Identifier name,
+                           bool isCascading, SourceLoc diagLoc)
+      : fileOrModule(fileOrModule), name(name), isCascading(isCascading),
+        diagLoc(diagLoc) {}
+
+public:
+  /// Retrieves the files to perform lookup in.
+  ArrayRef<FileUnit *> getFiles() const;
+
+  /// If this is for a module lookup, returns the module. Otherwise returns
+  /// \c nullptr.
+  ModuleDecl *getModule() const {
+    return fileOrModule.dyn_cast<ModuleDecl *>();
+  }
 
   friend llvm::hash_code hash_value(const OperatorLookupDescriptor &desc) {
-    return llvm::hash_combine(desc.SF, desc.name, desc.isCascading);
+    return llvm::hash_combine(desc.fileOrModule, desc.name, desc.isCascading);
   }
 
   friend bool operator==(const OperatorLookupDescriptor &lhs,
                          const OperatorLookupDescriptor &rhs) {
-    return lhs.SF == rhs.SF && lhs.name == rhs.name &&
+    return lhs.fileOrModule == rhs.fileOrModule && lhs.name == rhs.name &&
            lhs.isCascading == rhs.isCascading;
   }
 
   friend bool operator!=(const OperatorLookupDescriptor &lhs,
                          const OperatorLookupDescriptor &rhs) {
     return !(lhs == rhs);
   }
+
+  static OperatorLookupDescriptor forFile(FileUnit *file, Identifier name,
+                                          bool isCascading, SourceLoc diagLoc) {
+    return OperatorLookupDescriptor(file, name, isCascading, diagLoc);
+  }
+
+  static OperatorLookupDescriptor forModule(ModuleDecl *mod, Identifier name,
+                                            bool isCascading,
+                                            SourceLoc diagLoc) {
+    return OperatorLookupDescriptor(mod, name, isCascading, diagLoc);
+  }
 };
 
 void simple_display(llvm::raw_ostream &out,

lib/AST/Module.cpp

@@ -1154,22 +1154,20 @@ lookupOperatorDeclForName(ModuleDecl *M, SourceLoc Loc, Identifier Name) {
 template <typename OperatorType>
 llvm::Expected<OperatorType *> LookupOperatorRequest<OperatorType>::evaluate(
     Evaluator &evaluator, OperatorLookupDescriptor desc) const {
-  auto result = lookupOperatorDeclForName<OperatorType>(*desc.SF, desc.diagLoc,
+  auto *file = desc.fileOrModule.get<FileUnit *>();
+  auto result = lookupOperatorDeclForName<OperatorType>(*file, desc.diagLoc,
                                                         desc.name,
                                                         /*includePrivate*/ true);
   if (!result.hasValue())
     return nullptr;
-  if (auto *tracker = desc.SF->getReferencedNameTracker()) {
-    if (!result.getValue() ||
-        result.getValue()->getDeclContext()->getModuleScopeContext() !=
-            desc.SF) {
-      tracker->addTopLevelName(desc.name, desc.isCascading);
-    }
+
+  if (!result.getValue() ||
+      result.getValue()->getDeclContext()->getModuleScopeContext() != file) {
+    namelookup::recordLookupOfTopLevelName(file, desc.name, desc.isCascading);
   }
   if (!result.getValue()) {
-    result = lookupOperatorDeclForName<OperatorType>(desc.SF->getParentModule(),
-                                                     desc.diagLoc,
-                                                     desc.name);
+    result = lookupOperatorDeclForName<OperatorType>(file->getParentModule(),
+                                                     desc.diagLoc, desc.name);
   }
   return result.hasValue() ? result.getValue() : nullptr;
 }

lib/AST/NameLookupRequests.cpp

@@ -211,12 +211,20 @@ SourceLoc swift::extractNearestSourceLoc(const DirectLookupDescriptor &desc) {
 // LookupOperatorRequest computation.
 //----------------------------------------------------------------------------//
 
+ArrayRef<FileUnit *> OperatorLookupDescriptor::getFiles() const {
+  if (auto *module = getModule())
+    return module->getFiles();
+
+  // Return an ArrayRef pointing to the FileUnit in the union.
+  return llvm::makeArrayRef(*fileOrModule.getAddrOfPtr1());
+}
+
 void swift::simple_display(llvm::raw_ostream &out,
                            const OperatorLookupDescriptor &desc) {
   out << "looking up operator ";
   simple_display(out, desc.name);
   out << " in ";
-  simple_display(out, desc.SF);
+  simple_display(out, desc.fileOrModule);
 }
 
 SourceLoc swift::extractNearestSourceLoc(const OperatorLookupDescriptor &desc) {

lib/Sema/TypeCheckDecl.cpp

@@ -1253,12 +1253,9 @@ static PrecedenceGroupDecl *
 lookupPrecedenceGroup(const PrecedenceGroupDescriptor &descriptor) {
   auto *dc = descriptor.dc;
   if (auto sf = dc->getParentSourceFile()) {
-    OperatorLookupDescriptor desc{
-      sf,
-      descriptor.ident,
-      dc->isCascadingContextForLookup(false),
-      descriptor.nameLoc
-    };
+    auto desc = OperatorLookupDescriptor::forFile(
+        sf, descriptor.ident, dc->isCascadingContextForLookup(false),
+        descriptor.nameLoc);
     return evaluateOrDefault(sf->getASTContext().evaluator,
                              LookupPrecedenceGroupRequest{desc}, nullptr);
   } else {
@@ -1730,12 +1727,9 @@ FunctionOperatorRequest::evaluate(Evaluator &evaluator, FuncDecl *FD) const {
     FD->diagnose(diag::operator_in_local_scope);
   }
 
-  OperatorLookupDescriptor desc{
-    FD->getDeclContext()->getParentSourceFile(),
-    operatorName,
-    FD->isCascadingContextForLookup(false),
-    FD->getLoc()
-  };
+  auto desc = OperatorLookupDescriptor::forFile(
+      FD->getDeclContext()->getParentSourceFile(), operatorName,
+      FD->isCascadingContextForLookup(false), FD->getLoc());
   OperatorDecl *op = nullptr;
   if (FD->isUnaryOperator()) {
     if (FD->getAttrs().hasAttribute<PrefixAttr>()) {

lib/Sema/TypeCheckExpr.cpp

@@ -130,12 +130,9 @@ Expr *TypeChecker::substituteInputSugarTypeForResult(ApplyExpr *E) {
 static PrecedenceGroupDecl *lookupPrecedenceGroupForOperator(DeclContext *DC,
                                                              Identifier name,
                                                              SourceLoc loc) {
-  OperatorLookupDescriptor desc{
-    DC->getParentSourceFile(),
-    name,
-    DC->isCascadingContextForLookup(true),
-    loc
-  };
+  auto desc = OperatorLookupDescriptor::forFile(
+      DC->getParentSourceFile(), name, DC->isCascadingContextForLookup(true),
+      loc);
   auto &Ctx = DC->getASTContext();
   if (auto op = evaluateOrDefault(Ctx.evaluator,
                                   LookupInfixOperatorRequest{desc},