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

When lowering closures, we avoid capturing the enclosing generic context when possible. However, the generic context may still be necessary to canonicalize types mentioned inside the closure, such as when an associated type is referred to that is same-typed to a concrete type. Fixes rdar://problem/30254048.

include/swift/AST/AnyFunctionRef.h

@@ -176,6 +176,16 @@ public:
     }
     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"
@@ -757,9 +758,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 {
+    // NB: The generic signature may be elided from the lowered function type
-      if (genericSig)
+    // if the function is in a fully-specialized context, but we still need to
-        return genericSig->getCanonicalTypeInContext(t, *M.getSwiftModule());
+    // 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

@@ -409,6 +409,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

@@ -346,7 +346,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();
@@ -359,7 +361,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);
   };
@@ -448,7 +455,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()
+  }
+}