[CSDiagnostics] Diagnose invalid partial application

lib/Sema/CSDiagnostics.cpp

@@ -1477,3 +1477,24 @@ bool MissingMemberFailure::diagnoseAsError() {
   corrections.noteAllCandidates();
   return true;
 }
+
+bool PartialApplicationFailure::diagnoseAsError() {
+  auto &cs = getConstraintSystem();
+  auto *anchor = cast<UnresolvedDotExpr>(getRawAnchor());
+
+  RefKind kind = RefKind::MutatingMethod;
+
+  // If this is initializer delegation chain, we have a tailored message.
+  if (getOverloadChoiceIfAvailable(cs.getConstraintLocator(
+          anchor, ConstraintLocator::ConstructorMember))) {
+    kind = anchor->getBase()->isSuperExpr() ? RefKind::SuperInit
+                                            : RefKind::SelfInit;
+  }
+
+  auto diagnostic = CompatibilityWarning
+                        ? diag::partial_application_of_function_invalid_swift4
+                        : diag::partial_application_of_function_invalid;
+
+  emitDiagnostic(anchor->getNameLoc(), diagnostic, kind);
+  return true;
+}

lib/Sema/CSDiagnostics.h

@@ -674,6 +674,23 @@ class MissingMemberFailure final : public FailureDiagnostic {
                                           DeclName memberName);
 };
 
