diff --git a/transform/allocs.go b/transform/allocs.go index 4f006c3a..f21dde12 100644 --- a/transform/allocs.go +++ b/transform/allocs.go @@ -69,8 +69,13 @@ func OptimizeAllocs(mod llvm.Module) { sizeInWords := (size + uint64(alignment) - 1) / uint64(alignment) allocaType := llvm.ArrayType(mod.Context().IntType(alignment*8), int(sizeInWords)) alloca := builder.CreateAlloca(allocaType, "stackalloc.alloca") + + // Zero the allocation inside the block where the value was originally allocated. zero := llvm.ConstNull(alloca.Type().ElementType()) + builder.SetInsertPointBefore(bitcast) builder.CreateStore(zero, alloca) + + // Replace heap alloc bitcast with stack alloc bitcast. stackalloc := builder.CreateBitCast(alloca, bitcast.Type(), "stackalloc") bitcast.ReplaceAllUsesWith(stackalloc) if heapalloc != bitcast { diff --git a/transform/testdata/allocs.ll b/transform/testdata/allocs.ll index 87718aee..0511d79d 100644 --- a/transform/testdata/allocs.ll +++ b/transform/testdata/allocs.ll @@ -54,6 +54,20 @@ define i32* @testEscapingReturn() { ret i32* %2 } +; Do a non-escaping allocation in a loop. +define void @testNonEscapingLoop() { +entry: + br label %loop +loop: + %0 = call i8* @runtime.alloc(i32 4) + %1 = bitcast i8* %0 to i32* + %2 = call i32* @noescapeIntPtr(i32* %1) + %3 = icmp eq i32* null, %2 + br i1 %3, label %loop, label %end +end: + ret void +} + declare i32* @escapeIntPtr(i32*) declare i32* @noescapeIntPtr(i32* nocapture) diff --git a/transform/testdata/allocs.out.ll b/transform/testdata/allocs.out.ll index eb92d6ed..e788beeb 100644 --- a/transform/testdata/allocs.out.ll +++ b/transform/testdata/allocs.out.ll @@ -50,6 +50,20 @@ define i32* @testEscapingReturn() { ret i32* %2 } +define void @testNonEscapingLoop() { +entry: + %stackalloc.alloca = alloca [1 x i32] + br label %loop +loop: + store [1 x i32] zeroinitializer, [1 x i32]* %stackalloc.alloca + %stackalloc = bitcast [1 x i32]* %stackalloc.alloca to i32* + %0 = call i32* @noescapeIntPtr(i32* %stackalloc) + %1 = icmp eq i32* null, %0 + br i1 %1, label %loop, label %end +end: + ret void +} + declare i32* @escapeIntPtr(i32*) declare i32* @noescapeIntPtr(i32* nocapture) diff --git a/transform/transform_test.go b/transform/transform_test.go index 12240ac5..d0f12569 100644 --- a/transform/transform_test.go +++ b/transform/transform_test.go @@ -70,8 +70,9 @@ func fuzzyEqualIR(s1, s2 string) bool { func filterIrrelevantIRLines(lines []string) []string { var out []string for _, line := range lines { - line = strings.TrimSpace(line) // drop '\r' on Windows - if line == "" || line[0] == ';' { + line = strings.Split(line, ";")[0] // strip out comments/info + line = strings.TrimRight(line, "\r ") // drop '\r' on Windows and remove trailing spaces from comments + if line == "" { continue } if strings.HasPrefix(line, "source_filename = ") {