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.

cgo/security.go

@@ -142,7 +142,6 @@ var validLinkerFlags = []*regexp.Regexp{
 	re(`-L([^@\-].*)`),
 	re(`-O`),
 	re(`-O([^@\-].*)`),
-	re(`--export=(.+)`), // for wasm-ld
 	re(`-f(no-)?(pic|PIC|pie|PIE)`),
 	re(`-f(no-)?openmp(-simd)?`),
 	re(`-fsanitize=([^@\-].*)`),

compiler/compiler.go

@@ -1059,12 +1059,11 @@ func (b *builder) createFunctionStart(intrinsic bool) {
 	if b.info.section != "" {
 		b.llvmFn.SetSection(b.info.section)
 	}
-	if b.info.exported && b.info.module != "" && strings.HasPrefix(b.Triple, "wasm") {
+	if b.info.exported && strings.HasPrefix(b.Triple, "wasm") {
 		// Set the exported name. This is necessary for WebAssembly because
 		// otherwise the function is not exported.
-		b.llvmFn.AddFunctionAttr(b.ctx.CreateStringAttribute("wasm-export-name", b.info.linkName))
-		// Set the export module.
-		b.llvmFn.AddFunctionAttr(b.ctx.CreateStringAttribute("wasm-export-module", b.info.module))
+		functionAttr := b.ctx.CreateStringAttribute("wasm-export-name", b.info.linkName)
+		b.llvmFn.AddFunctionAttr(functionAttr)
 	}
 
 	// Some functions have a pragma controlling the inlining level.

compiler/symbol.go

@@ -210,9 +210,8 @@ func (c *compilerContext) getFunction(fn *ssa.Function) llvm.Value {
 // exported.
 func (c *compilerContext) getFunctionInfo(f *ssa.Function) functionInfo {
 	info := functionInfo{
-		module:     "env",
-		importName: f.Name(),
-		linkName:   f.RelString(nil), // pick the default linkName
+		// Pick the default linkName.
+		linkName: f.RelString(nil),
 	}
 	// Check for //go: pragmas, which may change the link name (among others).
 	info.parsePragmas(f)
@@ -226,6 +225,10 @@ func (info *functionInfo) parsePragmas(f *ssa.Function) {
 		return
 	}
 	if decl, ok := f.Syntax().(*ast.FuncDecl); ok && decl.Doc != nil {
+
+		// Our importName for a wasm module (if we are compiling to wasm), or llvm link name
+		var importName string
+
 		for _, comment := range decl.Doc.List {
 			text := comment.Text
 			if strings.HasPrefix(text, "//export ") {
@@ -243,8 +246,7 @@ func (info *functionInfo) parsePragmas(f *ssa.Function) {
 					continue
 				}
 
-				info.importName = parts[1]
-				info.linkName = parts[1]
+				importName = parts[1]
 				info.exported = true
 			case "//go:interrupt":
 				if hasUnsafeImport(f.Pkg.Pkg) {
@@ -252,13 +254,10 @@ func (info *functionInfo) parsePragmas(f *ssa.Function) {
 				}
 			case "//go:wasm-module":
 				// Alternative comment for setting the import module.
-				if len(parts) == 1 {
-					// Function must not be exported outside of the WebAssembly
-					// module (but only be made available for linking).
-					info.module = ""
-				} else if len(parts) == 2 {
-					info.module = parts[1]
+				if len(parts) != 2 {
+					continue
 				}
+				info.module = parts[1]
 			case "//go:inline":
 				info.inline = inlineHint
 			case "//go:noinline":
@@ -298,6 +297,17 @@ func (info *functionInfo) parsePragmas(f *ssa.Function) {
 				}
 			}
 		}
+
+		// Set the importName for our exported function if we have one
+		if importName != "" {
+			if info.module == "" {
+				info.linkName = importName
+			} else {
+				// WebAssembly import
+				info.importName = importName
+			}
+		}
+
 	}
 }
 

compiler/testdata/pragma.ll

@@ -62,7 +62,7 @@ declare void @main.undefinedFunctionNotInSection(i8*) #0
 
 attributes #0 = { "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" }
 attributes #1 = { nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" }
-attributes #2 = { nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" "wasm-export-module"="env" "wasm-export-name"="extern_func" "wasm-import-module"="env" "wasm-import-name"="extern_func" }
+attributes #2 = { nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" "wasm-export-name"="extern_func" "wasm-import-module"="env" "wasm-import-name"="extern_func" }
 attributes #3 = { inlinehint nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" }
 attributes #4 = { noinline nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" }
-attributes #5 = { nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" "wasm-export-module"="env" "wasm-export-name"="exportedFunctionInSection" "wasm-import-module"="env" "wasm-import-name"="exportedFunctionInSection" }
+attributes #5 = { nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" "wasm-export-name"="exportedFunctionInSection" "wasm-import-module"="env" "wasm-import-name"="exportedFunctionInSection" }

src/runtime/arch_tinygowasm.go

@@ -62,57 +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
-//go:wasm-module
 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
-//go:wasm-module
 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
-//go:wasm-module
 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
-//go:wasm-module
 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)
 		})
 	}
 }