wasm: improve malloc/free heap tracking

Using a slice requires a lot less in code size than a map. This is
visible when compiling a very small "hello world" style program.

Before tracking memory in malloc/free:  2873 bytes
With tracking using a map:              6551 bytes
With a slice instead of a map:          3532 bytes

Of course, most of this code size increase won't be visible with
https://github.com/tinygo-org/tinygo/pull/3142, but it's still a saving
of around 3kB in this minimal example.

src/runtime/arch_tinygowasm.go

@@ -62,53 +62,63 @@ func growHeap() bool {
 }
 
 // The below functions override the default allocator of wasi-libc. This ensures
-// code linked from other languages can allocate memory without colliding with
-// our GC allocations.
+// code linked from other languages can allocate memory on the same heap as the
+// TinyGo heap.
 
-var allocs = make(map[uintptr][]byte)
+// Keep track of all the heap allocations while they're in use.
+// This is not the most efficient solution but it costs a lot less in code size
+// compared to a map.
+var allocs []unsafe.Pointer
+
+func trackAlloc(ptr unsafe.Pointer) {
+	// Try to find some empty space in the allocs slice.
+	for i, slot := range allocs {
+		if slot == nil {
+			allocs[i] = slot
+			return
+		}
+	}
+	// Couldn't find this space. Fall back to appending to the end.
+	allocs = append(allocs, ptr)
+}
+
+func removeAlloc(ptr unsafe.Pointer) {
+	// Remove the pointer so it can be garbage collected.
+	for i, slot := range allocs {
+		if ptr == slot {
+			allocs[i] = nil
+			return
+		}
+	}
+}
 
 //export malloc
 func libc_malloc(size uintptr) unsafe.Pointer {
-	buf := make([]byte, size)
-	ptr := unsafe.Pointer(&buf[0])
-	allocs[uintptr(ptr)] = buf
+	ptr := alloc(size, nil)
+	trackAlloc(ptr)
 	return ptr
 }
 
 //export free
 func libc_free(ptr unsafe.Pointer) {
-	if ptr == nil {
-		return
-	}
-	if _, ok := allocs[uintptr(ptr)]; ok {
-		delete(allocs, uintptr(ptr))
-	} else {
-		panic("free: invalid pointer")
-	}
+	removeAlloc(ptr)
+	free(ptr)
 }
 
 //export calloc
 func libc_calloc(nmemb, size uintptr) unsafe.Pointer {
-	// No difference between calloc and malloc.
+	// Note: we could be even more correct here and check that nmemb * size
+	// doesn't overflow. However the current implementation should normally work
+	// fine.
 	return libc_malloc(nmemb * size)
 }
 
 //export realloc
 func libc_realloc(oldPtr unsafe.Pointer, size uintptr) unsafe.Pointer {
-	// It's hard to optimize this to expand the current buffer with our GC, but
-	// it is theoretically possible. For now, just always allocate fresh.
-	buf := make([]byte, size)
-
-	if oldPtr != nil {
-		if oldBuf, ok := allocs[uintptr(oldPtr)]; ok {
-			copy(buf, oldBuf)
-			delete(allocs, uintptr(oldPtr))
-		} else {
-			panic("realloc: invalid pointer")
-		}
+	newPtr := realloc(oldPtr, size)
+	if newPtr != oldPtr {
+		removeAlloc(oldPtr)
+		trackAlloc(newPtr)
 	}
-
-	ptr := unsafe.Pointer(&buf[0])
-	allocs[uintptr(ptr)] = buf
-	return ptr
+	return newPtr
 }

tests/runtime_wasi/malloc_test.go

@@ -23,13 +23,13 @@ func libc_calloc(nmemb, size uintptr) unsafe.Pointer
 //export realloc
 func libc_realloc(ptr unsafe.Pointer, size uintptr) unsafe.Pointer
 
-func getFilledBuffer_malloc() uintptr {
+func getFilledBuffer_malloc() unsafe.Pointer {
 	ptr := libc_malloc(5)
 	fillPanda(ptr)
-	return uintptr(ptr)
+	return ptr
 }
 
-func getFilledBuffer_calloc() uintptr {
+func getFilledBuffer_calloc() unsafe.Pointer {
 	ptr := libc_calloc(2, 5)
 	fillPanda(ptr)
 	*(*byte)(unsafe.Add(ptr, 5)) = 'b'
@@ -37,23 +37,23 @@ func getFilledBuffer_calloc() uintptr {
 	*(*byte)(unsafe.Add(ptr, 7)) = 'a'
 	*(*byte)(unsafe.Add(ptr, 8)) = 'r'
 	*(*byte)(unsafe.Add(ptr, 9)) = 's'
-	return uintptr(ptr)
+	return ptr
 }
 
-func getFilledBuffer_realloc() uintptr {
+func getFilledBuffer_realloc() unsafe.Pointer {
 	origPtr := getFilledBuffer_malloc()
-	ptr := libc_realloc(unsafe.Pointer(origPtr), 9)
+	ptr := libc_realloc(origPtr, 9)
 	*(*byte)(unsafe.Add(ptr, 5)) = 'b'
 	*(*byte)(unsafe.Add(ptr, 6)) = 'e'
 	*(*byte)(unsafe.Add(ptr, 7)) = 'a'
 	*(*byte)(unsafe.Add(ptr, 8)) = 'r'
-	return uintptr(ptr)
+	return ptr
 }
 
-func getFilledBuffer_reallocNil() uintptr {
+func getFilledBuffer_reallocNil() unsafe.Pointer {
 	ptr := libc_realloc(nil, 5)
 	fillPanda(ptr)
-	return uintptr(ptr)
+	return ptr
 }
 
 func fillPanda(ptr unsafe.Pointer) {
@@ -64,10 +64,10 @@ func fillPanda(ptr unsafe.Pointer) {
 	*(*byte)(unsafe.Add(ptr, 4)) = 'a'
 }
 
-func checkFilledBuffer(t *testing.T, ptr uintptr, content string) {
+func checkFilledBuffer(t *testing.T, ptr unsafe.Pointer, content string) {
 	t.Helper()
 	buf := *(*string)(unsafe.Pointer(&reflect.StringHeader{
-		Data: ptr,
+		Data: uintptr(ptr),
 		Len:  uintptr(len(content)),
 	}))
 	if buf != content {
@@ -78,7 +78,7 @@ func checkFilledBuffer(t *testing.T, ptr uintptr, content string) {
 func TestMallocFree(t *testing.T) {
 	tests := []struct {
 		name      string
-		getBuffer func() uintptr
+		getBuffer func() unsafe.Pointer
 		content   string
 	}{
 		{
@@ -121,7 +121,7 @@ func TestMallocFree(t *testing.T) {
 
 			checkFilledBuffer(t, bufPtr, tt.content)
 
-			libc_free(unsafe.Pointer(bufPtr))
+			libc_free(bufPtr)
 		})
 	}
 }