fix: Reset lens state when calling Reset (#148)

Author: AndrewSisley

host-go/engine/module/instance.go

@@ -22,6 +22,10 @@ type Instance struct {
 	// module after this function has been called are not guaranteed to be visible to the returned io.Reader.
 	Memory func() Memory
 
+	// Reset resets the wasm memory to the state it was in immediately after configuration, before any items were
+	// pulled through it.
+	Reset func()
+
 	// OwnedBy hosts a reference to any object(s) that may be required to live in memory for the lifetime of this Module.
 	//
 	// This is very important when working with some libraries (such as wasmer-go), as without this, dependencies of other members

host-go/engine/pipes/fromPipe.go

@@ -90,6 +90,7 @@ func (p *fromPipe[TSource, TResult]) Bytes() ([]byte, error) {
 }
 
 func (p *fromPipe[TSource, TResult]) Reset() {
+	p.instance.Reset()
 	p.source.Reset()
 }
 

host-go/engine/pipes/fromSource.go

@@ -91,6 +91,7 @@ func (s *fromSource[TSource, TResult]) Bytes() ([]byte, error) {
 }
 
 func (s *fromSource[TSource, TResult]) Reset() {
+	s.instance.Reset()
 	s.source.Reset()
 }
 

host-go/runtimes/js/runtime.go

@@ -156,6 +156,11 @@ func (m *wModule) NewInstance(functionName string, paramSets ...map[string]any)
 		}
 	}
 
+	jsMemory := js.Global().Get("Uint8Array").New(memory.Get("buffer"))
+	initialLen := jsMemory.Get("length").Int()
+	initialState := make([]byte, initialLen)
+	js.CopyBytesToGo(initialState, jsMemory)
+
 	return module.Instance{
 		Alloc: func(u module.MemSize) (module.MemSize, error) {
 			result := alloc.Invoke(int32(u))
@@ -173,6 +178,16 @@ func (m *wModule) NewInstance(functionName string, paramSets ...map[string]any)
 			buffer := memory.Get("buffer")
 			return newMemory(buffer)
 		},
+		Reset: func() {
+			initialLen := len(initialState)
+			currentLen := jsMemory.Get("length").Int()
+
+			js.CopyBytesToJS(jsMemory, initialState)
+
+			for i := initialLen; i < currentLen; i++ {
+				jsMemory.SetIndex(i, js.ValueOf(0))
+			}
+		},
 		OwnedBy: instance,
 	}, nil
 }

host-go/runtimes/wasmer/runtime.go

@@ -151,6 +151,10 @@ func (m *wModule) NewInstance(functionName string, paramSets ...map[string]any)
 		}
 	}
 
+	memSlice := memory.Data()
+	initialState := make([]byte, len(memSlice))
+	copy(initialState, memSlice)
+
 	return module.Instance{
 		Alloc: func(u module.MemSize) (module.MemSize, error) {
 			r, err := alloc.Call(u)
@@ -173,6 +177,16 @@ func (m *wModule) NewInstance(functionName string, paramSets ...map[string]any)
 		Memory: func() module.Memory {
 			return module.NewBytesMemory(memory.Data())
 		},
+		Reset: func() {
+			initialLen := len(initialState)
+			currentMemory := memory.Data()
+
+			copy(currentMemory, initialState)
+
+			for i := initialLen; i < len(currentMemory); i++ {
+				currentMemory[i] = 0
+			}
+		},
 		OwnedBy: instance,
 	}, nil
 }

host-go/runtimes/wasmtime/runtime.go

@@ -144,6 +144,10 @@ func (m *wModule) NewInstance(functionName string, paramSets ...map[string]any)
 		}
 	}
 
