From acf923b9fe02ebd1dd25bf95136a15737885de93 Mon Sep 17 00:00:00 2001 From: Erik Eckstein Date: Thu, 11 Jun 2020 18:56:44 +0200 Subject: [PATCH] MandatoryInlining: fix a memory lifetime bug related to partial_apply with in_guaranteed parameters. Because partial_apply consumes it's arguments we need to copy them. This was done for "direct" parameters but not for address parameters. rdar://problem/64035105 --- .../Mandatory/MandatoryInlining.cpp | 53 ++++++++++++++----- test/SILOptimizer/mandatory_inlining.sil | 35 ++++++++++++ 2 files changed, 75 insertions(+), 13 deletions(-) diff --git a/lib/SILOptimizer/Mandatory/MandatoryInlining.cpp b/lib/SILOptimizer/Mandatory/MandatoryInlining.cpp index ce7c2dcf951ed..6af5002074429 100644 --- a/lib/SILOptimizer/Mandatory/MandatoryInlining.cpp +++ b/lib/SILOptimizer/Mandatory/MandatoryInlining.cpp @@ -59,7 +59,10 @@ static SILValue stripCopiesAndBorrows(SILValue v) { /// It is important to note that, we can not assume that the partial apply, the /// apply site, or the callee value are control dependent in any way. This /// requires us to need to be very careful. See inline comments. -static void fixupReferenceCounts( +/// +/// Returns true if the stack nesting is invalidated and must be corrected +/// afterwards. +static bool fixupReferenceCounts( PartialApplyInst *pai, FullApplySite applySite, SILValue calleeValue, ArrayRef captureArgConventions, MutableArrayRef capturedArgs, bool isCalleeGuaranteed) { @@ -72,6 +75,7 @@ static void fixupReferenceCounts( // FIXME: Can we cache this in between inlining invocations? DeadEndBlocks deadEndBlocks(pai->getFunction()); SmallVector leakingBlocks; + bool invalidatedStackNesting = false; // Add a copy of each non-address type capture argument to lifetime extend the // captured argument over at least the inlined function and till the end of a @@ -83,14 +87,6 @@ static void fixupReferenceCounts( for (unsigned i : indices(captureArgConventions)) { auto convention = captureArgConventions[i]; SILValue &v = capturedArgs[i]; - if (v->getType().isAddress()) { - // FIXME: What about indirectly owned parameters? The invocation of the - // closure would perform an indirect copy which we should mimick here. - assert(convention != ParameterConvention::Indirect_In && - "Missing indirect copy"); - continue; - } - auto *f = applySite.getFunction(); // See if we have a trivial value. In such a case, just continue. We do not @@ -102,11 +98,40 @@ static void fixupReferenceCounts( switch (convention) { case ParameterConvention::Indirect_In: + llvm_unreachable("Missing indirect copy"); + case ParameterConvention::Indirect_In_Constant: case ParameterConvention::Indirect_Inout: case ParameterConvention::Indirect_InoutAliasable: - case ParameterConvention::Indirect_In_Guaranteed: - llvm_unreachable("Should be handled above"); + break; + + case ParameterConvention::Indirect_In_Guaranteed: { + // Do the same as for Direct_Guaranteed, just the address version. + // (See comment below). + SILBuilderWithScope builder(pai); + auto *stackLoc = builder.createAllocStack(loc, v->getType().getObjectType()); + builder.createCopyAddr(loc, v, stackLoc, IsNotTake, IsInitialization); + visitedBlocks.clear(); + + LinearLifetimeChecker checker(visitedBlocks, deadEndBlocks); + bool consumedInLoop = checker.completeConsumingUseSet( + pai, applySite.getCalleeOperand(), + [&](SILBasicBlock::iterator insertPt) { + SILBuilderWithScope builder(insertPt); + builder.createDestroyAddr(loc, stackLoc); + builder.createDeallocStack(loc, stackLoc); + }); + + if (!consumedInLoop) { + applySite.insertAfterInvocation([&](SILBasicBlock::iterator iter) { + SILBuilderWithScope(iter).createDestroyAddr(loc, stackLoc); + SILBuilderWithScope(iter).createDeallocStack(loc, stackLoc); + }); + } + v = stackLoc; + invalidatedStackNesting = true; + break; + } case ParameterConvention::Direct_Guaranteed: { // If we have a direct_guaranteed value, the value is being taken by the @@ -242,6 +267,7 @@ static void fixupReferenceCounts( SILBuilderWithScope(iter).emitDestroyValueOperation(loc, calleeValue); }); } + return invalidatedStackNesting; } static SILValue cleanupLoadedCalleeValue(SILValue calleeValue, LoadInst *li) { @@ -917,8 +943,9 @@ runOnFunctionRecursively(SILOptFunctionBuilder &FuncBuilder, // We need to insert the copies before the partial_apply since if we can // not remove the partial_apply the captured values will be dead by the // time we hit the call site. - fixupReferenceCounts(PAI, InnerAI, CalleeValue, CapturedArgConventions, - CapturedArgs, IsCalleeGuaranteed); + invalidatedStackNesting |= fixupReferenceCounts(PAI, InnerAI, + CalleeValue, CapturedArgConventions, + CapturedArgs, IsCalleeGuaranteed); } // Register a callback to record potentially unused function values after diff --git a/test/SILOptimizer/mandatory_inlining.sil b/test/SILOptimizer/mandatory_inlining.sil index 911343961cc62..ad225f0ae1683 100644 --- a/test/SILOptimizer/mandatory_inlining.sil +++ b/test/SILOptimizer/mandatory_inlining.sil @@ -1428,3 +1428,38 @@ bb0(%0 : $@callee_guaranteed () -> ()): %9999 = tuple() return %9999 : $() } + +struct Mystruct { + var s: String +} + +sil @use_string : $@convention(thin) (@guaranteed String) -> () + +sil private [transparent] [ossa] @callee_with_guaranteed : $@convention(thin) (Bool, @in_guaranteed Mystruct) -> () { +bb0(%0 : $Bool, %1 : $*Mystruct): + %4 = struct_element_addr %1 : $*Mystruct, #Mystruct.s + %5 = load [copy] %4 : $*String + %6 = function_ref @use_string : $@convention(thin) (@guaranteed String) -> () + %7 = apply %6(%5) : $@convention(thin) (@guaranteed String) -> () + destroy_value %5 : $String + %19 = tuple () + return %19 : $() +} + +// Make sure that this doesn't cause a memory lifetime failure. +// +// CHECK-LABEL: sil [ossa] @partial_apply_with_indirect_param +// CHECK: } // end sil function 'partial_apply_with_indirect_param' +sil [ossa] @partial_apply_with_indirect_param : $@convention(thin) (@in_guaranteed Mystruct, Bool) -> () { +bb0(%0 : $*Mystruct, %1 : $Bool): + %216 = function_ref @callee_with_guaranteed : $@convention(thin) (Bool, @in_guaranteed Mystruct) -> () + %217 = alloc_stack $Mystruct + copy_addr %0 to [initialization] %217 : $*Mystruct + %219 = partial_apply [callee_guaranteed] %216(%217) : $@convention(thin) (Bool, @in_guaranteed Mystruct) -> () + %220 = apply %219(%1) : $@callee_guaranteed (Bool) -> () + destroy_value %219 : $@callee_guaranteed (Bool) -> () + dealloc_stack %217 : $*Mystruct + %19 = tuple () + return %19 : $() +} +