transform: track 0-index GEPs

It appears that LLVM is turning bitcasts into 0-index GEPs.
This caused stuff to not be tracked, resulting in use-after-free issues.
This solution is sub-optimal, but is the most reasonable solution I could come up with without redesigning the stack slots pass.

3 files changed, 61 insertions, 13 deletions

diff --git a/transform/gc.go b/transform/gc.go
index 27300bad9..932964452 100644
--- a/transform/gc.go
+++ b/transform/gc.go
@@ -151,10 +151,24 @@ func MakeGCStackSlots(mod llvm.Module) bool {
 			}
 			switch ptr.InstructionOpcode() {
 			case llvm.GetElementPtr:
-				// These values do not create new values: the values already
-				// existed locally in this function so must have been tracked
-				// already.
-				continue
+				// Check for all zero offsets.
+				// Sometimes LLVM rewrites bitcasts to zero-index GEPs, and we still need to track the GEP.
+				n := ptr.OperandsCount()
+				var hasOffset bool
+				for i := 1; i < n; i++ {
+					offset := ptr.Operand(i)
+					if offset.IsAConstantInt().IsNil() || offset.ZExtValue() != 0 {
+						hasOffset = true
+						break
+					}
+				}
+
+				if hasOffset {
+					// These values do not create new values: the values already
+					// existed locally in this function so must have been tracked
+					// already.
+					continue
+				}
 			case llvm.PHI:
 				// While the value may have already been tracked, it may be overwritten in a loop.
 				// Therefore, a second copy must be created to ensure that it is tracked over the entirety of its lifetime.
diff --git a/transform/testdata/gc-stackslots.ll b/transform/testdata/gc-stackslots.ll
index 10d10975d..deba7b128 100644
--- a/transform/testdata/gc-stackslots.ll
+++ b/transform/testdata/gc-stackslots.ll
@@ -84,3 +84,14 @@ loop:
 end:
   ret i8* %next.x
 }
+
+declare [32 x i8]* @arrayAlloc()
+
+define void @testGEPBitcast() {
+  %arr = call [32 x i8]* @arrayAlloc()
+  %arr.bitcast = getelementptr [32 x i8], [32 x i8]* %arr, i32 0, i32 0
+  call void @runtime.trackPointer(i8* %arr.bitcast)
+  %other = call i8* @runtime.alloc(i32 1)
+  call void @runtime.trackPointer(i8* %other)
+  ret void
+}
diff --git a/transform/testdata/gc-stackslots.out.ll b/transform/testdata/gc-stackslots.out.ll
index 88ee8c02a..efcc3ed25 100644
--- a/transform/testdata/gc-stackslots.out.ll
+++ b/transform/testdata/gc-stackslots.out.ll
@@ -30,24 +30,26 @@ define i8* @needsStackSlots() {
 }
 
 define i8* @needsStackSlots2() {
-  %gc.stackobject = alloca { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }
-  store { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* } { %runtime.stackChainObject* null, i32 4, i8* null, i8* null, i8* null, i8* null }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }* %gc.stackobject
+  %gc.stackobject = alloca { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }
+  store { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* } { %runtime.stackChainObject* null, i32 5, i8* null, i8* null, i8* null, i8* null, i8* null }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }* %gc.stackobject
   %1 = load %runtime.stackChainObject*, %runtime.stackChainObject** @runtime.stackChainStart
-  %2 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }* %gc.stackobject, i32 0, i32 0
+  %2 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }* %gc.stackobject, i32 0, i32 0
   store %runtime.stackChainObject* %1, %runtime.stackChainObject** %2
-  %3 = bitcast { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }* %gc.stackobject to %runtime.stackChainObject*
+  %3 = bitcast { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }* %gc.stackobject to %runtime.stackChainObject*
   store %runtime.stackChainObject* %3, %runtime.stackChainObject** @runtime.stackChainStart
   %ptr1 = call i8* @getPointer()
-  %4 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }* %gc.stackobject, i32 0, i32 4
+  %4 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }* %gc.stackobject, i32 0, i32 4
   store i8* %ptr1, i8** %4
-  %5 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }* %gc.stackobject, i32 0, i32 3
+  %5 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }* %gc.stackobject, i32 0, i32 3
   store i8* %ptr1, i8** %5
-  %6 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }* %gc.stackobject, i32 0, i32 2
+  %6 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }* %gc.stackobject, i32 0, i32 2
   store i8* %ptr1, i8** %6
   %ptr2 = getelementptr i8, i8* @someGlobal, i32 0
+  %7 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }* %gc.stackobject, i32 0, i32 5
+  store i8* %ptr2, i8** %7
   %unused = call i8* @runtime.alloc(i32 4)
-  %7 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8* }* %gc.stackobject, i32 0, i32 5
-  store i8* %unused, i8** %7
+  %8 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8*, i8*, i8*, i8* }* %gc.stackobject, i32 0, i32 6
+  store i8* %unused, i8** %8
   store %runtime.stackChainObject* %1, %runtime.stackChainObject** @runtime.stackChainStart
   ret i8* %ptr1
 }
@@ -112,3 +114,24 @@ end:                                              ; preds = %loop
   store %runtime.stackChainObject* %0, %runtime.stackChainObject** @runtime.stackChainStart
   ret i8* %next.x
 }
+
+declare [32 x i8]* @arrayAlloc()
+
+define void @testGEPBitcast() {
+  %gc.stackobject = alloca { %runtime.stackChainObject*, i32, i8*, i8* }
+  store { %runtime.stackChainObject*, i32, i8*, i8* } { %runtime.stackChainObject* null, i32 2, i8* null, i8* null }, { %runtime.stackChainObject*, i32, i8*, i8* }* %gc.stackobject
+  %1 = load %runtime.stackChainObject*, %runtime.stackChainObject** @runtime.stackChainStart
+  %2 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8* }* %gc.stackobject, i32 0, i32 0
+  store %runtime.stackChainObject* %1, %runtime.stackChainObject** %2
+  %3 = bitcast { %runtime.stackChainObject*, i32, i8*, i8* }* %gc.stackobject to %runtime.stackChainObject*
+  store %runtime.stackChainObject* %3, %runtime.stackChainObject** @runtime.stackChainStart
+  %arr = call [32 x i8]* @arrayAlloc()
+  %arr.bitcast = getelementptr [32 x i8], [32 x i8]* %arr, i32 0, i32 0
+  %4 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8* }* %gc.stackobject, i32 0, i32 2
+  store i8* %arr.bitcast, i8** %4
+  %other = call i8* @runtime.alloc(i32 1)
+  %5 = getelementptr { %runtime.stackChainObject*, i32, i8*, i8* }, { %runtime.stackChainObject*, i32, i8*, i8* }* %gc.stackobject, i32 0, i32 3
+  store i8* %other, i8** %5
+  store %runtime.stackChainObject* %1, %runtime.stackChainObject** @runtime.stackChainStart
+  ret void
+}