SILGen: When forming a closure, apply substitutions before applying captures

slavapestov · slavapestov · commit ac110c0fc0d7 · 2016-05-28T22:30:41.000-07:00
We now have enough machinery in place to reference local generic
functions which have captures, to get a value of function type
that can be passed around.

Generic local functions still cannot be directly called from
function call expressions, since those go down a different
path in SILGenApply.cpp -- the next patch will add support for
this case.
diff --git a/lib/SILGen/SILGenExpr.cpp b/lib/SILGen/SILGenExpr.cpp
@@ -445,57 +445,9 @@ emitRValueForDecl(SILLocation loc, ConcreteDeclRef declRef, Type ncRefType,
   }
 
   auto silDeclRef = SILDeclRef(decl, ResilienceExpansion::Minimal, uncurryLevel);
-  auto constantInfo = getConstantInfo(silDeclRef);
-
-  ManagedValue result = emitFunctionRef(loc, silDeclRef, constantInfo);
-
-  // Get the lowered AST types:
-  //  - the original type
-  auto origLoweredFormalType =
-      AbstractionPattern(constantInfo.LoweredInterfaceType);
-  if (hasLocalCaptures) {
-    // Get the unlowered formal type of the constant, stripping off
-    // the first level of function application, which applies captures.
-    origLoweredFormalType =
-      AbstractionPattern(constantInfo.FormalInterfaceType)
-          .getFunctionResultType();
-
-    // Lower it, being careful to use the right generic signature.
-    origLoweredFormalType =
-      AbstractionPattern(
-          origLoweredFormalType.getGenericSignature(),
-          SGM.Types.getLoweredASTFunctionType(
-              cast<FunctionType>(origLoweredFormalType.getType()),
-              0, silDeclRef));
-  }
-
-  // - the substituted type
-  auto substFormalType = cast<AnyFunctionType>(refType);
-  auto substLoweredFormalType =
-    SGM.Types.getLoweredASTFunctionType(substFormalType, 0, silDeclRef);
-
-  // If the declaration reference is specialized, create the partial
-  // application.
-  if (declRef.isSpecialized()) {
-    // Substitute the function type.
-    auto origFnType = result.getType().castTo<SILFunctionType>();
-    auto substFnType = origFnType->substGenericArgs(
-                                                    SGM.M, SGM.SwiftModule,
-                                                    declRef.getSubstitutions());
-    auto closureType = adjustFunctionType(substFnType,
-                                        SILFunctionType::Representation::Thick);
-
-    SILValue spec = B.createPartialApply(loc, result.forward(*this),
-                                SILType::getPrimitiveObjectType(substFnType),
-                                         declRef.getSubstitutions(),
-                                         { },
-                                SILType::getPrimitiveObjectType(closureType));
-    result = emitManagedRValueWithCleanup(spec);
-  }
-
-  // Generalize if necessary.
-  result = emitOrigToSubstValue(loc, result, origLoweredFormalType,
-                                substLoweredFormalType);
+
+  ManagedValue result = emitClosureValue(loc, silDeclRef, refType,
+                                         declRef.getSubstitutions());
   return RValue(*this, loc, refType, result);
 }
 
@@ -1965,9 +1917,16 @@ RValue RValueEmitter::visitAbstractClosureExpr(AbstractClosureExpr *e,
   // Emit the closure body.
   SGF.SGM.emitClosure(e);
 
+  ArrayRef<Substitution> subs;
+  if (e->getCaptureInfo().hasGenericParamCaptures())
+    subs = SGF.getForwardingSubstitutions();
+
   // Generate the closure value (if any) for the closure expr's function
   // reference.
-  return RValue(SGF, e, SGF.emitClosureValue(e, SILDeclRef(e), e));
+  auto refType = e->getType()->getCanonicalType();
+  ManagedValue result = SGF.emitClosureValue(e, SILDeclRef(e),
+                                             refType, subs);
+  return RValue(SGF, e, refType, result);
 }
 
 RValue RValueEmitter::
diff --git a/lib/SILGen/SILGenFunction.cpp b/lib/SILGen/SILGenFunction.cpp
@@ -207,27 +207,11 @@ SILGenFunction::emitSiblingMethodRef(SILLocation loc,
                          methodTy, subs);
 }
 
