Skip to content

Commit 4acffab

Browse files
eecksteineeckstein
committed
SILOptimizer: fix a crash in ArrayElementValuePropagation.
This optimization is doing 2 things: replacing getElement calls and replacing append(contentOf:) calls. Now if the argument to a append(contentOf:) is a previously replaced getElement call, we ended up in a use-after-free crash. The main change here is to do the transformations immediately after gathering the data (and that means: separately) and not collecting all the data and do both transformation afterwards The pass does not use any Analysis (where invalidation could be a problem). Also iterator invalidation is not a problem here. SR-10003 rdar://problem/48445856
1 parent e80c4d1 commit 4acffab

File tree

2 files changed

+131
-105
lines changed

2 files changed

+131
-105
lines changed

‎lib/SILOptimizer/Transforms/ArrayElementValuePropagation.cpp

Lines changed: 111 additions & 105 deletions
Original file line numberDiff line numberDiff line change
@@ -76,20 +76,6 @@ class ArrayAllocation {
7676

7777
public:
7878

79-
/// Specifies the value with which a get-element call can be replaced.
80-
struct GetElementReplacement {
81-
ApplyInst *GetElementCall;
82-
SILValue Replacement;
83-
};
84-
85-
/// Specifies the set of elements with which an append-contentof call can be
86-
/// replaced.
87-
struct AppendContentOfReplacement {
88-
ApplyInst *AppendContentOfCall;
89-
llvm::SmallVector<SILValue, 4> ReplacementValues;
90-
SILValue Array;
91-
};
92-
9379
ArrayAllocation() {}
9480

9581
/// Analyzes an array allocation call.
@@ -99,14 +85,11 @@ class ArrayAllocation {
9985
/// or append(contentof) calls.
10086
bool analyze(ApplyInst *Alloc);
10187

102-
/// Gets the list of get_element calls which can be replaced.
103-
void getGetElementReplacements(
104-
llvm::SmallVectorImpl<GetElementReplacement> &Replacements);
88+
/// Replace getElement calls with the actual values.
89+
bool replaceGetElements();
10590

106-
/// Gets the list of append(contentof) calls which can be replaced by a
107-
/// set of values.
108-
void getAppendContentOfReplacements(
109-
llvm::SmallVectorImpl<AppendContentOfReplacement> &Replacements);
91+
/// Replace append(contentsOf:) with multiple append(element:)
92+
bool replaceAppendContentOf();
11093
};
11194

11295
/// Map the indices of array element initialization stores to their values.
@@ -222,6 +205,10 @@ bool ArrayAllocation::recursivelyCollectUses(ValueBase *Def) {
222205
}
223206

224207
bool ArrayAllocation::analyze(ApplyInst *Alloc) {
208+
GetElementCalls.clear();
209+
AppendContentsOfCalls.clear();
210+
ElementValueMap.clear();
211+
225212
ArraySemanticsCall Uninitialized(Alloc, "array.uninitialized");
226213
if (!Uninitialized)
227214
return false;
@@ -245,8 +232,16 @@ bool ArrayAllocation::analyze(ApplyInst *Alloc) {
245232
return true;
246233
}
247234

248-
void ArrayAllocation::getGetElementReplacements(
249-
llvm::SmallVectorImpl<GetElementReplacement> &Replacements) {
235+
/// Replace getElement calls with the actual values.
236+
///
237+
/// \code
238+
/// store %x to %element_address
239+
/// ...
240+
/// %e = apply %getElement(%array, %constant_index)
241+
/// \endcode
242+
/// The value %e is replaced with %x.
243+
bool ArrayAllocation::replaceGetElements() {
244+
bool Changed = false;
250245
for (auto *GetElementCall : GetElementCalls) {
251246
ArraySemanticsCall GetElement(GetElementCall);
252247
assert(GetElement.getKind() == ArrayCallKind::kGetElement);
@@ -261,85 +256,107 @@ void ArrayAllocation::getGetElementReplacements(
261256
if (EltValueIt == ElementValueMap.end())
262257
continue;
263258

264-
Replacements.push_back({GetElementCall, EltValueIt->second});
259+
Changed |= GetElement.replaceByValue(EltValueIt->second);
265260
}
261+
return Changed;
266262
}
267263

268-
void ArrayAllocation::getAppendContentOfReplacements(
269-
llvm::SmallVectorImpl<AppendContentOfReplacement> &Replacements) {
264+
/// Replace append(contentsOf:) with multiple append(element:)
265+
///
266+
/// \code
267+
/// store %x to %source_array_element_address_0
268+
/// store %y to %source_array_element_address_1
269+
/// ...
270+
/// apply %append_contentsOf(%dest_array, %source_array)
271+
/// \endcode
272+
/// is replaced by
273+
/// \code
274+
/// store %x to %source_array_element_address_0
275+
/// store %y to %source_array_element_address_1
276+
/// ...
277+
/// apply %reserveCapacityForAppend(%dest_array, %number_of_values)
278+
/// apply %append_element(%dest_array, %x)
279+
/// apply %append_element(%dest_array, %y)
280+
/// ...
281+
/// \endcode
282+
/// The source_array and its initialization code can then be deleted (if not
283+
/// used otherwise).
284+
bool ArrayAllocation::replaceAppendContentOf() {
270285
if (AppendContentsOfCalls.empty())
271-
return;
286+
return false;
287+
if (ElementValueMap.empty())
288+
return false;
289+
290+
// Check if there is a store to each element.
272291
if (!replacementsAreValid())
273-
return;
292+
return false;
274293

275294
llvm::SmallVector<SILValue, 4> ElementValueVector;
276-
for (unsigned i = 0; i < ElementValueMap.size(); ++i)
277-
ElementValueVector.push_back(ElementValueMap[i]);
295+
for (unsigned i = 0; i < ElementValueMap.size(); ++i) {
296+
SILValue V = ElementValueMap[i];
297+
ElementValueVector.push_back(V);
298+
}
278299

279-
for (auto *Call : AppendContentsOfCalls)
280-
Replacements.push_back({Call, ElementValueVector, ArrayValue});
281-
}
300+
SILFunction *Fn = AppendContentsOfCalls[0]->getFunction();
301+
SILModule &M = Fn->getModule();
302+
ASTContext &Ctx = M.getASTContext();
282303

283-
// =============================================================================
284-
// Driver
285-
// =============================================================================
304+
LLVM_DEBUG(llvm::dbgs() << "Array append contentsOf calls replaced in "
305+
<< Fn->getName() << "\n");
286306

287-
class ArrayElementPropagation : public SILFunctionTransform {
288-
public:
289-
ArrayElementPropagation() {}
307+
// Get the needed Array helper functions.
308+
FuncDecl *AppendFnDecl = Ctx.getArrayAppendElementDecl();
309+
if (!AppendFnDecl)
310+
return false;
290311

291-
bool replaceAppendCalls(
292-
ArrayRef<ArrayAllocation::AppendContentOfReplacement> Repls) {
293-
auto &Fn = *getFunction();
294-
auto &M = Fn.getModule();
295-
auto &Ctx = M.getASTContext();
312+
FuncDecl *ReserveFnDecl = Ctx.getArrayReserveCapacityDecl();
313+
if (!ReserveFnDecl)
314+
return false;
296315

297-
if (Repls.empty())
298-
return false;
316+
auto Mangled = SILDeclRef(AppendFnDecl, SILDeclRef::Kind::Func).mangle();
317+
SILFunction *AppendFn = M.findFunction(Mangled, SILLinkage::PublicExternal);
318+
if (!AppendFn)
319+
return false;
299320

300-
LLVM_DEBUG(llvm::dbgs() << "Array append contentsOf calls replaced in "
301-
<< Fn.getName() << " (" << Repls.size() << ")\n");
321+
Mangled = SILDeclRef(ReserveFnDecl, SILDeclRef::Kind::Func).mangle();
322+
SILFunction *ReserveFn = M.findFunction(Mangled, SILLinkage::PublicExternal);
323+
if (!ReserveFn)
324+
return false;
302325

303-
FuncDecl *AppendFnDecl = Ctx.getArrayAppendElementDecl();
304-
if (!AppendFnDecl)
305-
return false;
326+
bool Changed = false;
327+
// Usually there is only a single append(contentsOf:) call. But there might
328+
// be multiple - with the same source array to append.
329+
for (ApplyInst *AppendContentOfCall : AppendContentsOfCalls) {
330+
ArraySemanticsCall AppendContentsOf(AppendContentOfCall);
331+
assert(AppendContentsOf && "Must be AppendContentsOf call");
306332

307-
FuncDecl *ReserveFnDecl = Ctx.getArrayReserveCapacityDecl();
308-
if (!ReserveFnDecl)
309-
return false;
333+
NominalTypeDecl *AppendSelfArray = AppendContentsOf.getSelf()->getType().
334+
getASTType()->getAnyNominal();
310335

311-
auto Mangled = SILDeclRef(AppendFnDecl, SILDeclRef::Kind::Func).mangle();
312-
SILFunction *AppendFn = M.findFunction(Mangled, SILLinkage::PublicExternal);
313-
if (!AppendFn)
314-
return false;
315-
316-
Mangled = SILDeclRef(ReserveFnDecl, SILDeclRef::Kind::Func).mangle();
317-
SILFunction *ReserveFn = M.findFunction(Mangled, SILLinkage::PublicExternal);
318-
if (!ReserveFn)
319-
return false;
336+
// In case if it's not an Array, but e.g. an ContiguousArray
337+
if (AppendSelfArray != Ctx.getArrayDecl())
338+
continue;
320339

321-
for (const ArrayAllocation::AppendContentOfReplacement &Repl : Repls) {
322-
ArraySemanticsCall AppendContentsOf(Repl.AppendContentOfCall);
323-
assert(AppendContentsOf && "Must be AppendContentsOf call");
340+
SILType ArrayType = ArrayValue->getType();
341+
auto *NTD = ArrayType.getASTType()->getAnyNominal();
342+
SubstitutionMap ArraySubMap = ArrayType.getASTType()
343+
->getContextSubstitutionMap(M.getSwiftModule(), NTD);
324344

325-
NominalTypeDecl *AppendSelfArray = AppendContentsOf.getSelf()->getType().
326-
getASTType()->getAnyNominal();
345+
AppendContentsOf.replaceByAppendingValues(M, AppendFn, ReserveFn,
346+
ElementValueVector,
347+
ArraySubMap);
348+
Changed = true;
349+
}
350+
return Changed;
351+
}
327352

328-
// In case if it's not an Array, but e.g. an ContiguousArray
329-
if (AppendSelfArray != Ctx.getArrayDecl())
330-
continue;
353+
// =============================================================================
354+
// Driver
355+
// =============================================================================
331356

332-
SILType ArrayType = Repl.Array->getType();
333-
auto *NTD = ArrayType.getASTType()->getAnyNominal();
334-
SubstitutionMap ArraySubMap = ArrayType.getASTType()
335-
->getContextSubstitutionMap(M.getSwiftModule(), NTD);
336-
337-
AppendContentsOf.replaceByAppendingValues(M, AppendFn, ReserveFn,
338-
Repl.ReplacementValues,
339-
ArraySubMap);
340-
}
341-
return true;
342-
}
357+
class ArrayElementPropagation : public SILFunctionTransform {
358+
public:
359+
ArrayElementPropagation() {}
343360

344361
void run() override {
345362
auto &Fn = *getFunction();
@@ -348,39 +365,28 @@ class ArrayElementPropagation : public SILFunctionTransform {
348365
if (Fn.hasOwnership())
349366
return;
350367

351-
// Propagate the elements an of array value to its users.
352-
llvm::SmallVector<ArrayAllocation::GetElementReplacement, 16>
353-
GetElementReplacements;
354-
llvm::SmallVector<ArrayAllocation::AppendContentOfReplacement, 4>
355-
AppendContentsOfReplacements;
368+
bool Changed = false;
356369

357370
for (auto &BB :Fn) {
358371
for (auto &Inst : BB) {
359372
if (auto *Apply = dyn_cast<ApplyInst>(&Inst)) {
360373
ArrayAllocation ALit;
361-
if (ALit.analyze(Apply)) {
362-
ALit.getGetElementReplacements(GetElementReplacements);
363-
ALit.getAppendContentOfReplacements(AppendContentsOfReplacements);
374+
if (!ALit.analyze(Apply))
375+
continue;
376+
377+
// First optimization: replace getElemente calls.
378+
if (ALit.replaceGetElements()) {
379+
Changed = true;
380+
// Re-do the analysis if the SIL changed.
381+
if (!ALit.analyze(Apply))
382+
continue;
364383
}
384+
// Second optimization: replace append(contentsOf:) calls.
385+
Changed |= ALit.replaceAppendContentOf();
365386
}
366387
}
367388
}
368389

369-
LLVM_DEBUG(if (!GetElementReplacements.empty()) {
370-
llvm::dbgs() << "Array elements replaced in " << Fn.getName() << " ("
371-
<< GetElementReplacements.size() << ")\n";
372-
});
373-
374-
bool Changed = false;
375-
376-
// Perform the actual replacement of the get_element call by its value.
377-
for (ArrayAllocation::GetElementReplacement &Repl : GetElementReplacements) {
378-
ArraySemanticsCall GetElement(Repl.GetElementCall);
379-
Changed |= GetElement.replaceByValue(Repl.Replacement);
380-
}
381-
382-
Changed |= replaceAppendCalls(AppendContentsOfReplacements);
383-
384390
if (Changed) {
385391
PM->invalidateAnalysis(
386392
&Fn, SILAnalysis::InvalidationKind::CallsAndInstructions);

‎test/SILOptimizer/array_element_propagation_crash.swift

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
// RUN: %empty-directory(%t)
2+
// RUN: %target-build-swift -O %s -o %t/a.out
3+
// RUN: %target-run %t/a.out | %FileCheck %s
4+
5+
// REQUIRES: executable_test
6+
7+
// Check that ArrayElementPropagation does not crash when compiling this function.
8+
@inline(never)
9+
func testit(_ arr: inout [Int]) {
10+
let a = [28]
11+
arr += [a[0]]
12+
}
13+
14+
var a = [27]
15+
testit(&a)
16+
17+
// As a bonus, also check if the code works as expected.
18+
// CHECK: [27, 28]
19+
print(a)
20+

0 commit comments

Comments
 (0)