+class PartialApplicationFailure final : public FailureDiagnostic {
+  enum RefKind : unsigned {
+    MutatingMethod,
+    SuperInit,
+    SelfInit,
+  };
+
+  bool CompatibilityWarning;
+
+public:
+  PartialApplicationFailure(Expr *root, bool warning, ConstraintSystem &cs,
+                            ConstraintLocator *locator)
+      : FailureDiagnostic(root, cs, locator), CompatibilityWarning(warning) {}
+
+  bool diagnoseAsError() override;
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSFix.cpp

@@ -258,3 +258,16 @@ DefineMemberBasedOnUse::create(ConstraintSystem &cs, Type baseType,
   return new (cs.getAllocator())
       DefineMemberBasedOnUse(cs, baseType, member, locator);
 }
+
+bool AllowInvalidPartialApplication::diagnose(Expr *root, bool asNote) const {
+  auto failure = PartialApplicationFailure(root, isWarning(),
+                                           getConstraintSystem(), getLocator());
+  return failure.diagnose(asNote);
+}
+
+AllowInvalidPartialApplication *
+AllowInvalidPartialApplication::create(bool isWarning, ConstraintSystem &cs,
+                                       ConstraintLocator *locator) {
+  return new (cs.getAllocator())
+      AllowInvalidPartialApplication(isWarning, cs, locator);
+}

lib/Sema/CSFix.h

@@ -104,6 +104,14 @@ enum class FixKind : uint8_t {
   /// fix this issue by pretending that member exists and matches
   /// given arguments/result types exactly.
   DefineMemberBasedOnUse,
+
+  /// Allow expressions where 'mutating' method is only partially applied,
+  /// which means either not applied at all e.g. `Foo.bar` or only `Self`
+  /// is applied e.g. `foo.bar` or `Foo.bar(&foo)`.
+  ///
+  /// Allow expressions where initializer call (either `self.init` or
+  /// `super.init`) is only partially applied.
+  AllowInvalidPartialApplication,
 };
 
 class ConstraintFix {
@@ -501,6 +509,24 @@ class DefineMemberBasedOnUse final : public ConstraintFix {
                                         ConstraintLocator *locator);
 };
 
+class AllowInvalidPartialApplication final : public ConstraintFix {
+public:
+  AllowInvalidPartialApplication(bool isWarning, ConstraintSystem &cs,
+                                 ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::AllowInvalidPartialApplication, locator,
+                      isWarning) {}
+
+  std::string getName() const override {
+    return "allow partially applied 'mutating' method";
+  }
+
+  bool diagnose(Expr *root, bool asNote = false) const override;
+
+  static AllowInvalidPartialApplication *create(bool isWarning,
+                                                ConstraintSystem &cs,
+                                                ConstraintLocator *locator);
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSGen.cpp

@@ -1633,7 +1633,7 @@ namespace {
                                      expr->getFunctionRefKind(),
                                      expr->getOuterAlternatives());
     }
-    
+
     Type visitUnresolvedSpecializeExpr(UnresolvedSpecializeExpr *expr) {
       auto baseTy = CS.getType(expr->getSubExpr());
 

lib/Sema/CSSimplify.cpp

@@ -5502,6 +5502,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::AutoClosureForwarding:
   case FixKind::RemoveUnwrap:
   case FixKind::DefineMemberBasedOnUse:
+  case FixKind::AllowInvalidPartialApplication:
     llvm_unreachable("handled elsewhere");
   }
 

lib/Sema/ConstraintSystem.cpp

@@ -1632,6 +1632,77 @@ resolveOverloadForDeclWithSpecialTypeCheckingSemantics(ConstraintSystem &CS,
   llvm_unreachable("Unhandled DeclTypeCheckingSemantics in switch.");
 }
 
+/// \returns true if given declaration is an instance method marked as
+/// `mutating`, false otherwise.
+bool isMutatingMethod(const ValueDecl *decl) {
+  if (!(decl->isInstanceMember() && isa<FuncDecl>(decl)))
+    return false;
+  return cast<FuncDecl>(decl)->isMutating();
+}
+
+static bool shouldCheckForPartialApplication(ConstraintSystem &cs,
+                                             const ValueDecl *decl,
+                                             ConstraintLocator *locator) {
+  auto *anchor = locator->getAnchor();
+  if (!(anchor && isa<UnresolvedDotExpr>(anchor)))
+    return false;
+
+  // FIXME(diagnostics): This check should be removed together with
+  // expression based diagnostics.
+  if (cs.TC.isExprBeingDiagnosed(anchor))
+    return false;
+
+  // If this is a reference to instance method marked as 'mutating'
+  // it should be checked for invalid partial application.
+  if (isMutatingMethod(decl))
+    return true;
+
+  // Another unsupported partial application is related
+  // to constructor delegation via `self.init` or `super.init`.
+
+  if (!isa<ConstructorDecl>(decl))
+    return false;
+
+  auto *UDE = cast<UnresolvedDotExpr>(anchor);
+  // This is `super.init`
+  if (UDE->getBase()->isSuperExpr())
+    return true;
+
+  // Or this might be `self.init`.
+  if (auto *DRE = dyn_cast<DeclRefExpr>(UDE->getBase())) {
+    if (auto *baseDecl = DRE->getDecl())
+      return baseDecl->getBaseName() == cs.getASTContext().Id_self;
+  }
+
+  return false;
+}
+
+/// Try to identify and fix failures related to partial function application
+/// e.g. partial application of `init` or 'mutating' instance methods.
+static std::pair<bool, unsigned>
+isInvalidPartialApplication(ConstraintSystem &cs, const ValueDecl *member,
+                            ConstraintLocator *locator) {
+  if (!shouldCheckForPartialApplication(cs, member, locator))
+    return {false, 0};
+
+  auto anchor = cast<UnresolvedDotExpr>(locator->getAnchor());
+  // If this choice is a partial application of `init` or
+  // `mutating` instance method we should report that it's not allowed.
+  auto baseTy =
+      cs.simplifyType(cs.getType(anchor->getBase()))->getWithoutSpecifierType();
+
+  // If base is a metatype it would be ignored (unless this is an initializer
+  // call), but if it is some other type it means that we have a single
+  // application level already.
+  unsigned level =
+      baseTy->is<MetatypeType>() && !isa<ConstructorDecl>(member) ? 0 : 1;
+  if (auto *call = dyn_cast_or_null<CallExpr>(cs.getParentExpr(anchor))) {
+    level += dyn_cast_or_null<CallExpr>(cs.getParentExpr(call)) ? 2 : 1;
+  }
+
+  return {true, level};
+}
+
 void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
                                        Type boundType,
                                        OverloadChoice choice,
@@ -1872,10 +1943,10 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
   }
   assert(!refType->hasTypeParameter() && "Cannot have a dependent type here");
 