-ManagedValue SILGenFunction::emitFunctionRef(SILLocation loc,
-                                             SILDeclRef constant,
-                                             SILConstantInfo constantInfo) {
-  // If the function has captures, apply them.
-  if (auto fn = constant.getAnyFunctionRef()) {
-    if (fn->getCaptureInfo().hasLocalCaptures() ||
-        fn->getCaptureInfo().hasGenericParamCaptures()) {
-      return emitClosureValue(loc, constant, *fn);
-    }
-  }
-
-  // Otherwise, use a global FunctionRefInst.
-  SILValue c = emitGlobalFunctionRef(loc, constant, constantInfo);
-  return ManagedValue::forUnmanaged(c);
-}
-
 void SILGenFunction::emitCaptures(SILLocation loc,
-                                  AnyFunctionRef TheClosure,
+                                  AnyFunctionRef closure,
                                   CaptureEmission purpose,
                                   SmallVectorImpl<ManagedValue> &capturedArgs) {
-  auto captureInfo = SGM.Types.getLoweredLocalCaptures(TheClosure);
+  auto captureInfo = SGM.Types.getLoweredLocalCaptures(closure);
   // For boxed captures, we need to mark the contained variables as having
   // escaped for DI diagnostics.
   SmallVector<SILValue, 2> escapesToMark;
@@ -349,31 +333,38 @@ void SILGenFunction::emitCaptures(SILLocation loc,
 
 ManagedValue
 SILGenFunction::emitClosureValue(SILLocation loc, SILDeclRef constant,
-                                 AnyFunctionRef TheClosure) {
+                                 CanType expectedType,
+                                 ArrayRef<Substitution> subs) {
+  auto closure = *constant.getAnyFunctionRef();
+  auto captureInfo = closure.getCaptureInfo();
+
   assert(((constant.uncurryLevel == 1 &&
-           TheClosure.getCaptureInfo().hasLocalCaptures()) ||
+           captureInfo.hasLocalCaptures()) ||
           (constant.uncurryLevel == 0 &&
-           !TheClosure.getCaptureInfo().hasLocalCaptures())) &&
+           !captureInfo.hasLocalCaptures())) &&
          "curried local functions not yet supported");
 
   auto constantInfo = getConstantInfo(constant);
   SILValue functionRef = emitGlobalFunctionRef(loc, constant, constantInfo);
   SILType functionTy = functionRef->getType();
 
-  auto expectedType =
-    cast<FunctionType>(TheClosure.getType()->getCanonicalType());
-
-  // Forward substitutions from the outer scope.
-
+  // Apply substitutions.
   auto pft = constantInfo.SILFnType;
 
-  auto forwardSubs = constantInfo.getForwardingSubstitutions(getASTContext());
-
   bool wasSpecialized = false;
-  if (pft->isPolymorphic() && !forwardSubs.empty()) {
+  if (pft->isPolymorphic()) {
+    // If the lowered function type is generic but Sema did not hand us any
+    // substitutions, the function is a local function that appears in a
+    // generic context but does not have a generic parameter list of its own;
+    // just use our forwarding substitutions.
+    if (subs.empty()) {
+      assert(closure.getAsDeclContext()->isLocalContext() &&
+             "cannot reference generic global function without substitutions");
+      subs = getForwardingSubstitutions();
+    }
     auto specialized = pft->substGenericArgs(F.getModule(),
                                              F.getModule().getSwiftModule(),
-                                             forwardSubs);
+                                             subs);
     functionTy = SILType::getPrimitiveObjectType(specialized);
     wasSpecialized = true;
   }