+	memSlice := memory.UnsafeData(m.rt.store)
+	initialState := make([]byte, len(memSlice))
+	copy(initialState, memSlice)
+
 	return module.Instance{
 		Alloc: func(u module.MemSize) (module.MemSize, error) {
 			r, err := alloc.Call(m.rt.store, u)
@@ -166,6 +170,16 @@ func (m *wModule) NewInstance(functionName string, paramSets ...map[string]any)
 		Memory: func() module.Memory {
 			return module.NewBytesMemory(memory.UnsafeData(m.rt.store))
 		},
+		Reset: func() {
+			initialLen := len(initialState)
+			currentMemory := memory.UnsafeData(m.rt.store)
+
+			copy(currentMemory, initialState)
+
+			for i := initialLen; i < len(currentMemory); i++ {
+				currentMemory[i] = 0
+			}
+		},
 		OwnedBy: instance,
 	}, nil
 }

host-go/runtimes/wazero/runtime.go

@@ -143,6 +143,8 @@ func (m *wModule) NewInstance(functionName string, paramSets ...map[string]any)
 		}
 	}
 
+	initialState, _ := memory.Read(0, memory.Size())
+
 	return module.Instance{
 		Alloc: func(u module.MemSize) (module.MemSize, error) {
 			r, err := alloc.Call(ctx, uint64(u))
@@ -165,6 +167,15 @@ func (m *wModule) NewInstance(functionName string, paramSets ...map[string]any)
 		Memory: func() module.Memory {
 			return newMemory(memory)
 		},
+		Reset: func() {
+			initialLen := uint32(len(initialState))
+			currentLen := memory.Size()
+			memory.Write(0, initialState)
+
+			for i := initialLen; i < currentLen; i++ {
+				memory.WriteByte(i, 0)
+			}
+		},
 		OwnedBy: instance,
 	}, nil
 }

tests/action/transform.go

@@ -0,0 +1,66 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+package action
+
+import (
+	"fmt"
+
+	"github.com/sourcenetwork/immutable/enumerable"
+	"github.com/sourcenetwork/lens/host-go/store"
+	"github.com/stretchr/testify/require"
+)
+
+// Transform action executes the transform of the given ID with the given input
+// and asserts that the output matches the given expected output.
+type Transform struct {
+	Nodeful
+
+	LensID   string
+	Input    enumerable.Enumerable[store.Document]
+	Expected enumerable.Enumerable[store.Document]
+}
+
+var _ Action = (*Add)(nil)
+var _ Stateful = (*Add)(nil)
+
+func (a *Transform) Execute() {
+	lensID := replace(a.s, a.LensID)
+	for nodeIndex, n := range a.Nodes() {
+		output, err := n.Store.Transform(a.s.Ctx, a.Input, lensID)
+		require.NoError(a.s.T, err)
+
+		n.Transforms = append(n.Transforms, output)
+
+		for i := 0; true; i++ {
+			println(fmt.Sprintf("TransformEval: Node: {%v} item: {%v}", nodeIndex, i))
+
+			hasNext, err := output.Next()
+			require.NoError(a.s.T, err)
+
+			expectedHasNext, err := a.Expected.Next()
+			require.NoError(a.s.T, err)
+
+			require.Equal(a.s.T, expectedHasNext, hasNext)
+
+			if !hasNext {
+				break
+			}
+
+			value, err := output.Value()
+			require.NoError(a.s.T, err)
+
+			expectedValue, err := a.Expected.Value()
+			require.NoError(a.s.T, err)
+
+			require.Equal(a.s.T, expectedValue, value)
+		}
+
+		expectedHasNext, err := a.Expected.Next()
+		require.NoError(a.s.T, err)
+		require.False(a.s.T, expectedHasNext)
+
+		a.Expected.Reset()
+	}
+}

tests/action/transform_reset.go

@@ -0,0 +1,21 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+package action
+
+// TransformReset resets the given transform.
+type TransformReset struct {
+	Nodeful
+
+	TransformIndex int
+}
+
+var _ Action = (*Add)(nil)
+var _ Stateful = (*Add)(nil)
+
+func (a *TransformReset) Execute() {
+	for _, n := range a.Nodes() {
+		n.Transforms[a.TransformIndex].Reset()
+	}
+}