runtime.go

Documentation: github.com/tetratelabs/wazero/internal/integration_test/vs

     1  package vs
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  
     8  	"github.com/tetratelabs/wazero"
     9  	"github.com/tetratelabs/wazero/api"
    10  	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
    11  	"github.com/tetratelabs/wazero/internal/wasm"
    12  )
    13  
    14  type RuntimeConfig struct {
    15  	Name              string
    16  	ModuleName        string
    17  	ModuleWasm        []byte
    18  	FuncNames         []string
    19  	NeedsWASI         bool
    20  	NeedsMemoryExport bool
    21  	// LogFn requires the implementation to export a function "env.log" which accepts i32i32_v.
    22  	// The implementation invoke this with a byte slice allocated from the offset, length pair.
    23  	// This function simulates a host function that logs a message.
    24  	LogFn func([]byte) error
    25  	// EnvFReturnValue is set to non-zero if we want the runtime to instantiate "env" module with the function "f"
    26  	// which accepts one i64 value and returns the EnvFReturnValue as i64. This is mutually exclusive to LogFn.
    27  	EnvFReturnValue uint64
    28  }
    29  
    30  type Runtime interface {
    31  	Name() string
    32  	Compile(context.Context, *RuntimeConfig) error
    33  	Instantiate(context.Context, *RuntimeConfig) (Module, error)
    34  	Close(context.Context) error
    35  }
    36  
    37  type Module interface {
    38  	CallI32_I32(ctx context.Context, funcName string, param uint32) (uint32, error)
    39  	CallI32I32_V(ctx context.Context, funcName string, x, y uint32) error
    40  	CallI32_V(ctx context.Context, funcName string, param uint32) error
    41  	CallV_V(ctx context.Context, funcName string) error
    42  	CallI64_I64(ctx context.Context, funcName string, param uint64) (uint64, error)
    43  	WriteMemory(offset uint32, bytes []byte) error
    44  	Memory() []byte
    45  	Close(context.Context) error
    46  }
    47  
    48  func NewWazeroInterpreterRuntime() Runtime {
    49  	return newWazeroRuntime("wazero-interpreter", wazero.NewRuntimeConfigInterpreter())
    50  }
    51  
    52  func NewWazeroCompilerRuntime() Runtime {
    53  	return newWazeroRuntime(compilerRuntime, wazero.NewRuntimeConfigCompiler())
    54  }
    55  
    56  func newWazeroRuntime(name string, config wazero.RuntimeConfig) *wazeroRuntime {
    57  	return &wazeroRuntime{name: name, config: config}
    58  }
    59  
    60  type wazeroRuntime struct {
    61  	name          string
    62  	config        wazero.RuntimeConfig
    63  	runtime       wazero.Runtime
    64  	logFn         func([]byte) error
    65  	env, compiled wazero.CompiledModule
    66  }
    67  
    68  type wazeroModule struct {
    69  	wasi     api.Closer
    70  	env, mod api.Module
    71  	funcs    map[string]api.Function
    72  }
    73  
    74  func (r *wazeroRuntime) Name() string {
    75  	return r.name
    76  }
    77  
    78  func (m *wazeroModule) Memory() []byte {
    79  	return m.mod.Memory().(*wasm.MemoryInstance).Buffer
    80  }
    81  
    82  func (r *wazeroRuntime) log(_ context.Context, mod api.Module, stack []uint64) {
    83  	offset, byteCount := uint32(stack[0]), uint32(stack[1])
    84  
    85  	buf, ok := mod.Memory().Read(offset, byteCount)
    86  	if !ok {
    87  		panic("out of memory reading log buffer")
    88  	}
    89  	if err := r.logFn(buf); err != nil {
    90  		panic(err)
    91  	}
    92  }
    93  
    94  func (r *wazeroRuntime) Compile(ctx context.Context, cfg *RuntimeConfig) (err error) {
    95  	r.runtime = wazero.NewRuntimeWithConfig(ctx, r.config)
    96  	if cfg.LogFn != nil {
    97  		r.logFn = cfg.LogFn
    98  		if r.env, err = r.runtime.NewHostModuleBuilder("env").
    99  			NewFunctionBuilder().
   100  			WithGoModuleFunction(api.GoModuleFunc(r.log), []api.ValueType{api.ValueTypeI32, api.ValueTypeI32}, []api.ValueType{}).
   101  			Export("log").
   102  			Compile(ctx); err != nil {
   103  			return err
   104  		}
   105  	} else if cfg.EnvFReturnValue != 0 {
   106  		if r.env, err = r.runtime.NewHostModuleBuilder("env").
   107  			NewFunctionBuilder().
   108  			WithGoFunction(api.GoFunc(func(ctx context.Context, stack []uint64) {
   109  				stack[0] = cfg.EnvFReturnValue
   110  			}), []api.ValueType{api.ValueTypeI64}, []api.ValueType{api.ValueTypeI64}).
   111  			Export("f").
   112  			Compile(ctx); err != nil {
   113  			return err
   114  		}
   115  	}
   116  	r.compiled, err = r.runtime.CompileModule(ctx, cfg.ModuleWasm)
   117  	return
   118  }
   119  
   120  func (r *wazeroRuntime) Instantiate(ctx context.Context, cfg *RuntimeConfig) (mod Module, err error) {
   121  	wazeroCfg := wazero.NewModuleConfig().WithName(cfg.ModuleName)
   122  	m := &wazeroModule{funcs: map[string]api.Function{}}
   123  
   124  	// Instantiate WASI, if configured.
   125  	if cfg.NeedsWASI {
   126  		if m.wasi, err = wasi_snapshot_preview1.Instantiate(ctx, r.runtime); err != nil {
   127  			return
   128  		}
   129  	}
   130  
   131  	// Instantiate the host module, "env", if configured.
   132  	if env := r.env; env != nil {
   133  		if m.env, err = r.runtime.InstantiateModule(ctx, env, wazero.NewModuleConfig()); err != nil {
   134  			return
   135  		}
   136  	}
   137  
   138  	// Instantiate the module.
   139  	if m.mod, err = r.runtime.InstantiateModule(ctx, r.compiled, wazeroCfg); err != nil {
   140  		return
   141  	}
   142  
   143  	// Ensure function exports exist.
   144  	for _, funcName := range cfg.FuncNames {
   145  		if fn := m.mod.ExportedFunction(funcName); fn == nil {
   146  			return nil, fmt.Errorf("%s is not an exported function", funcName)
   147  		} else {
   148  			m.funcs[funcName] = fn
   149  		}
   150  	}
   151  	mod = m
   152  	return
   153  }
   154  
   155  func (r *wazeroRuntime) Close(ctx context.Context) (err error) {
   156  	if compiled := r.compiled; compiled != nil {
   157  		err = compiled.Close(ctx)
   158  	}
   159  	r.compiled = nil
   160  	if env := r.env; env != nil {
   161  		err = env.Close(ctx)
   162  	}
   163  	r.env = nil
   164  	return
   165  }
   166  
   167  func (m *wazeroModule) CallV_V(ctx context.Context, funcName string) (err error) {
   168  	_, err = m.funcs[funcName].Call(ctx)
   169  	return
   170  }
   171  
   172  func (m *wazeroModule) CallI32_I32(ctx context.Context, funcName string, param uint32) (uint32, error) {
   173  	if results, err := m.funcs[funcName].Call(ctx, uint64(param)); err != nil {
   174  		return 0, err
   175  	} else if len(results) > 0 {
   176  		return uint32(results[0]), nil
   177  	}
   178  	return 0, nil
   179  }
   180  
   181  func (m *wazeroModule) CallI32I32_V(ctx context.Context, funcName string, x, y uint32) (err error) {
   182  	_, err = m.funcs[funcName].Call(ctx, uint64(x), uint64(y))
   183  	return
   184  }
   185  
   186  func (m *wazeroModule) CallI32_V(ctx context.Context, funcName string, param uint32) (err error) {
   187  	_, err = m.funcs[funcName].Call(ctx, uint64(param))
   188  	return
   189  }
   190  
   191  func (m *wazeroModule) CallI64_I64(ctx context.Context, funcName string, param uint64) (uint64, error) {
   192  	if results, err := m.funcs[funcName].Call(ctx, param); err != nil {
   193  		return 0, err
   194  	} else {
   195  		return results[0], nil
   196  	}
   197  }
   198  
   199  func (m *wazeroModule) WriteMemory(offset uint32, bytes []byte) error {
   200  	if !m.mod.Memory().Write(offset, bytes) {
   201  		return errors.New("out of memory writing name")
   202  	}
   203  	return nil
   204  }
   205  
   206  func (m *wazeroModule) Close(ctx context.Context) (err error) {
   207  	if mod := m.mod; mod != nil {
   208  		err = mod.Close(ctx)
   209  	}
   210  	m.mod = nil
   211  	if env := m.env; env != nil {
   212  		err = env.Close(ctx)
   213  	}
   214  	m.env = nil
   215  	if wasi := m.wasi; wasi != nil {
   216  		err = wasi.Close(ctx)
   217  	}
   218  	m.wasi = nil
   219  	return
   220  }
   221
View as plain text