-  // If we're binding to an init member, the 'throws' need to line up between
-  // the bound and reference types.
   if (choice.isDecl()) {
     auto decl = choice.getDecl();
+    // If we're binding to an init member, the 'throws' need to line up between
+    // the bound and reference types.
     if (auto CD = dyn_cast<ConstructorDecl>(decl)) {
       auto boundFunctionType = boundType->getAs<AnyFunctionType>();
 
@@ -1885,6 +1956,41 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
             boundFunctionType->getExtInfo().withThrows());
       }
     }
+
+    // Check whether applying this overload would result in invalid
+    // partial function application e.g. partial application of
+    // mutating method or initializer.
+
+    // This check is supposed to be performed without
+    // `shouldAttemptFixes` because name lookup can't
+    // detect that particular partial application is
+    // invalid, so it has to return all of the candidates.
+
+    bool isInvalidPartialApply;
+    unsigned level;
+
+    std::tie(isInvalidPartialApply, level) =
+        isInvalidPartialApplication(*this, decl, locator);
+
+    if (isInvalidPartialApply) {
+      // No application at all e.g. `Foo.bar`.
+      if (level == 0) {
+        // Swift 4 and earlier failed to diagnose a reference to a mutating
+        // method without any applications at all, which would get
+        // miscompiled into a function with undefined behavior. Warn for
+        // source compatibility.
+        bool isWarning = !getASTContext().isSwiftVersionAtLeast(5);
+        (void)recordFix(
+            AllowInvalidPartialApplication::create(isWarning, *this, locator));
+      } else if (level == 1) {
+        // `Self` parameter is applied, e.g. `foo.bar` or `Foo.bar(&foo)`
+        (void)recordFix(AllowInvalidPartialApplication::create(
+            /*isWarning=*/false, *this, locator));
+      }
+
+      // Otherwise both `Self` and arguments are applied,
+      // e.g. `foo.bar()` or `Foo.bar(&foo)()`, and there is nothing to do.
+    }
   }
 
   // Note that we have resolved this overload.

lib/Sema/ConstraintSystem.h

