Requestify Operator Decl Attachment

Author: CodaFi

include/swift/AST/ASTTypeIDZone.def

@@ -27,6 +27,7 @@ SWIFT_TYPEID(CtorInitializerKind)
 SWIFT_TYPEID(ResilienceExpansion)
 SWIFT_TYPEID_NAMED(Optional<PropertyWrapperMutability>, PropertyWrapperMutability)
 SWIFT_TYPEID_NAMED(CustomAttr *, CustomAttr)
+SWIFT_TYPEID_NAMED(OperatorDecl *, OperatorDecl)
 SWIFT_TYPEID_NAMED(TypeAliasDecl *, TypeAliasDecl)
 SWIFT_TYPEID(AncestryFlags)
 SWIFT_TYPEID_NAMED(GenericSignature *, GenericSignature)

include/swift/AST/Decl.h

@@ -5926,7 +5926,6 @@ class FuncDecl : public AbstractFunctionDecl {
 
   TypeLoc FnRetType;
 
-  OperatorDecl *Operator = nullptr;
   OpaqueTypeDecl *OpaqueReturn = nullptr;
 
 protected:
@@ -6073,13 +6072,7 @@ class FuncDecl : public AbstractFunctionDecl {
     return cast_or_null<FuncDecl>(AbstractFunctionDecl::getOverriddenDecl());
   }
 
-  OperatorDecl *getOperatorDecl() const {
-    return Operator;
-  }
-  void setOperatorDecl(OperatorDecl *o) {
-    assert(isOperator() && "can't set an OperatorDecl for a non-operator");
-    Operator = o;
-  }
+  OperatorDecl *getOperatorDecl() const;
 
   OpaqueTypeDecl *getOpaqueResultTypeDecl() const {
     return OpaqueReturn;

include/swift/AST/TypeCheckRequests.h

@@ -1067,7 +1067,6 @@ class ClassAncestryFlagsRequest :
   // Evaluation.
   llvm::Expected<AncestryFlags>
   evaluate(Evaluator &evaluator, ClassDecl *value) const;
-
 public:
   // Caching.
   bool isCached() const { return true; }
@@ -1116,7 +1115,25 @@ class ExtendedTypeRequest
 
   // Evaluation.
   llvm::Expected<Type> evaluate(Evaluator &eval, ExtensionDecl *) const;
+public:
+  // Caching.
+  bool isCached() const { return true; }
+};
 
+class FunctionOperatorRequest :
+    public SimpleRequest<FunctionOperatorRequest,
+                         OperatorDecl *(FuncDecl *),
+                         CacheKind::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+  
+private:
+  friend SimpleRequest;
+  
+  // Evaluation.
+  llvm::Expected<OperatorDecl *>
+  evaluate(Evaluator &evaluator, FuncDecl *value) const;
+  
 public:
   // Caching.
   bool isCached() const { return true; }

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -59,3 +59,4 @@ SWIFT_TYPEID(IsImplicitlyUnwrappedOptionalRequest)
 SWIFT_TYPEID(ClassAncestryFlagsRequest)
 SWIFT_TYPEID(AbstractGenericSignatureRequest)
 SWIFT_TYPEID(ExtendedTypeRequest)
+SWIFT_TYPEID(FunctionOperatorRequest)

lib/AST/Decl.cpp

@@ -6887,6 +6887,18 @@ FuncDecl *FuncDecl::create(ASTContext &Context, SourceLoc StaticLoc,
   FD->getBodyResultTypeLoc() = FnRetType;
   return FD;
 }
+      
+OperatorDecl *FuncDecl::getOperatorDecl() const {
+  // Fast-path: Most functions are not operators.
+  if (!isOperator()) {
+    return nullptr;
+  }
+  return evaluateOrDefault(getASTContext().evaluator,
+                           FunctionOperatorRequest{
+                             const_cast<FuncDecl *>(this)
+                           },
+                           nullptr);
+ }
 
 AccessorDecl *AccessorDecl::createImpl(ASTContext &ctx,
                                        SourceLoc declLoc,

lib/Sema/TypeCheckDecl.cpp

@@ -3430,19 +3430,19 @@ static void validateTypealiasType(TypeChecker &tc, TypeAliasDecl *typeAlias) {
 }
 
 
-/// Bind the given function declaration, which declares an operator, to
-/// the corresponding operator declaration.
-void bindFuncDeclToOperator(TypeChecker &TC, FuncDecl *FD) {
-  OperatorDecl *op = nullptr;
+/// Bind the given function declaration, which declares an operator, to the corresponding operator declaration.
+llvm::Expected<OperatorDecl *>
+FunctionOperatorRequest::evaluate(Evaluator &evaluator, FuncDecl *FD) const {  
+  auto &C = FD->getASTContext();
+  auto &diags = C.Diags;
   auto operatorName = FD->getFullName().getBaseIdentifier();
 
   // Check for static/final/class when we're in a type.
   auto dc = FD->getDeclContext();
   if (dc->isTypeContext()) {
     if (!FD->isStatic()) {
-      TC.diagnose(FD->getLoc(), diag::nonstatic_operator_in_type,
-                  operatorName,
-                  dc->getDeclaredInterfaceType())
+      FD->diagnose(diag::nonstatic_operator_in_type,
+                   operatorName, dc->getDeclaredInterfaceType())
         .fixItInsert(FD->getAttributeInsertionLoc(/*forModifier=*/true),
                      "static ");
 
@@ -3451,17 +3451,18 @@ void bindFuncDeclToOperator(TypeChecker &TC, FuncDecl *FD) {
       // For a class, we also need the function or class to be 'final'.
       if (!classDecl->isFinal() && !FD->isFinal() &&
           FD->getStaticSpelling() != StaticSpellingKind::KeywordStatic) {
-        TC.diagnose(FD->getLoc(), diag::nonfinal_operator_in_class,
-                    operatorName, dc->getDeclaredInterfaceType())
+        FD->diagnose(diag::nonfinal_operator_in_class,
+                     operatorName, dc->getDeclaredInterfaceType())
           .fixItInsert(FD->getAttributeInsertionLoc(/*forModifier=*/true),
                        "final ");
-        FD->getAttrs().add(new (TC.Context) FinalAttr(/*IsImplicit=*/true));
+        FD->getAttrs().add(new (C) FinalAttr(/*IsImplicit=*/true));
       }
     }
   } else if (!dc->isModuleScopeContext()) {
-    TC.diagnose(FD, diag::operator_in_local_scope);
+    FD->diagnose(diag::operator_in_local_scope);
   }
 
+  OperatorDecl *op = nullptr;
   SourceFile &SF = *FD->getDeclContext()->getParentSourceFile();
   if (FD->isUnaryOperator()) {
     if (FD->getAttrs().hasAttribute<PrefixAttr>()) {
@@ -3485,21 +3486,21 @@ void bindFuncDeclToOperator(TypeChecker &TC, FuncDecl *FD) {
       // If we found both prefix and postfix, or neither prefix nor postfix,
       // complain. We can't fix this situation.
       if (static_cast<bool>(prefixOp) == static_cast<bool>(postfixOp)) {
-        TC.diagnose(FD, diag::declared_unary_op_without_attribute);
+        diags.diagnose(FD, diag::declared_unary_op_without_attribute);
 
         // If we found both, point at them.
         if (prefixOp) {
-          TC.diagnose(prefixOp, diag::unary_operator_declaration_here, false)
+          diags.diagnose(prefixOp, diag::unary_operator_declaration_here, false)
             .fixItInsert(FD->getLoc(), "prefix ");
-          TC.diagnose(postfixOp, diag::unary_operator_declaration_here, true)
+          diags.diagnose(postfixOp, diag::unary_operator_declaration_here, true)
             .fixItInsert(FD->getLoc(), "postfix ");
         } else {
           // FIXME: Introduce a Fix-It that adds the operator declaration?
         }
 
         // FIXME: Errors could cascade here, because name lookup for this
         // operator won't find this declaration.
-        return;
+        return nullptr;
       }
 
       // We found only one operator declaration, so we know whether this
@@ -3519,29 +3520,33 @@ void bindFuncDeclToOperator(TypeChecker &TC, FuncDecl *FD) {
       }
 
       // Emit diagnostic with the Fix-It.
-      TC.diagnose(FD->getFuncLoc(), diag::unary_op_missing_prepos_attribute,
+      diags.diagnose(FD->getFuncLoc(), diag::unary_op_missing_prepos_attribute,
                   static_cast<bool>(postfixOp))
         .fixItInsert(FD->getFuncLoc(), insertionText);
-      TC.diagnose(op, diag::unary_operator_declaration_here,
+      diags.diagnose(op, diag::unary_operator_declaration_here,
                   static_cast<bool>(postfixOp));
     }
   } else if (FD->isBinaryOperator()) {
     op = SF.lookupInfixOperator(operatorName,
                                 FD->isCascadingContextForLookup(false),
                                 FD->getLoc());
   } else {
-    TC.diagnose(FD, diag::invalid_arg_count_for_operator);
-    return;
+    diags.diagnose(FD, diag::invalid_arg_count_for_operator);
+    return nullptr;
   }
 
   if (!op) {
     SourceLoc insertionLoc;
     if (isa<SourceFile>(FD->getParent())) {
-      // Parent context is SourceFile, insertion location is start of func declaration
-      // or unary operator
-      insertionLoc = FD->isUnaryOperator() ? FD->getAttrs().getStartLoc() : FD->getStartLoc();
+      // Parent context is SourceFile, insertion location is start of func
+      // declaration or unary operator
+      if (FD->isUnaryOperator()) {
+        insertionLoc = FD->getAttrs().getStartLoc();
+      } else {
+        insertionLoc = FD->getStartLoc();
+      }
     } else {
-      // Finding top-level decl context before SourceFile and inserting before it
+      // Find the topmost non-file decl context and insert there.
       for (DeclContext *CurContext = FD->getLocalContext();
            !isa<SourceFile>(CurContext);
            CurContext = CurContext->getParent()) {
@@ -3552,25 +3557,28 @@ void bindFuncDeclToOperator(TypeChecker &TC, FuncDecl *FD) {
     }
 
     SmallString<128> insertion;
-    auto numOfParams = FD->getParameters()->size();
-    if (numOfParams == 1) {
-      if (FD->getAttrs().hasAttribute<PrefixAttr>())
-        insertion += "prefix operator ";
-      else
-        insertion += "postfix operator ";
-    } else if (numOfParams == 2) {
-      insertion += "infix operator ";
-    }
+    {
+      llvm::raw_svector_ostream str(insertion);
+      assert(FD->isUnaryOperator() || FD->isBinaryOperator());
+      if (FD->isUnaryOperator()) {
+        if (FD->getAttrs().hasAttribute<PrefixAttr>())
+          str << "prefix operator ";
+        else
+          str << "postfix operator ";
+      } else {
+        str << "infix operator ";
+      }
 
-    insertion += operatorName.str();
-    insertion += " : <# Precedence Group #>\n";
-    InFlightDiagnostic opDiagnostic = TC.diagnose(FD, diag::declared_operator_without_operator_decl);
+       str << operatorName.str() << " : <# Precedence Group #>\n";
+    }
+    InFlightDiagnostic opDiagnostic =
+        diags.diagnose(FD, diag::declared_operator_without_operator_decl);
     if (insertionLoc.isValid())
       opDiagnostic.fixItInsert(insertionLoc, insertion);
-    return;
+    return nullptr;
   }
 
-  FD->setOperatorDecl(op);
+  return op;
 }
 
 bool swift::isMemberOperator(FuncDecl *decl, Type type) {
@@ -3868,10 +3876,9 @@ void TypeChecker::validateDecl(ValueDecl *D) {
 
     DeclValidationRAII IBV(FD);
 
-    // Bind operator functions to the corresponding operator declaration.
-    if (FD->isOperator())
-      bindFuncDeclToOperator(*this, FD);
-
+    // Force computing the operator decl in case it emits diagnostics.
+    (void) FD->getOperatorDecl();
+    
     // Validate 'static'/'class' on functions in extensions.
     auto StaticSpelling = FD->getStaticSpelling();
     if (StaticSpelling != StaticSpellingKind::None &&

lib/Sema/TypeCheckDecl.h

@@ -42,4 +42,5 @@ void validatePatternBindingEntries(TypeChecker &tc,
                                    PatternBindingDecl *binding);
 }
 
-#endif
+#endif
+

lib/Serialization/Deserialization.cpp

@@ -3058,7 +3058,8 @@ class swift::DeclDeserializer {
     if (!isAccessor) {
       if (Decl *associated = MF.getDecl(associatedDeclID)) {
         if (auto op = dyn_cast<OperatorDecl>(associated)) {
-          fn->setOperatorDecl(op);
+          ctx.evaluator.cacheOutput(FunctionOperatorRequest{fn},
+                                    std::move(op));
 
           if (isa<PrefixOperatorDecl>(op))
             fn->getAttrs().add(new (ctx) PrefixAttr(/*implicit*/false));