github.com/bananabytelabs/wazero@v0.0.0-20240105073314-54b22a776da8/internal/wasm/module_instance_lookup_test.go (about)

     1  package wasm
     2  
     3  import (
     4  	"context"
     5  	"math"
     6  	"testing"
     7  
     8  	"github.com/bananabytelabs/wazero/api"
     9  	"github.com/bananabytelabs/wazero/internal/testing/require"
    10  )
    11  
    12  func TestModuleInstance_LookupFunction(t *testing.T) {
    13  	var called int
    14  	hostModule := &Module{
    15  		IsHostModule: true,
    16  		CodeSection: []Code{
    17  			{GoFunc: api.GoFunc(func(context.Context, []uint64) {
    18  				called++
    19  			})},
    20  			{GoFunc: api.GoModuleFunc(func(context.Context, api.Module, []uint64) {
    21  				called++
    22  			})},
    23  		},
    24  		FunctionDefinitionSection: []FunctionDefinition{{}, {}},
    25  	}
    26  
    27  	me := &mockModuleEngine{
    28  		lookupEntries: map[Index]mockModuleEngineLookupEntry{
    29  			0: {m: &ModuleInstance{Source: hostModule}, index: 0},
    30  			1: {m: &ModuleInstance{Source: hostModule}, index: 1},
    31  		},
    32  	}
    33  	m := &ModuleInstance{Engine: me}
    34  
    35  	t.Run("host", func(t *testing.T) {
    36  		gf, ok := m.LookupFunction(nil, 0, 0).(*lookedUpGoFunction)
    37  		require.True(t, ok)
    38  		require.Nil(t, gf.lookedUpModule) // GoFunction doesn't need looked up module.
    39  		require.Equal(t, &hostModule.FunctionDefinitionSection[0], gf.def)
    40  		err := gf.CallWithStack(context.Background(), nil)
    41  		require.NoError(t, err)
    42  
    43  		gmf, ok := m.LookupFunction(nil, 0, 1).(*lookedUpGoFunction)
    44  		require.True(t, ok)
    45  		require.Equal(t, m, gmf.lookedUpModule)
    46  		require.Equal(t, hostModule.CodeSection[1].GoFunc, gmf.g)
    47  		require.Equal(t, &hostModule.FunctionDefinitionSection[1], gmf.def)
    48  		err = gmf.CallWithStack(context.Background(), nil)
    49  		require.NoError(t, err)
    50  
    51  		require.Equal(t, 2, called)
    52  	})
    53  
    54  	t.Run("wasm", func(t *testing.T) {
    55  		me.lookupEntries[2] = mockModuleEngineLookupEntry{
    56  			m: &ModuleInstance{
    57  				Source: &Module{},
    58  				Engine: me,
    59  			},
    60  			index: 100,
    61  		}
    62  		wf, ok := m.LookupFunction(nil, 0, 2).(*mockCallEngine)
    63  		require.True(t, ok)
    64  		require.Equal(t, Index(100), wf.index)
    65  	})
    66  }
    67  
    68  func Test_lookedUpGoModuleFunction(t *testing.T) {
    69  	def := &FunctionDefinition{
    70  		Functype: &FunctionType{
    71  			ParamNumInUint64: 2, ResultNumInUint64: 1,
    72  			Params: []ValueType{ValueTypeI64, ValueTypeF64}, Results: []ValueType{ValueTypeI64},
    73  		},
    74  	}
    75  
    76  	expMod := &ModuleInstance{}
    77  	var called bool
    78  	l := &lookedUpGoFunction{
    79  		lookedUpModule: expMod,
    80  		def:            def,
    81  		g: api.GoModuleFunc(
    82  			func(ctx context.Context, mod api.Module, stack []uint64) {
    83  				require.Equal(t, expMod, mod)
    84  				require.Equal(t, []uint64{math.MaxUint64, math.Float64bits(math.Pi)}, stack)
    85  				stack[0] = 1
    86  				called = true
    87  			},
    88  		),
    89  	}
    90  	require.Equal(t, l.Definition().(*FunctionDefinition), l.def)
    91  	result, err := l.Call(context.Background(), math.MaxUint64, math.Float64bits(math.Pi))
    92  	require.NoError(t, err)
    93  	require.True(t, called)
    94  	require.Equal(t, []uint64{1}, result)
    95  	require.Equal(t, 1, len(result))
    96  }