Update wazero and use custom allocator

go.mod

@@ -5,7 +5,7 @@ go 1.19
 require (
 	github.com/google/go-cmp v0.5.5
 	github.com/pganalyze/pg_query_go/v5 v5.1.0
-	github.com/tetratelabs/wazero v1.7.0
+	github.com/tetratelabs/wazero v1.7.1
 	google.golang.org/protobuf v1.31.0
 )
 

go.sum

@@ -3,8 +3,8 @@ github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/pganalyze/pg_query_go/v5 v5.1.0 h1:MlxQqHZnvA3cbRQYyIrjxEjzo560P6MyTgtlaf3pmXg=
 github.com/pganalyze/pg_query_go/v5 v5.1.0/go.mod h1:FsglvxidZsVN+Ltw3Ai6nTgPVcK2BPukH3jCDEqc1Ug=
-github.com/tetratelabs/wazero v1.7.0 h1:jg5qPydno59wqjpGrHph81lbtHzTrWzwwtD4cD88+hQ=
-github.com/tetratelabs/wazero v1.7.0/go.mod h1:ytl6Zuh20R/eROuyDaGPkp82O9C/DJfXAwJfQ3X6/7Y=
+github.com/tetratelabs/wazero v1.7.1 h1:QtSfd6KLc41DIMpDYlJdoMc6k7QTN246DM2+n2Y/Dx8=
+github.com/tetratelabs/wazero v1.7.1/go.mod h1:ytl6Zuh20R/eROuyDaGPkp82O9C/DJfXAwJfQ3X6/7Y=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=

internal/alloc/alloc.go

@@ -0,0 +1,97 @@
+//go:build unix
+
+// Mostly copied from https://github.com/ncruces/go-sqlite3/blob/main/internal/util/alloc.go#L12
+
+// MIT License
+//
+// Copyright (c) 2023 Nuno Cruces
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+package alloc
+
+import (
+	"math"
+	"syscall"
+
+	"github.com/tetratelabs/wazero/experimental"
+)
+
+func Allocator() experimental.MemoryAllocator {
+	return experimental.MemoryAllocatorFunc(mmappedAllocator)
+}
+
+func mmappedAllocator(cap, max uint64) experimental.LinearMemory {
+	// Round up to the page size.
+	rnd := uint64(syscall.Getpagesize() - 1)
+	max = (max + rnd) &^ rnd
+	cap = (cap + rnd) &^ rnd
+
+	if max > math.MaxInt {
+		// This ensures int(max) overflows to a negative value,
+		// and syscall.Mmap returns EINVAL.
+		max = math.MaxUint64
+	}
+	// Reserve max bytes of address space, to ensure we won't need to move it.
+	// A protected, private, anonymous mapping should not commit memory.
+	b, err := syscall.Mmap(-1, 0, int(max), syscall.PROT_NONE, syscall.MAP_PRIVATE|syscall.MAP_ANON)
+	if err != nil {
+		panic(err)
+	}
+	// Commit the initial cap bytes of memory.
+	err = syscall.Mprotect(b[:cap], syscall.PROT_READ|syscall.PROT_WRITE)
+	if err != nil {
+		_ = syscall.Munmap(b)
+		panic(err)
+	}
+	return &mmappedMemory{buf: b[:cap]}
+}
+
+// The slice covers the entire mmapped memory:
+//   - len(buf) is the already committed memory,
+//   - cap(buf) is the reserved address space.
+type mmappedMemory struct {
+	buf []byte
+}
+
+func (m *mmappedMemory) Reallocate(size uint64) []byte {
+	if com := uint64(len(m.buf)); com < size {
+		// Round up to the page size.
+		rnd := uint64(syscall.Getpagesize() - 1)
+		new := (size + rnd) &^ rnd
+
+		// Commit additional memory up to new bytes.
+		err := syscall.Mprotect(m.buf[com:new], syscall.PROT_READ|syscall.PROT_WRITE)
+		if err != nil {
+			panic(err)
+		}
+
+		// Update committed memory.
+		m.buf = m.buf[:new]
+	}
+	return m.buf[:size]
+}
+
+func (m *mmappedMemory) Free() {
+	err := syscall.Munmap(m.buf[:cap(m.buf)])
+	if err != nil {
+		panic(err)
+	}
+	m.buf = nil
+}

internal/alloc/alloc_notunix.go

@@ -0,0 +1,9 @@
+//go:build !unix
+
+package alloc
+
+import "github.com/tetratelabs/wazero/experimental"
+
+func Allocator() experimental.MemoryAllocator {
+	return nil
+}

internal/wasix_32v1/wasix.go

@@ -153,7 +153,7 @@ var stackCheckpointFn = api.GoModuleFunc(func(ctx context.Context, mod api.Modul
 	cstack := make([]byte, len(cstackView))
 	copy(cstack, cstackView)
 
-	sc := ctx.Value(experimental.SnapshotterKey{}).(experimental.Snapshotter)
+	sc := experimental.GetSnapshotter(ctx)
 	s := sc.Snapshot()
 
 	idx := len(d.checkpoints)
@@ -221,7 +221,7 @@ type wasixDataKey struct{}
 
 func BackgroundContext() context.Context {
 	ctx := context.Background()
-	ctx = context.WithValue(ctx, experimental.EnableSnapshotterKey{}, true)
+	ctx = experimental.WithSnapshotter(ctx)
 	ctx = context.WithValue(ctx, wasixDataKey{}, &wasixData{})
 	return ctx
 }

parser/parser_wazero.go

@@ -16,6 +16,7 @@ import (
 	"github.com/tetratelabs/wazero/experimental"
 	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
 
+	"github.com/wasilibs/go-pgquery/internal/alloc"
 	"github.com/wasilibs/go-pgquery/internal/wasix_32v1"
 	"github.com/wasilibs/go-pgquery/internal/wasm"
 )
@@ -28,6 +29,10 @@ var (
 // TODO(anuraaga): Use shared memory with child modules instead of fresh runtimes per call.
 func newRT() (wazero.Runtime, wazero.CompiledModule) {
 	ctx := context.Background()
+	if a := alloc.Allocator(); a != nil {
+		ctx = experimental.WithMemoryAllocator(ctx, a)
+	}
+
 	rt := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfig().
 		WithCompilationCache(wazero.NewCompilationCache()).
 		WithCoreFeatures(api.CoreFeaturesV2|experimental.CoreFeaturesThreads))