[3.1] SIL: Canonicalize capture types with the AST function's generic signature.

include/swift/AST/AnyFunctionRef.h

@@ -176,6 +176,16 @@ class AnyFunctionRef {
     }
     llvm_unreachable("unexpected AnyFunctionRef representation");
   }
+
+  GenericSignature *getGenericSignature() const {
+    if (auto afd = TheFunction.dyn_cast<AbstractFunctionDecl *>()) {
+      return afd->getGenericSignature();
+    }
+    if (auto ce = TheFunction.dyn_cast<AbstractClosureExpr *>()) {
+      return ce->getGenericSignatureOfContext();
+    }
+    llvm_unreachable("unexpected AnyFunctionRef representation");
+  }
 };
 #if SWIFT_COMPILER_IS_MSVC
 #pragma warning(pop)

lib/SIL/SILFunctionType.cpp

@@ -23,6 +23,7 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/DiagnosticsSIL.h"
 #include "swift/AST/ForeignErrorConvention.h"
+#include "swift/AST/GenericEnvironment.h"
 #include "swift/Basic/Fallthrough.h"
 #include "clang/Analysis/DomainSpecific/CocoaConventions.h"
 #include "clang/AST/Attr.h"
@@ -742,9 +743,17 @@ static CanSILFunctionType getSILFunctionType(SILModule &M,
   // from the function to which the argument is attached.
   if (constant && !constant->isDefaultArgGenerator())
   if (auto function = constant->getAnyFunctionRef()) {
-    auto getCanonicalType = [&](Type t) -> CanType {
-      if (genericSig)
-        return genericSig->getCanonicalTypeInContext(t, *M.getSwiftModule());
+    // NB: The generic signature may be elided from the lowered function type
+    // if the function is in a fully-specialized context, but we still need to
+    // canonicalize references to the generic parameters that may appear in
+    // non-canonical types in that context. We need the original generic
+    // signature from the AST for that.
+    auto origGenericSig
+      = function->getGenericSignature();
+    auto getCanonicalType = [origGenericSig, &M](Type t) -> CanType {
+      if (origGenericSig)
+        return origGenericSig->getCanonicalTypeInContext(t,
+                                                         *M.getSwiftModule());
       return t->getCanonicalType();
     };
 

lib/SIL/SILInstructions.cpp

@@ -408,6 +408,11 @@ ApplyInst *ApplyInst::create(SILDebugLocation Loc, SILValue Callee,
                              ArrayRef<SILValue> Args, bool isNonThrowing,
                              SILFunction &F,
                              SILOpenedArchetypesState &OpenedArchetypes) {
+  if (!F.getModule().Types.getCurGenericContext())
+    assert(Callee->getType().castTo<SILFunctionType>()
+                 ->substGenericArgs(F.getModule(), Subs)
+           == SubstCalleeTy.getSwiftRValueType());
+
   SmallVector<SILValue, 32> TypeDependentOperands;
   collectTypeDependentOperands(TypeDependentOperands, OpenedArchetypes, F,
                                SubstCalleeTy.getSwiftRValueType(), Subs);

lib/SILGen/SILGenProlog.cpp

@@ -344,7 +344,9 @@ void SILGenFunction::bindParametersForForwarding(const ParameterList *params,
   }
 }
 
-static void emitCaptureArguments(SILGenFunction &gen, CapturedValue capture,
+static void emitCaptureArguments(SILGenFunction &gen,
+                                 AnyFunctionRef closure,
+                                 CapturedValue capture,
                                  unsigned ArgNo) {
 
   auto *VD = capture.getDecl();
@@ -357,7 +359,12 @@ static void emitCaptureArguments(SILGenFunction &gen, CapturedValue capture,
     auto interfaceType = cast<VarDecl>(VD)->getInterfaceType();
     if (!interfaceType->hasTypeParameter()) return interfaceType;
 
-    auto genericEnv = gen.F.getGenericEnvironment();
+    // NB: The generic signature may be elided from the lowered function type
+    // if the function is in a fully-specialized context, but we still need to
+    // canonicalize references to the generic parameters that may appear in
+    // non-canonical types in that context. We need the original generic
+    // environment from the AST for that.
+    auto genericEnv = closure.getGenericEnvironment();
     return genericEnv->mapTypeIntoContext(gen.F.getModule().getSwiftModule(),
                                           interfaceType);
   };
@@ -446,7 +453,7 @@ void SILGenFunction::emitProlog(AnyFunctionRef TheClosure,
       return;
     }
 
-    emitCaptureArguments(*this, capture, ++ArgNo);
+    emitCaptureArguments(*this, TheClosure, capture, ++ArgNo);
   }
 }
 

test/SILGen/capture-canonicalization.swift

@@ -0,0 +1,14 @@
+// RUN: %target-swift-frontend -emit-silgen %s | %FileCheck %s
+
+struct Foo<T> {}
+struct Bar {}
+
+extension Foo where T == Bar {
+  func foo(x: T) -> Bar {
+    // CHECK-LABEL: sil shared @{{.*}}3foo{{.*}}4foo2{{.*}} : $@convention(thin) (Bar) -> Bar
+    func foo2() -> Bar {
+      return x
+    }
+    return foo2()
+  }
+}