compiler: remove intermediary copy of compiled functions code (#1466)

Signed-off-by: Achille Roussel <achille.roussel@gmail.com>
Co-authored-by: Crypt Keeper <64215+codefromthecrypt@users.noreply.github.com>

internal/engine/compiler/engine.go

@@ -520,12 +520,34 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, listeners
 		return e.addCompiledModule(module, cm, withGoFunc)
 	}
 
-	bodies := make([][]byte, localFuncs)
 	// As this uses mmap, we need to munmap on the compiled machine code when it's GCed.
 	e.setFinalizer(cm, releaseCompiledModule)
 	ln := len(listeners)
 	cmp := newCompiler()
 	asmNodes := new(asmNodes)
+
+	// The executable code is allocated in memory mappings of executableLength,
+	// and grown on demand when we exhaust the memory mapping capacity.
+	//
+	// The executableOffset variable tracks the position where the next function
+	// code will be written, and is always aligned on 16 bytes boundaries.
+	var executableOffset int
+	var executableLength int
+	var executable []byte
+
+	defer func() {
+		// At the end of the function, the executable is set on the compiled
+		// module and the local variable cleared; until then, the function owns
+		// the memory mapping and is reponsible for clearing it if it returns
+		// due to an error. Note that an error at this stage is not recoverable
+		// so we panic if we fail to unmap the memory segment.
+		if executable != nil {
+			if err := platform.MunmapCodeSegment(executable); err != nil {
+				panic(fmt.Errorf("compiler: failed to munmap code segment: %w", err))
+			}
+		}
+	}()
+
 	for i := range module.CodeSection {
 		typ := &module.TypeSection[module.FunctionSection[i]]
 		var lsn experimental.FunctionListener
@@ -534,6 +556,7 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, listeners
 		}
 		funcIndex := wasm.Index(i)
 		compiledFn := &cm.functions[i]
+
 		var body []byte
 		if codeSeg := &module.CodeSection[i]; codeSeg.GoFunc != nil {
 			cmp.Init(typ, nil, lsn != nil)
@@ -557,39 +580,38 @@ func (e *engine) CompileModule(_ context.Context, module *wasm.Module, listeners
 			}
 		}
 
-		// The `body` here is the view owned by assembler and will be overridden by the next iteration, so copy the body here.
-		bodyCopied := make([]byte, len(body))
-		copy(bodyCopied, body)
-		bodies[i] = bodyCopied
+		functionEndOffset := executableOffset + len(body)
+		if executableLength < functionEndOffset {
+			if executableLength == 0 {
+				executableLength = 65536
+			}
+			for executableLength < functionEndOffset {
+				executableLength *= 2
+			}
+			b, err := platform.RemapCodeSegment(executable, executableLength)
+			if err != nil {
+				return err
+			}
+			executable = b
+		}
+
+		compiledFn.executableOffset = executableOffset
 		compiledFn.listener = lsn
 		compiledFn.parent = cm
 		compiledFn.index = importedFuncs + funcIndex
-	}
 
-	var executableOffset int
-	for i, b := range bodies {
-		cm.functions[i].executableOffset = executableOffset
+		copy(executable[executableOffset:], body)
 		// Align 16-bytes boundary.
-		executableOffset = (executableOffset + len(b) + 15) &^ 15
-	}
-
-	executable, err := platform.MmapCodeSegment(executableOffset)
-	if err != nil {
-		return err
-	}
-
-	for i, b := range bodies {
-		offset := cm.functions[i].executableOffset
-		copy(executable[offset:], b)
+		executableOffset = (executableOffset + len(body) + 15) &^ 15
 	}
 
 	if runtime.GOARCH == "arm64" {
 		// On arm64, we cannot give all of rwx at the same time, so we change it to exec.
-		if err = platform.MprotectRX(executable); err != nil {
+		if err := platform.MprotectRX(executable); err != nil {
 			return err
 		}
 	}
-	cm.executable = executable
+	cm.executable, executable = executable, nil
 	return e.addCompiledModule(module, cm, withGoFunc)
 }
 

