Use the right generic signature when producing a substituted function…

include/swift/SIL/TypeLowering.h

@@ -598,6 +598,8 @@ class TypeConverter {
 
     TypeExpansionContext expansionContext;
 
+    bool IsCacheable;
+
     CachingTypeKey getCachingKey() const {
       assert(isCacheable());
       return { (OrigType.hasGenericSignature()
@@ -609,19 +611,37 @@ class TypeConverter {
     }
 
     bool isCacheable() const {
-      return OrigType.hasCachingKey();
+      return IsCacheable;
     }
 
     TypeKey getKeyForMinimalExpansion() const {
-      return {OrigType, SubstType, TypeExpansionContext::minimal()};
+      return {OrigType, SubstType, TypeExpansionContext::minimal(),
+              IsCacheable};
+    }
+
+    void computeCacheable() {
+      IsCacheable = (OrigType.hasCachingKey() &&
+                     !isTreacherousInterfaceType(SubstType));
+    }
+
+    static bool isTreacherousInterfaceType(CanType type) {
+      // Don't cache lowerings for interface function types that involve
+      // type parameters; we might need a contextual generic signature to
+      // handle them correctly.
+      if (!type->hasTypeParameter()) return false;
+      return type.findIf([](CanType type) {
+        return isa<FunctionType>(type) && type->hasTypeParameter();
+      });
     }
   };
 
   friend struct llvm::DenseMapInfo<CachingTypeKey>;
 
   TypeKey getTypeKey(AbstractionPattern origTy, CanType substTy,
                      TypeExpansionContext context) {
-    return {origTy, substTy, context};
+    TypeKey result = {origTy, substTy, context, false};
+    result.computeCacheable();
+    return result;
   }
 
   struct OverrideKey {
@@ -643,14 +663,16 @@ class TypeConverter {
 
   /// Find a cached TypeLowering by TypeKey, or return null if one doesn't
   /// exist.
-  const TypeLowering *find(TypeKey k);
+  const TypeLowering *find(const TypeKey &k);
   /// Insert a mapping into the cache.
-  void insert(TypeKey k, const TypeLowering *tl);
+  void insert(const TypeKey &k, const TypeLowering *tl);
 #ifndef NDEBUG
   /// Remove the nullptr entry from the type map.
-  void removeNullEntry(TypeKey k);
+  void removeNullEntry(const TypeKey &k);
 #endif
 
+  CanGenericSignature CurGenericSignature;
+
   /// Mapping for types independent on contextual generic parameters.
   llvm::DenseMap<CachingTypeKey, const TypeLowering *> LoweredTypes;
 
@@ -696,6 +718,27 @@ class TypeConverter {
   TypeConverter(TypeConverter const &) = delete;
   TypeConverter &operator=(TypeConverter const &) = delete;
 
+  CanGenericSignature getCurGenericSignature() const {
+    return CurGenericSignature;
+  }
+
+  class GenericContextRAII {
+    TypeConverter &TC;
+    CanGenericSignature SavedSig;
+  public:
+    GenericContextRAII(TypeConverter &TC, CanGenericSignature sig)
+        : TC(TC), SavedSig(TC.CurGenericSignature) {
+      TC.CurGenericSignature = sig;
+    }
+
+    GenericContextRAII(const GenericContextRAII &) = delete;
+    GenericContextRAII &operator=(const GenericContextRAII &) = delete;
+
+    ~GenericContextRAII() {
+      TC.CurGenericSignature = SavedSig;
+    }
+  };
+
   /// Return the CaptureKind to use when capturing a decl.
   CaptureKind getDeclCaptureKind(CapturedValue capture,
                                  TypeExpansionContext expansion);

lib/SIL/SILFunctionType.cpp

@@ -912,7 +912,10 @@ class SubstFunctionTypeCollector {
     CanGenericSignature origSig = origType.getGenericSignature();
     if (substituteBindingsInSubstType) {
       origContextType = substType;
-      origSig = GenericSig;
+      origSig = TC.getCurGenericSignature();
+      assert((!substType->hasTypeParameter() || origSig) &&
+             "lowering mismatched interface types in a context without "
+             "a generic signature");
     }
 
     if (!origContextType->hasTypeParameter()
@@ -925,10 +928,11 @@ class SubstFunctionTypeCollector {
     }
 
     // Extract structural substitutions.
-    if (origContextType->hasTypeParameter())
+    if (origContextType->hasTypeParameter()) {
       origContextType = origSig->getGenericEnvironment()
         ->mapTypeIntoContext(origContextType)
         ->getCanonicalType(origSig);
+    }
 
     auto result = origContextType
       ->substituteBindingsTo(substType,
@@ -1648,6 +1652,9 @@ static CanSILFunctionType getSILFunctionType(
   CanGenericSignature genericSig =
     substFnInterfaceType.getOptGenericSignature();
 
+  Optional<TypeConverter::GenericContextRAII> contextRAII;
+  if (genericSig) contextRAII.emplace(TC, genericSig);
+
   // Per above, only fully honor opaqueness in the abstraction pattern
   // for thick or polymorphic functions.  We don't need to worry about
   // non-opaque patterns because the type-checker forbids non-thick

lib/SIL/TypeLowering.cpp

@@ -1385,7 +1385,7 @@ void *TypeLowering::operator new(size_t size, TypeConverter &tc) {
   return tc.TypeLoweringBPA.Allocate(size, alignof(TypeLowering&));
 }
 
-const TypeLowering *TypeConverter::find(TypeKey k) {
+const TypeLowering *TypeConverter::find(const TypeKey &k) {
   if (!k.isCacheable()) return nullptr;
 
   auto ck = k.getCachingKey();
@@ -1399,7 +1399,7 @@ const TypeLowering *TypeConverter::find(TypeKey k) {
 }
 
 #ifndef NDEBUG
-void TypeConverter::removeNullEntry(TypeKey k) {
+void TypeConverter::removeNullEntry(const TypeKey &k) {
   if (!k.isCacheable())
     return;
 
@@ -1413,7 +1413,7 @@ void TypeConverter::removeNullEntry(TypeKey k) {
 }
 #endif
 
-void TypeConverter::insert(TypeKey k, const TypeLowering *tl) {
+void TypeConverter::insert(const TypeKey &k, const TypeLowering *tl) {
   if (!k.isCacheable()) return;
 
   LoweredTypes[k.getCachingKey()] = tl;

test/SILGen/variant_overrides.swift

@@ -0,0 +1,17 @@
+// RUN: %target-swift-emit-silgen %s | %FileCheck %s
+
+// CHECK-LABEL: sil hidden [ossa] @$s17variant_overrides1AC3foo5blockyyACc_tF :
+// CHECK-SAME:    $@convention(method) (@guaranteed @callee_guaranteed (@guaranteed A) -> (), @guaranteed A) -> ()
+class A {
+  func foo(block: @escaping (A) -> Void) {}
+}
+
+// CHECK-LABEL: sil hidden [ossa] @$s17variant_overrides1BC3foo5blockyyACyxGc_tF :
+// CHECK-SAME:    $@convention(method) <T> (@guaranteed @callee_guaranteed @substituted <τ_0_0> (@guaranteed B<τ_0_0>) -> () for <T>, @guaranteed B<T>) -> () 
+class B<T> : A {
+  override func foo(block: @escaping (B<T>) -> Void) {}
+}
+
+//   FIXME: allowing this without a thunk silently reinterprets the function to a different abstraction!
+// CHECK-LABEL: sil_vtable B {
+// CHECK:          #A.foo!1: (A) -> (@escaping (A) -> ()) -> () : @$s17variant_overrides1BC3foo5blockyyACyxGc_tF [override]