go-libwebp

lib.go (8961B)

---

 // Package webp provides a libwebp backend that runs the codec as a
 // WebAssembly module under wazero.
 //
 // libwebp is compiled once to wasm32-wasi (see tools/build-wasm.sh) and
 // embedded here, so a single artifact serves every GOOS/GOARCH combination
 // rather than one transpiled Go file per target.
 package webp
 
 import (
 	"context"
 	_ "embed"
 	"fmt"
 	"image"
 	"io"
 	"sync"
 
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
 )
 
 //go:embed libwebp.wasm
 var moduleBytes []byte
 
 var (
 	initOnce sync.Once
 	initErr  error
 	embedded *Compiled
 )
 
 // Compiled is a module compiled to machine code: the expensive, immutable,
 // shareable half. Instantiating it is cheap by comparison.
 type Compiled struct {
 	runtime  wazero.Runtime
 	compiled wazero.CompiledModule
 }
 
 // CompileBytes compiles an arbitrary libwebp wasm module. It exists so
 // modules built by different toolchains can be compared side by side; most
 // callers want the embedded one via Init.
 func CompileBytes(ctx context.Context, wasm []byte) (*Compiled, error) {
 	// The codec touches no host resources, so the compiler needs no
 	// filesystem or clock access. WASI is instantiated only because
 	// wasi-libc's abort path imports proc_exit.
 	rt := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfig())
 	if _, err := wasi_snapshot_preview1.Instantiate(ctx, rt); err != nil {
 		return nil, fmt.Errorf("instantiating wasi: %w", err)
 	}
 
 	// Emscripten's standalone output imports a single memory-growth callback
 	// that only matters to its JS glue. wasi-sdk builds import nothing here,
 	// so defining it unconditionally costs them nothing.
 	if _, err := rt.NewHostModuleBuilder("env").
 		NewFunctionBuilder().
 		WithFunc(func(context.Context, uint32) {}).
 		Export("emscripten_notify_memory_growth").
 		Instantiate(ctx); err != nil {
 		return nil, fmt.Errorf("instantiating env stub: %w", err)
 	}
 	cm, err := rt.CompileModule(ctx, wasm)
 	if err != nil {
 		return nil, fmt.Errorf("compiling module: %w", err)
 	}
 	return &Compiled{runtime: rt, compiled: cm}, nil
 }
 
 // Init compiles the embedded module. It is safe to call repeatedly; the
 // compile happens once. This is the expensive step (machine code generation),
 // as distinct from instantiation, which allocates a fresh linear memory.
 func Init() error {
 	initOnce.Do(func() {
 		embedded, initErr = CompileBytes(context.Background(), moduleBytes)
 	})
 	return initErr
 }
 
 // instance is a single module instantiation: one linear memory, one copy of
 // libwebp's globals, one shadow stack. Exactly one call may be in flight.
 type instance struct {
 	mod api.Module
 	mem api.Memory
 
 	malloc         api.Function
 	free           api.Function
 	encodeRGBA     api.Function
 	encodeLossless api.Function
 	decodeRGBA     api.Function
 	webpFree       api.Function
 }
 
 func newInstance(ctx context.Context) (*instance, error) {
 	if err := Init(); err != nil {
 		return nil, err
 	}
 	return embedded.newInstance(ctx)
 }
 
 func (c *Compiled) newInstance(ctx context.Context) (*instance, error) {
 	// Anonymous: not registered in the runtime's module namespace, so
 	// dropping it does not leak a reference.
 	mod, err := c.runtime.InstantiateModule(ctx, c.compiled,
 		wazero.NewModuleConfig().WithName(""))
 	if err != nil {
 		return nil, fmt.Errorf("instantiating module: %w", err)
 	}
 	i := &instance{
 		mod:            mod,
 		mem:            mod.Memory(),
 		malloc:         mod.ExportedFunction("malloc"),
 		free:           mod.ExportedFunction("free"),
 		encodeRGBA:     mod.ExportedFunction("WebPEncodeRGBA"),
 		encodeLossless: mod.ExportedFunction("WebPEncodeLosslessRGBA"),
 		decodeRGBA:     mod.ExportedFunction("WebPDecodeRGBA"),
 		webpFree:       mod.ExportedFunction("WebPFree"),
 	}
 	for name, fn := range map[string]api.Function{
 		"malloc": i.malloc, "free": i.free,
 		"WebPEncodeRGBA": i.encodeRGBA, "WebPEncodeLosslessRGBA": i.encodeLossless,
 		"WebPDecodeRGBA": i.decodeRGBA, "WebPFree": i.webpFree,
 	} {
 		if fn == nil {
 			mod.Close(ctx)
 			return nil, fmt.Errorf("module missing export: %s", name)
 		}
 	}
 	return i, nil
 }
 
 func (i *instance) close(ctx context.Context) { i.mod.Close(ctx) }
 
 // alloc reserves n bytes of guest memory and returns the offset.
 func (i *instance) alloc(ctx context.Context, n int) (uint32, error) {
 	res, err := i.malloc.Call(ctx, uint64(n))
 	if err != nil {
 		return 0, fmt.Errorf("guest malloc: %w", err)
 	}
 	if uint32(res[0]) == 0 {
 		return 0, fmt.Errorf("guest malloc: out of memory (%d bytes)", n)
 	}
 	return uint32(res[0]), nil
 }
 
 func (i *instance) readU32(off uint32) uint32 {
 	v, _ := i.mem.ReadUint32Le(off)
 	return v
 }
 
 func (i *instance) encode(ctx context.Context, m *image.NRGBA, q float32) ([]byte, error) {
 	b := m.Bounds()
 	// Guest pointers are 32-bit; an out-param needs 4 bytes of scratch.
 	in, err := i.alloc(ctx, len(m.Pix))
 	if err != nil {
 		return nil, err
 	}
 	defer i.free.Call(ctx, uint64(in))
 
 	outPtr, err := i.alloc(ctx, 4)
 	if err != nil {
 		return nil, err
 	}
 	defer i.free.Call(ctx, uint64(outPtr))
 
 	if !i.mem.Write(in, m.Pix) {
 		return nil, fmt.Errorf("writing pixels into guest memory")
 	}
 
 	var res []uint64
 	if q >= 1.0 {
 		res, err = i.encodeLossless.Call(ctx,
 			uint64(in), uint64(b.Dx()), uint64(b.Dy()), uint64(m.Stride), uint64(outPtr))
 	} else {
 		res, err = i.encodeRGBA.Call(ctx,
 			uint64(in), uint64(b.Dx()), uint64(b.Dy()), uint64(m.Stride),
 			uint64(api.EncodeF32(q*100)), uint64(outPtr))
 	}
 	if err != nil {
 		return nil, fmt.Errorf("WebPEncode: %w", err)
 	}
 
 	size := uint32(res[0])
 	if size == 0 {
 		return nil, fmt.Errorf("empty result")
 	}
 	out := i.readU32(outPtr)
 	if out == 0 {
 		return nil, fmt.Errorf("failed to allocate output buffer")
 	}
 	defer i.webpFree.Call(ctx, uint64(out))
 
 	buf, ok := i.mem.Read(out, size)
 	if !ok {
 		return nil, fmt.Errorf("reading encoded data out of guest memory")
 	}
 	// mem.Read aliases the guest memory; copy before it is reused or grown.
 	cp := make([]byte, size)
 	copy(cp, buf)
 	return cp, nil
 }
 
 func (i *instance) decode(ctx context.Context, data []byte) (image.Image, error) {
 	in, err := i.alloc(ctx, len(data))
 	if err != nil {
 		return nil, err
 	}
 	defer i.free.Call(ctx, uint64(in))
 
 	dims, err := i.alloc(ctx, 8) // int width, int height
 	if err != nil {
 		return nil, err
 	}
 	defer i.free.Call(ctx, uint64(dims))
 
 	if !i.mem.Write(in, data) {
 		return nil, fmt.Errorf("writing webp data into guest memory")
 	}
 
 	res, err := i.decodeRGBA.Call(ctx,
 		uint64(in), uint64(len(data)), uint64(dims), uint64(dims+4))
 	if err != nil {
 		return nil, fmt.Errorf("WebPDecodeRGBA: %w", err)
 	}
 
 	samples := uint32(res[0])
 	if samples == 0 {
 		return nil, fmt.Errorf("failed decoding webp into rgba")
 	}
 	defer i.webpFree.Call(ctx, uint64(samples))
 
 	w, h := int(i.readU32(dims)), int(i.readU32(dims+4))
 	raw, ok := i.mem.Read(samples, uint32(w*h*4))
 	if !ok {
 		return nil, fmt.Errorf("reading pixels out of guest memory")
 	}
 	pix := make([]uint8, w*h*4)
 	copy(pix, raw)
 
 	return &image.NRGBA{
 		Pix:    pix,
 		Rect:   image.Rectangle{Max: image.Point{X: w, Y: h}},
 		Stride: w * 4,
 	}, nil
 }
 
 // EncodeImpl encodes m, instantiating a fresh module for the call.
 func EncodeImpl(w io.Writer, m *image.NRGBA, quality float32) error {
 	ctx := context.Background()
 	i, err := newInstance(ctx)
 	if err != nil {
 		return err
 	}
 	defer i.close(ctx)
 	buf, err := i.encode(ctx, m, quality)
 	if err != nil {
 		return err
 	}
 	_, err = w.Write(buf)
 	return err
 }
 
 // DecodeImpl decodes buf, instantiating a fresh module for the call.
 func DecodeImpl(buf []byte) (image.Image, error) {
 	ctx := context.Background()
 	i, err := newInstance(ctx)
 	if err != nil {
 		return nil, err
 	}
 	defer i.close(ctx)
 	return i.decode(ctx, buf)
 }
 
 // Instance is a reusable module instantiation. Exactly one call may be in
 // flight on it at a time; callers needing concurrency should keep a pool.
 // It exists so the cost of instantiation can be separated from the cost of
 // the codec itself.
 type Instance struct{ inner *instance }
 
 // NewInstance instantiates this compiled module.
 func (c *Compiled) NewInstance(ctx context.Context) (*Instance, error) {
 	i, err := c.newInstance(ctx)
 	if err != nil {
 		return nil, err
 	}
 	return &Instance{inner: i}, nil
 }
 
 func NewInstance(ctx context.Context) (*Instance, error) {
 	i, err := newInstance(ctx)
 	if err != nil {
 		return nil, err
 	}
 	return &Instance{inner: i}, nil
 }
 
 func (i *Instance) Encode(ctx context.Context, m *image.NRGBA, q float32) ([]byte, error) {
 	return i.inner.encode(ctx, m, q)
 }
 
 func (i *Instance) Decode(ctx context.Context, data []byte) (image.Image, error) {
 	return i.inner.decode(ctx, data)
 }
 
 // MemorySize reports the instance's current linear memory in bytes. It only
 // ever grows, which is what makes pooled instances worth evicting.
 func (i *Instance) MemorySize() uint32 { return i.inner.mem.Size() }
 
 func (i *Instance) Close(ctx context.Context) { i.inner.close(ctx) }