Merge pull request #70707 from rjmccall/variadic-generic-class-method-dispatch

Fix variadic-generic class method dispatch

Author: rjmccall

include/swift/AST/ConcreteDeclRef.h

@@ -68,6 +68,9 @@ class ConcreteDeclRef {
   /// Retrieve a reference to the declaration this one overrides.
   ConcreteDeclRef getOverriddenDecl() const;
 
+  /// Retrieve a reference to the given declaration that this one overrides.
+  ConcreteDeclRef getOverriddenDecl(ValueDecl *overriddenDecl) const;
+
   /// Determine whether this reference specializes the declaration to which
   /// it refers.
   bool isSpecialized() const { return !substitutions.empty(); }

lib/AST/ConcreteDeclRef.cpp

@@ -43,6 +43,29 @@ ConcreteDeclRef ConcreteDeclRef::getOverriddenDecl() const {
   return ConcreteDeclRef(baseDecl, subs);
 }
 
+ConcreteDeclRef ConcreteDeclRef::getOverriddenDecl(ValueDecl *baseDecl) const {
+  auto *derivedDecl = getDecl();
+  if (baseDecl == derivedDecl) return *this;
+
+#ifndef NDEBUG
+  {
+    auto cur = derivedDecl;
+    for (; cur && cur != baseDecl; cur = cur->getOverriddenDecl()) {}
+    assert(cur && "decl is not an indirect override of baseDecl");
+  }
+#endif
+
+  if (!baseDecl->getInnermostDeclContext()->isGenericContext()) {
+    return ConcreteDeclRef(baseDecl);
+  }
+
+  auto subs = SubstitutionMap::getOverrideSubstitutions(baseDecl, derivedDecl);
+  if (auto derivedSubs = getSubstitutions()) {
+    subs = subs.subst(derivedSubs);
+  }
+  return ConcreteDeclRef(baseDecl, subs);
+}
+
 void ConcreteDeclRef::dump(raw_ostream &os) const {
   if (!getDecl()) {
     os << "**NULL**";

lib/SILGen/SILGenApply.cpp

@@ -444,17 +444,34 @@ class Callee {
     auto base = c.getOverriddenVTableEntry();
     auto &baseCI = SGF.getConstantInfo(SGF.getTypeExpansionContext(), base);
     auto &derivedCI = SGF.getConstantInfo(SGF.getTypeExpansionContext(), c);
+    subs = subs.mapIntoTypeExpansionContext(SGF.getTypeExpansionContext());
+
+    // We use an orig function type based on the overridden vtable entry, but
+    // the substitutions we have are for the current function.  To get subs
+    // that will work on the overridden entry, we need to construct the
+    // override substitutions.
+    auto origFunctionType = baseCI.FormalPattern;
+    if (base.getDecl() == c.getDecl()) {
+      origFunctionType = origFunctionType.withSubstitutions(subs);
+    } else {
+      auto derivedCDR = ConcreteDeclRef(c.getDecl(), subs);
+      auto baseCDR = derivedCDR.getOverriddenDecl(base.getDecl());
+      origFunctionType =
+        origFunctionType.withSubstitutions(baseCDR.getSubstitutions());
+    }
+
     return Callee(
-        Kind::ClassMethod, SGF, c, baseCI.FormalPattern, derivedCI.FormalType,
-        subs.mapIntoTypeExpansionContext(SGF.getTypeExpansionContext()), l);
+        Kind::ClassMethod, SGF, c, origFunctionType, derivedCI.FormalType,
+        subs, l);
   }
   static Callee forSuperMethod(SILGenFunction &SGF,
                                SILDeclRef c, SubstitutionMap subs,
                                SILLocation l) {
     auto &ci = SGF.getConstantInfo(SGF.getTypeExpansionContext(), c);
+    subs = subs.mapIntoTypeExpansionContext(SGF.getTypeExpansionContext());
+    auto origFunctionType = ci.FormalPattern.withSubstitutions(subs);
     return Callee(
-        Kind::SuperMethod, SGF, c, ci.FormalPattern, ci.FormalType,
-        subs.mapIntoTypeExpansionContext(SGF.getTypeExpansionContext()), l);
+        Kind::SuperMethod, SGF, c, origFunctionType, ci.FormalType, subs, l);
   }
   static Callee forWitnessMethod(SILGenFunction &SGF,
                                  CanType protocolSelfType,

test/SILGen/variadic-generic-class-methods.swift

@@ -0,0 +1,56 @@
+// RUN: %target-swift-emit-silgen -disable-availability-checking %s | %FileCheck %s
+
+public class A<each T> {
+  public func f(_ action: (repeat each T) -> ()) {}
+}
+
+// CHECK-LABEL: sil hidden [ossa] @$s4main5test02fnyySi_SftXE_tF :
+// CHECK:       [[CTOR:%.*]] = function_ref @$s4main1ACACyxxQp_QPGycfC :
+// CHECK:       [[OBJECT:%.*]] = apply [[CTOR]]<Pack{Int, Float}>
+// CHECK:       [[METHOD:%.*]] = class_method [[OBJECT]] : $A<Int, Float>, #A.f : <each T> (A<repeat each T>) -> ((repeat each T) -> ()) -> ()
+// CHECK:       apply [[METHOD]]<Pack{Int, Float}>
+func test0(fn: (Int, Float) -> ()) {
+  A<Int, Float>().f(fn)
+}
+
+// CHECK-LABEL: sil hidden [ossa] @$s4main5test1yyF :
+// CHECK:       [[CTOR:%.*]] = function_ref @$s4main1ACACyxxQp_QPGycfC :
+// CHECK:       [[OBJECT:%.*]] = apply [[CTOR]]<Pack{Int, Float}>
+// CHECK:       [[METHOD:%.*]] = class_method [[OBJECT]] : $A<Int, Float>, #A.f : <each T> (A<repeat each T>) -> ((repeat each T) -> ()) -> ()
+// CHECK:       apply [[METHOD]]<Pack{Int, Float}>
+func test1() {
+  A<Int, Float>().f { a, b in }
+}
+
+// These are all currently forbidden
+#if false
+public class B<each T> : A<repeat each T> {
+  public override func f(_ action: (repeat each T) -> ()) {
+    super.f(action)
+  }
+}
+
+func test2() {
+  B<Int, Float>().f { a, b in }
+}
+
+public class C<T, each U> : A<T, repeat each U> {
+  public override func f(_ action: (T, repeat each U) -> ()) {
+    super.f(action)
+  }
+}
+
+func test3() {
+  C<Int, Float>().f { a, b in }
+}
+
+public class D : A<Int, Float> {
+  public override func f(_ action: (Int, Float) -> ()) {
+    super.f(action)
+  }
+}
+
+func test4() {
+  D().f { a, b in }
+}
+#endif