prepared.go

package spine

import (
	"context"
	"errors"
	"fmt"

	"github.com/tetratelabs/wazero"
	"github.com/tetratelabs/wazero/api"
)

// NewPrepared captures only our Guest's clean init, before any user code or tool.
// Each Run still has private allocated memory: this is full-copy, NOT page COW.
// This is not a generic Wasm checkpoint (Host resources/globals/tables aren't captured).
func NewPrepared(ctx context.Context, wasm []byte, tools map[string]Tool, earlyReads ...string) (*Runner, error) {
	r, err := New(ctx, wasm, tools, earlyReads...)
	if err != nil {
		return nil, err
	}
	m, err := r.newGuest(ctx, &boundedText{}, &boundedText{})
	if err != nil {
		r.Close(context.Background())
		return nil, err
	}
	defer m.Close(context.Background())
	memory := m.Memory()
	data, ok := memory.Read(0, memory.Size())
	if !ok || len(data) == 0 {
		r.Close(context.Background())
		return nil, errors.New("cannot capture initialized Guest memory")
	}
	r.image = append([]byte(nil), data...)
	return r, nil
}

// Both constructors use this lifecycle. Only CPython init vs memory copy differs.
func (r *Runner) newGuest(ctx context.Context, stdout, stderr *boundedText) (api.Module, error) {
	m, err := r.runtime.InstantiateModule(ctx, r.code, wazero.NewModuleConfig().WithName("").WithStartFunctions().WithStdout(stdout).WithStderr(stderr))
	if err != nil {
		return nil, err
	}
	failed := true
	defer func() {
		if failed {
			m.Close(context.Background())
		}
	}()
	if _, err = m.ExportedFunction("_initialize").Call(ctx); err != nil {
		return nil, fmt.Errorf("initialize Guest: %w%s", err, stderr.String())
	}
	if r.image == nil {
		status, err := m.ExportedFunction("init").Call(ctx)
		if err != nil {
			return nil, fmt.Errorf("init CPython: %w%s", err, stderr.String())
		}
		if status[0] != 0 {
			return nil, fmt.Errorf("CPython initialization failed: %s", stderr.String())
		}
	} else {
		// Instantiate/_initialize write data segments; restore the image AFTER them.
		memory := m.Memory()
		size := memory.Size()
		if int(size) > len(r.image) {
			return nil, errors.New("prepared memory shape changed")
		}
		if int(size) < len(r.image) {
			if _, ok := memory.Grow(uint32(len(r.image)-int(size)) / 65536); !ok {
				return nil, errors.New("cannot grow memory to prepared image")
			}
		}
		if !memory.Write(0, r.image) {
			return nil, errors.New("cannot copy prepared memory")
		}
	}
	failed = false
	return m, nil
}