@@ -1775,6 +1775,12 @@ class ConstraintSystem {
   ConstraintLocator *
   getConstraintLocator(const ConstraintLocatorBuilder &builder);
 
+  /// Lookup and return parent associated with given expression.
+  Expr *getParentExpr(Expr *expr) const {
+    auto e = ExprWeights.find(expr);
+    return e != ExprWeights.end() ? e->second.second : nullptr;
+  }
+
 public:
 
   /// Whether we should attempt to fix problems.

lib/Sema/MiscDiagnostics.cpp

@@ -100,92 +100,6 @@ static void diagSyntacticUseRestrictions(TypeChecker &TC, const Expr *E,
       unsigned level : 29;
     };
 
-    // Partial applications of functions that are not permitted.  This is
-    // tracked in post-order and unraveled as subsequent applications complete
-    // the call (or not).
-    llvm::SmallDenseMap<Expr*, PartialApplication,2> InvalidPartialApplications;
-
-    ~DiagnoseWalker() override {
-      for (auto &unapplied : InvalidPartialApplications) {
-        unsigned kind = unapplied.second.kind;
-        if (unapplied.second.compatibilityWarning) {
-          TC.diagnose(unapplied.first->getLoc(),
-                      diag::partial_application_of_function_invalid_swift4,
-                      kind);
-        } else {
-          TC.diagnose(unapplied.first->getLoc(),
-                      diag::partial_application_of_function_invalid,
-                      kind);
-        }
-      }
-    }
-
-    /// methods are fully applied when they can't support partial application.
-    void checkInvalidPartialApplication(Expr *E) {
-      if (auto AE = dyn_cast<ApplyExpr>(E)) {
-        Expr *fnExpr = AE->getSemanticFn();
-        if (auto forceExpr = dyn_cast<ForceValueExpr>(fnExpr))
-          fnExpr = forceExpr->getSubExpr()->getSemanticsProvidingExpr();
-        if (auto dotSyntaxExpr = dyn_cast<DotSyntaxBaseIgnoredExpr>(fnExpr))
-          fnExpr = dotSyntaxExpr->getRHS();
-
-        // Check to see if this is a potentially unsupported partial
-        // application of a constructor delegation.
-        if (isa<OtherConstructorDeclRefExpr>(fnExpr)) {
-          auto kind = AE->getArg()->isSuperExpr()
-                    ? PartialApplication::SuperInit
-                    : PartialApplication::SelfInit;
-
-          // Partial applications of delegated initializers aren't allowed, and
-          // don't really make sense to begin with.
-          InvalidPartialApplications.insert(
-            {E, {PartialApplication::Error, kind, 1}});
-          return;
-        }
-
-        // If this is adding a level to an active partial application, advance
-        // it to the next level.
-        auto foundApplication = InvalidPartialApplications.find(fnExpr);
-        if (foundApplication == InvalidPartialApplications.end())
-          return;
-
-        unsigned level = foundApplication->second.level;
-        auto kind = foundApplication->second.kind;
-        assert(level > 0);
-        InvalidPartialApplications.erase(foundApplication);
-        if (level > 1) {
-          // We have remaining argument clauses.
-          // Partial applications were always diagnosed in Swift 4 and before,
-          // so there's no need to preserve the compatibility warning bit.
-          InvalidPartialApplications.insert(
-            {AE, {PartialApplication::Error, kind, level - 1}});
-        }
-        return;
-      }
-
-      /// If this is a reference to a mutating method, it cannot be partially
-      /// applied or even referenced without full application, so arrange for
-      /// us to check that it gets fully applied.
-      auto fnDeclRef = dyn_cast<DeclRefExpr>(E);
-      if (!fnDeclRef)
-        return;
-
-      auto fn = dyn_cast<FuncDecl>(fnDeclRef->getDecl());
-      if (!fn || !fn->isInstanceMember() || !fn->isMutating())
-        return;
-
-      // Swift 4 and earlier failed to diagnose a reference to a mutating method
-      // without any applications at all, which would get miscompiled into a
-      // function with undefined behavior. Warn for source compatibility.
-      auto errorBehavior = TC.Context.LangOpts.isSwiftVersionAtLeast(5)
-        ? PartialApplication::Error
-        : PartialApplication::CompatibilityWarning;
-
-      InvalidPartialApplications.insert(
-        {fnDeclRef, {errorBehavior,
-                     PartialApplication::MutatingMethod, 2}});
-    }
-
     // Not interested in going outside a basic expression.
     std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
       return { false, S };
@@ -463,11 +377,6 @@ static void diagSyntacticUseRestrictions(TypeChecker &TC, const Expr *E,
       return arg;
     }
 
-    Expr *walkToExprPost(Expr *E) override {
-      checkInvalidPartialApplication(E);
-      return E;
-    }
-
     void checkConvertedPointerArgument(ConcreteDeclRef callee,
                                        unsigned uncurryLevel,
                                        unsigned argIndex,

test/Constraints/mutating_members.swift

@@ -1,5 +1,7 @@
 // RUN: %target-swift-frontend -typecheck -verify -swift-version 5 %s
 
+protocol P {}
+
 struct Foo {
   mutating func boom() {}
 }
@@ -20,3 +22,15 @@ func fromLocalContext2(x: inout Foo, y: Bool) -> () -> () {
   }
 }
 
+func bar() -> P.Type { fatalError() }
+func bar() -> Foo.Type { fatalError() }
+
+_ = bar().boom       // expected-error{{partial application of 'mutating' method}}
+_ = bar().boom(&y)   // expected-error{{partial application of 'mutating' method}}
+_ = bar().boom(&y)() // ok
+
+func foo(_ foo: Foo.Type) {
+  _ = foo.boom       // expected-error{{partial application of 'mutating' method}}
+  _ = foo.boom(&y)   // expected-error{{partial application of 'mutating' method}}
+  _ = foo.boom(&y)() // ok
+}