@@ -382,19 +373,19 @@ SILGenFunction::emitClosureValue(SILLocation loc, SILDeclRef constant,
   // globals, but we still need to mark them as escaping so that DI can flag
   // uninitialized uses.
   if (this == SGM.TopLevelSGF) {
-    SGM.emitMarkFunctionEscapeForTopLevelCodeGlobals(loc,
-                                                   TheClosure.getCaptureInfo());
+    SGM.emitMarkFunctionEscapeForTopLevelCodeGlobals(
+        loc, captureInfo);
   }
 
-  if (!TheClosure.getCaptureInfo().hasLocalCaptures() && !wasSpecialized) {
+  if (!captureInfo.hasLocalCaptures() && !wasSpecialized) {
     auto result = ManagedValue::forUnmanaged(functionRef);
     return emitOrigToSubstValue(loc, result,
                                 AbstractionPattern(expectedType),
                                 expectedType);
   }
 
   SmallVector<ManagedValue, 4> capturedArgs;
-  emitCaptures(loc, TheClosure, CaptureEmission::PartialApplication,
+  emitCaptures(loc, closure, CaptureEmission::PartialApplication,
                capturedArgs);
 
   // The partial application takes ownership of the context parameters.
@@ -405,15 +396,42 @@ SILGenFunction::emitClosureValue(SILLocation loc, SILDeclRef constant,
   SILType closureTy =
     SILGenBuilder::getPartialApplyResultType(functionRef->getType(),
                                              capturedArgs.size(), SGM.M,
-                                             forwardSubs);
+                                             subs);
   auto toClosure =
     B.createPartialApply(loc, functionRef, functionTy,
-                         forwardSubs, forwardedArgs, closureTy);
+                         subs, forwardedArgs, closureTy);
   auto result = emitManagedRValueWithCleanup(toClosure);
 
-  return emitOrigToSubstValue(loc, result,
-                              AbstractionPattern(expectedType),
-                              expectedType);
+  // Get the lowered AST types:
+  //  - the original type
+  auto origLoweredFormalType =
+      AbstractionPattern(constantInfo.LoweredInterfaceType);
+  if (captureInfo.hasLocalCaptures()) {
+    // Get the unlowered formal type of the constant, stripping off
+    // the first level of function application, which applies captures.
+    origLoweredFormalType =
+      AbstractionPattern(constantInfo.FormalInterfaceType)
+          .getFunctionResultType();
+
+    // Lower it, being careful to use the right generic signature.
+    origLoweredFormalType =
+      AbstractionPattern(
+          origLoweredFormalType.getGenericSignature(),
+          SGM.Types.getLoweredASTFunctionType(
+              cast<FunctionType>(origLoweredFormalType.getType()),
+              0, constant));
+  }
+
+  // - the substituted type
+  auto substFormalType = cast<FunctionType>(expectedType);
+  auto substLoweredFormalType =
+    SGM.Types.getLoweredASTFunctionType(substFormalType, 0, constant);
+
+  // Generalize if necessary.
+  result = emitOrigToSubstValue(loc, result, origLoweredFormalType,
+                                substLoweredFormalType);
+
+  return result;
 }
 
 void SILGenFunction::emitFunction(FuncDecl *fd) {
diff --git a/lib/SILGen/SILGenFunction.h b/lib/SILGen/SILGenFunction.h
@@ -1045,12 +1045,6 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   SILValue emitDynamicMethodRef(SILLocation loc, SILDeclRef constant,
                                 SILConstantInfo constantInfo);  
 
-  /// Returns a reference to a constant in local context. This will return a
-  /// retained closure object reference if the constant refers to a local func
-  /// decl.
-  ManagedValue emitFunctionRef(SILLocation loc, SILDeclRef constant,
-                               SILConstantInfo constantInfo);
-  
   /// Emit the specified VarDecl as an LValue if possible, otherwise return
   /// null.
   ManagedValue emitLValueForDecl(SILLocation loc, VarDecl *var,
@@ -1081,9 +1075,13 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
                     CaptureEmission purpose,
                     SmallVectorImpl<ManagedValue> &captures);
 
+  /// Produce a reference to a function, which may be a local function
+  /// with captures. If the function is generic, substitutions must be
+  /// given. The result is re-abstracted to the given expected type.
   ManagedValue emitClosureValue(SILLocation loc,
                                 SILDeclRef function,
-                                AnyFunctionRef TheClosure);
+                                CanType expectedType,
+                                ArrayRef<Substitution> subs);
   
   ArgumentSource prepareAccessorBaseArg(SILLocation loc, ManagedValue base,
                                         CanType baseFormalType,
diff --git a/test/SILGen/generic_closures.swift b/test/SILGen/generic_closures.swift
