perf: optimized host module compilation (#538)

# Description

optimized host module compilation

cli/serverless/wasm/wasmedge.go

@@ -70,6 +70,11 @@ func (r *wasmEdgeRuntime) Init(wasmFile string) error {
 		},
 		[]wasmedge.ValType{wasmedge.ValType_I32}), r.contextData, nil, 0)
 	r.module.AddFunction(WasmFuncContextData, contextDataFunc)
+	// context data size
+	contextDataSizeFunc := wasmedge.NewFunction(wasmedge.NewFunctionType(
+		[]wasmedge.ValType{},
+		[]wasmedge.ValType{wasmedge.ValType_I32}), r.contextDataSize, nil, 0)
+	r.module.AddFunction(WasmFuncContextDataSize, contextDataSizeFunc)
 
 	err := r.vm.RegisterModule(r.module)
 	if err != nil {
@@ -159,6 +164,11 @@ func (r *wasmEdgeRuntime) contextData(_ any, callframe *wasmedge.CallingFrame, p
 	return []any{dataLen}, wasmedge.Result_Success
 }
 
+func (r *wasmEdgeRuntime) contextDataSize(_ any, callframe *wasmedge.CallingFrame, params []any) ([]any, wasmedge.Result) {
+	dataLen := len(r.serverlessCtx.Data())
+	return []any{dataLen}, wasmedge.Result_Success
+}
+
 func (r *wasmEdgeRuntime) write(_ any, callframe *wasmedge.CallingFrame, params []any) ([]any, wasmedge.Result) {
 	tag := params[0].(int32)
 	pointer := params[1].(int32)

cli/serverless/wasm/wasmtime.go

@@ -66,6 +66,10 @@ func (r *wasmtimeRuntime) Init(wasmFile string) error {
 	if err := r.linker.FuncWrap("env", WasmFuncContextData, r.contextData); err != nil {
 		return fmt.Errorf("linker.FuncWrap: %s %v", WasmFuncContextData, err)
 	}
+	// context data size
+	if err := r.linker.FuncWrap("env", WasmFuncContextDataSize, r.contextDataSize); err != nil {
+		return fmt.Errorf("linker.FuncWrap: %s %v", WasmFuncContextDataSize, err)
+	}
 	// write
 	if err := r.linker.FuncWrap("env", WasmFuncWrite, r.write); err != nil {
 		return fmt.Errorf("linker.FuncWrap: %s %v", WasmFuncWrite, err)
@@ -135,6 +139,10 @@ func (r *wasmtimeRuntime) contextData(pointer int32, limit int32) (dataLen int32
 	return
 }
 
+func (r *wasmtimeRuntime) contextDataSize() int32 {
+	return int32(len(r.serverlessCtx.Data()))
+}
+
 func (r *wasmtimeRuntime) write(tag int32, pointer int32, length int32) int32 {
 	output := r.memory.UnsafeData(r.store)[pointer : pointer+length]
 	if len(output) == 0 {

cli/serverless/wasm/wazero.go

@@ -14,6 +14,8 @@ import (
 	"github.com/yomorun/yomo/serverless"
 )
 
+const i32 = api.ValueTypeI32
+
 type wazeroRuntime struct {
 	wazero.Runtime
 	conf   wazero.ModuleConfig
@@ -58,23 +60,23 @@ func (r *wazeroRuntime) Init(wasmFile string) error {
 	_, err = builder.
 		// observeDataTag
 		NewFunctionBuilder().
-		WithFunc(r.observeDataTag).
+		WithGoFunction(api.GoFunc(r.observeDataTag), []api.ValueType{i32}, []api.ValueType{}).
 		Export(WasmFuncObserveDataTag).
 		// write
 		NewFunctionBuilder().
-		WithFunc(r.write).
+		WithGoModuleFunction(api.GoModuleFunc(r.write), []api.ValueType{i32, i32, i32}, []api.ValueType{i32}).
 		Export(WasmFuncWrite).
 		// context tag
 		NewFunctionBuilder().
-		WithFunc(r.contextTag).
+		WithGoFunction(api.GoFunc(r.contextTag), []api.ValueType{}, []api.ValueType{i32}).
 		Export(WasmFuncContextTag).
 		// context data
 		NewFunctionBuilder().
-		WithFunc(r.contextData).
+		WithGoModuleFunction(api.GoModuleFunc(r.contextData), []api.ValueType{i32, i32}, []api.ValueType{i32}).
 		Export(WasmFuncContextData).
 		// context data size
 		NewFunctionBuilder().
-		WithFunc(r.contextDataSize).
+		WithGoFunction(api.GoFunc(r.contextDataSize), []api.ValueType{}, []api.ValueType{i32}).
 		Export(WasmFuncContextDataSize).
 		// Instantiate
 		Instantiate(r.ctx)
@@ -133,44 +135,55 @@ func (r *wazeroRuntime) Close() error {
 	return r.Runtime.Close(r.ctx)
 }
 
-func (r *wazeroRuntime) observeDataTag(ctx context.Context, tag uint32) {
+func (r *wazeroRuntime) observeDataTag(ctx context.Context, stack []uint64) {
+	tag := uint32(stack[0])
 	r.observed = append(r.observed, tag)
 }
 
-func (r *wazeroRuntime) write(ctx context.Context, m api.Module, tag uint32, pointer uint32, length int32) uint32 {
-	output, ok := m.Memory().Read(pointer, uint32(length))
+func (r *wazeroRuntime) write(ctx context.Context, m api.Module, stack []uint64) {
+	tag := uint32(stack[0])
+	pointer := uint32(stack[1])
+	length := uint32(stack[2])
+	output, ok := m.Memory().Read(pointer, length)
 	if !ok {
 		log.Printf("Memory.Read(%d, %d) out of range\n", pointer, length)
-		return 1
+		stack[0] = 1
+		return
 	}
 	buf := make([]byte, length)
 	copy(buf, output)
 
 	if err := r.serverlessCtx.Write(tag, buf); err != nil {
-		return 2
+		stack[0] = 2
+		return
 	}
-	return 0
+	stack[0] = 0
 }
 
-func (r *wazeroRuntime) contextTag(ctx context.Context, m api.Module) uint32 {
-	return r.serverlessCtx.Tag()
+func (r *wazeroRuntime) contextTag(ctx context.Context, stack []uint64) {
+	stack[0] = uint64(r.serverlessCtx.Tag())
 }
 
-func (r *wazeroRuntime) contextData(ctx context.Context, m api.Module, pointer uint32, limit uint32) (dataLen uint32) {
+func (r *wazeroRuntime) contextData(ctx context.Context, m api.Module, stack []uint64) {
+	pointer := uint32(stack[0])
+	limit := uint32(stack[1])
 	data := r.serverlessCtx.Data()
-	dataLen = uint32(len(data))
+	dataLen := uint32(len(data))
 	if dataLen > limit {
+		stack[0] = uint64(dataLen)
 		return
 	} else if dataLen == 0 {
+		stack[0] = 0
 		return
 	}
 	if ok := m.Memory().Write(pointer, data); !ok {
 		log.Printf("Memory.Write(%d, %d) out of range\n", pointer, dataLen)
-		return 0
+		stack[0] = 0
+		return
 	}
-	return
+	stack[0] = uint64(dataLen)
 }
 
-func (r *wazeroRuntime) contextDataSize(ctx context.Context, m api.Module) uint32 {
-	return uint32(len(r.serverlessCtx.Data()))
+func (r *wazeroRuntime) contextDataSize(ctx context.Context, stack []uint64) {
+	stack[0] = uint64(len(r.serverlessCtx.Data()))
 }

example/2-iopipe/Taskfile.yml

@@ -31,8 +31,7 @@ tasks:
   sfn:
     desc: sfn
     cmds:
-      # - "yomo run -n counter serverless/counter.go" # use latest yomo
-      - "yomo run -n counter serverless/sfn.wasm" # use custom go.mod for development
+      - "yomo run -n counter serverless/sfn.wasm" 
 
   source:
     desc: source