Make MapVariableStorage concurrent-safe #6

README.md

@@ -99,7 +99,7 @@ commands to the handler.
        vm := &yarn.VirtualMachine{
            Program: program,
            Handler: myHandler,
-           Vars:    make(yarn.MapVariableStorage), // or your own VariableStorage implementation
+           Vars:    yarn.NewMapVariableStorage(), // or your own VariableStorage implementation
            FuncMap: yarn.FuncMap{ // this is optional
                "last_value": func(x ...any) any {
                    return x[len(x)-1]

cmd/yarnrunner.go

@@ -52,7 +52,7 @@ func main() {
 		Handler: &dialogueHandler{
 			stringTable: stringTable,
 		},
-		Vars: make(yarn.MapVariableStorage),
+		Vars: yarn.NewMapVariableStorage(),
 	}
 	if err := vm.Run(*startNode); err != nil {
 		log.Printf("Yarn VM error: %v", err)

vars.go

@@ -14,30 +14,60 @@
 
 package yarn
 
+import "sync"
+
 // VariableStorage stores values of any kind.
 type VariableStorage interface {
 	Clear()
-	GetValue(name string) (value interface{}, ok bool)
-	SetValue(name string, value interface{})
+	GetValue(name string) (value any, ok bool)
+	SetValue(name string, value any)
 }
 
 // MapVariableStorage implements VariableStorage, in memory, using a map.
-type MapVariableStorage map[string]interface{}
+type MapVariableStorage struct {
+	mu sync.RWMutex
+	m  map[string]any
+}
+
+// NewMapVariableStorage creates a new empty MapVariableStorage.
+func NewMapVariableStorage() *MapVariableStorage {
+	return &MapVariableStorage{
+		m: make(map[string]any),
+	}
+}
 
 // Clear empties the storage of all values.
-func (m MapVariableStorage) Clear() {
-	for name := range m {
-		delete(m, name)
+func (m *MapVariableStorage) Clear() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	for name := range m.m {
+		delete(m.m, name)
 	}
 }
 
 // GetValue fetches a value from the map, returning (nil, false) if not present.
-func (m MapVariableStorage) GetValue(name string) (value interface{}, found bool) {
-	value, found = m[name]
+func (m *MapVariableStorage) GetValue(name string) (value any, found bool) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	value, found = m.m[name]
 	return value, found
 }
 
 // SetValue sets a value in the map.
-func (m MapVariableStorage) SetValue(name string, value interface{}) {
-	m[name] = value
+func (m *MapVariableStorage) SetValue(name string, value any) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.m[name] = value
+}
+
+// Copy returns a copy of the full contents of the storage, as a regular map.
+func (m *MapVariableStorage) Copy() map[string]any {
+	m2 := make(map[string]any, len(m.m))
+
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	for name, val := range m.m {
+		m2[name] = val
+	}
+	return m2
 }

vm_test.go

@@ -47,7 +47,7 @@ func TestAllTestPlans(t *testing.T) {
 			vm := &VirtualMachine{
 				Program: prog,
 				Handler: testplan,
-				Vars:    make(MapVariableStorage),
+				Vars:    NewMapVariableStorage(),
 				FuncMap: FuncMap{
 					// Used by various
 					"assert": func(x interface{}) error {