internal/platform/mremap_linux.go

@@ -0,0 +1,43 @@
+package platform
+
+import (
+	"syscall"
+	"unsafe"
+)
+
+const (
+	__MREMAP_MAYMOVE = 1
+)
+
+func remapCodeSegmentAMD64(code []byte, size int) ([]byte, error) {
+	return remapCodeSegment(code, size)
+}
+
+func remapCodeSegmentARM64(code []byte, size int) ([]byte, error) {
+	return remapCodeSegment(code, size)
+}
+
+func remapCodeSegment(code []byte, size int) ([]byte, error) {
+	p, err := mremap(*(*unsafe.Pointer)(unsafe.Pointer(&code)), len(code), size, __MREMAP_MAYMOVE)
+	if err != nil {
+		return nil, err
+	}
+	return unsafe.Slice((*byte)(p), size), nil
+}
+
+//go:nosplit
+func mremap(oldAddr unsafe.Pointer, oldSize, newSize, flags int) (unsafe.Pointer, error) {
+	p, _, err := syscall.Syscall6(
+		syscall.SYS_MREMAP,
+		uintptr(oldAddr),
+		uintptr(oldSize),
+		uintptr(newSize),
+		uintptr(flags),
+		uintptr(0),
+		uintptr(0),
+	)
+	if err != 0 {
+		return nil, err
+	}
+	return unsafe.Pointer(p), nil
+}

internal/platform/mremap_other.go

@@ -0,0 +1,23 @@
+//go:build !linux
+
+package platform
+
+func remapCodeSegmentAMD64(code []byte, size int) ([]byte, error) {
+	b, err := mmapCodeSegmentAMD64(size)
+	if err != nil {
+		return nil, err
+	}
+	copy(b, code)
+	munmapCodeSegment(code)
+	return b, nil
+}
+
+func remapCodeSegmentARM64(code []byte, size int) ([]byte, error) {
+	b, err := mmapCodeSegmentARM64(size)
+	if err != nil {
+		return nil, err
+	}
+	copy(b, code)
+	munmapCodeSegment(code)
+	return b, nil
+}

internal/platform/platform.go

@@ -5,7 +5,6 @@
 package platform
 
 import (
-	"errors"
 	"runtime"
 	"strings"
 )
@@ -34,7 +33,7 @@ func CompilerSupported() bool {
 // See https://man7.org/linux/man-pages/man2/mmap.2.html for mmap API and flags.
 func MmapCodeSegment(size int) ([]byte, error) {
 	if size == 0 {
-		panic(errors.New("BUG: MmapCodeSegment with zero length"))
+		panic("BUG: MmapCodeSegment with zero length")
 	}
 	if runtime.GOARCH == "amd64" {
 		return mmapCodeSegmentAMD64(size)
@@ -43,10 +42,32 @@ func MmapCodeSegment(size int) ([]byte, error) {
 	}
 }
 
+// RemapCodeSegment reallocates the memory mapping of an existing code segment
+// to increase its size. The previous code mapping is unmapped and must not be
+// reused after the function returns.
+//
+// This is similar to mremap(2) on linux, and emulated on platforms which do not
+// have this syscall.
+//
+// See https://man7.org/linux/man-pages/man2/mremap.2.html
+func RemapCodeSegment(code []byte, size int) ([]byte, error) {
+	if size < len(code) {
+		panic("BUG: RemapCodeSegment with size less than code")
+	}
+	if code == nil {
+		return MmapCodeSegment(size)
+	}
+	if runtime.GOARCH == "amd64" {
+		return remapCodeSegmentAMD64(code, size)
+	} else {
+		return remapCodeSegmentARM64(code, size)
+	}
+}
+
 // MunmapCodeSegment unmaps the given memory region.
 func MunmapCodeSegment(code []byte) error {
 	if len(code) == 0 {
-		panic(errors.New("BUG: MunmapCodeSegment with zero length"))
+		panic("BUG: MunmapCodeSegment with zero length")
 	}
 	return munmapCodeSegment(code)
 }