[AST] Add ValueDecl::getNumCurryLevels

Author: hamishknight

include/swift/AST/Decl.h

@@ -2704,6 +2704,9 @@ class ValueDecl : public Decl {
   /// parameter lists, for example an enum element without associated values.
   bool hasParameterList() const;
 
+  /// Returns the number of curry levels in the declaration's interface type.
+  unsigned getNumCurryLevels() const;
+
   /// Get the decl for this value's opaque result type, if it has one.
   OpaqueTypeDecl *getOpaqueResultTypeDecl() const;
 
@@ -7164,6 +7167,15 @@ inline bool ValueDecl::hasParameterList() const {
   return isa<AbstractFunctionDecl>(this) || isa<SubscriptDecl>(this);
 }
 
+inline unsigned ValueDecl::getNumCurryLevels() const {
+  unsigned curryLevels = 0;
+  if (hasParameterList())
+    curryLevels++;
+  if (hasCurriedSelf())
+    curryLevels++;
+  return curryLevels;
+}
+
 inline bool Decl::isPotentiallyOverridable() const {
   if (isa<VarDecl>(this) ||
       isa<SubscriptDecl>(this) ||

lib/AST/Decl.cpp

@@ -2401,7 +2401,7 @@ CanType ValueDecl::getOverloadSignatureType() const {
                            /*topLevelFunction=*/true,
                            isMethod,
                            /*isInitializer=*/isa<ConstructorDecl>(afd),
-                           isMethod ? 2 : 1)->getCanonicalType();
+                           getNumCurryLevels())->getCanonicalType();
   }
 
   if (isa<AbstractStorageDecl>(this)) {
@@ -2417,7 +2417,7 @@ CanType ValueDecl::getOverloadSignatureType() const {
           /*topLevelFunction=*/true,
           /*isMethod=*/false,
           /*isInitializer=*/false,
-          1)->getCanonicalType();
+          getNumCurryLevels())->getCanonicalType();
     }
 
     // We want to curry the default signature type with the 'self' type of the
@@ -2431,7 +2431,7 @@ CanType ValueDecl::getOverloadSignatureType() const {
   if (isa<EnumElementDecl>(this)) {
     auto mappedType = mapSignatureFunctionType(
         getASTContext(), getInterfaceType(), /*topLevelFunction=*/false,
-        /*isMethod=*/false, /*isInitializer=*/false, /*curryLevels=*/0);
+        /*isMethod=*/false, /*isInitializer=*/false, getNumCurryLevels());
     return mappedType->getCanonicalType();
   }
 

lib/ParseSIL/ParseSIL.cpp

@@ -2028,18 +2028,9 @@ bool SILParser::parseSILDeclRef(SILDeclRef &Member, bool FnTypeRequired) {
     for (unsigned I = 0, E = values.size(); I < E; I++) {
       auto *decl = values[I];
 
-      unsigned numArgumentLabels = 0;
-      if (auto *eed = dyn_cast<EnumElementDecl>(decl)) {
-        numArgumentLabels =
-          (eed->hasAssociatedValues() ? 2 : 1);
-      } else if (auto *afd = dyn_cast<AbstractFunctionDecl>(decl)) {
-        numArgumentLabels =
-          (decl->getDeclContext()->isTypeContext() ? 2 : 1);
-      }
-
       auto lookupTy =
         decl->getInterfaceType()
-            ->removeArgumentLabels(numArgumentLabels);
+            ->removeArgumentLabels(decl->getNumCurryLevels());
       if (declTy == lookupTy->getCanonicalType()) {
         TheDecl = decl;
         // Update SILDeclRef to point to the right Decl.

lib/SIL/SILDeclRef.cpp

@@ -1016,10 +1016,10 @@ unsigned SILDeclRef::getParameterListCount() const {
 
   auto *vd = getDecl();
 
-  if (auto *func = dyn_cast<AbstractFunctionDecl>(vd)) {
-    return func->hasImplicitSelfDecl() ? 2 : 1;
-  } else if (auto *ed = dyn_cast<EnumElementDecl>(vd)) {
-    return ed->hasAssociatedValues() ? 2 : 1;
+  if (isa<AbstractFunctionDecl>(vd) || isa<EnumElementDecl>(vd)) {
+    // For functions and enum elements, the number of parameter lists is the
+    // same as in their interface type.
+    return vd->getNumCurryLevels();
   } else if (isa<ClassDecl>(vd)) {
     return 2;
   } else if (isa<VarDecl>(vd)) {

lib/Sema/ConstraintSystem.cpp

@@ -877,29 +877,7 @@ void ConstraintSystem::recordOpenedTypes(
 static unsigned getNumRemovedArgumentLabels(TypeChecker &TC, ValueDecl *decl,
                                             bool isCurriedInstanceReference,
                                             FunctionRefKind functionRefKind) {
-  unsigned numParameterLists;
-
-  // Enum elements with associated values have to be treated
-  // as regular function values.
-  //
-  // enum E {
-  //   case foo(a: Int)
-  // }
-  // let bar: [Int] = []
-  // bar.map(E.foo)
-  //
-  // `E.foo` has to act as a regular function type passed as a value.
-  if (auto *EED = dyn_cast<EnumElementDecl>(decl)) {
-    numParameterLists = EED->hasAssociatedValues() ? 2 : 1;
-
-  // Only applicable to functions. Nothing else should have argument labels in
-  // the type.
-  } else if (auto func = dyn_cast<AbstractFunctionDecl>(decl)) {
-    numParameterLists = func->hasImplicitSelfDecl() ? 2 : 1;
-  } else {
-    return 0;
-  }
-
+  unsigned numParameterLists = decl->getNumCurryLevels();
   switch (functionRefKind) {
   case FunctionRefKind::Unapplied:
   case FunctionRefKind::Compound:

lib/Sema/TypeCheckError.cpp

@@ -55,7 +55,7 @@ class AbstractFunction {
   explicit AbstractFunction(AbstractFunctionDecl *fn)
     : TheKind(Kind::Function),
       IsRethrows(fn->getAttrs().hasAttribute<RethrowsAttr>()),
-      ParamCount(fn->hasImplicitSelfDecl() ? 2 : 1) {
+      ParamCount(fn->getNumCurryLevels()) {
     TheFunction = fn;
   }
 
@@ -896,7 +896,7 @@ class Context {
     }
 
     return Context(getKindForFunctionBody(
-        D->getInterfaceType(), D->hasImplicitSelfDecl() ? 2 : 1));
+        D->getInterfaceType(), D->getNumCurryLevels()));
   }
 
   static Context forDeferBody() {