Merge pull request #412 from nspcc-dev/feature/map

VM: implement maps, closes #359.

pkg/vm/stack_item.go

@@ -19,10 +19,18 @@ func makeStackItem(v interface{}) StackItem {
 		return &BigIntegerItem{
 			value: big.NewInt(int64(val)),
 		}
+	case int64:
+		return &BigIntegerItem{
+			value: big.NewInt(val),
+		}
 	case []byte:
 		return &ByteArrayItem{
 			value: val,
 		}
+	case string:
+		return &ByteArrayItem{
+			value: []byte(val),
+		}
 	case bool:
 		return &BoolItem{
 			value: val,
@@ -193,3 +201,50 @@ func (i *ArrayItem) MarshalJSON() ([]byte, error) {
 func (i *ArrayItem) String() string {
 	return "Array"
 }
+
+// MapItem represents Map object.
+type MapItem struct {
+	value map[interface{}]StackItem
+}
+
+// NewMapItem returns new MapItem object.
+func NewMapItem() *MapItem {
+	return &MapItem{
+		value: make(map[interface{}]StackItem),
+	}
+}
+
+// Value implements StackItem interface.
+func (i *MapItem) Value() interface{} {
+	return i.value
+}
+
+// MarshalJSON implements the json.Marshaler interface.
+func (i *MapItem) String() string {
+	return "Map"
+}
+
+// Has checks if map has specified key.
+func (i *MapItem) Has(key StackItem) (ok bool) {
+	_, ok = i.value[toMapKey(key)]
+	return
+}
+
+// Add adds key-value pair to the map.
+func (i *MapItem) Add(key, value StackItem) {
+	i.value[toMapKey(key)] = value
+}
+
+// toMapKey converts StackItem so that it can be used as a map key.
+func toMapKey(key StackItem) interface{} {
+	switch t := key.(type) {
+	case *BoolItem:
+		return t.value
+	case *BigIntegerItem:
+		return t.value.Int64()
+	case *ByteArrayItem:
+		return string(t.value)
+	default:
+		panic("wrong key type")
+	}
+}

pkg/vm/vm.go

@@ -478,6 +478,11 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 				v.estack.PushVal(ta == tb)
 				break
 			}
+		} else if ma, ok := a.value.(*MapItem); ok {
+			if mb, ok := b.value.(*MapItem); ok {
+				v.estack.PushVal(ma == mb)
+				break
+			}
 		}
 		v.estack.PushVal(reflect.DeepEqual(a, b))
 
@@ -665,10 +670,7 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 		itemElem := v.estack.Pop()
 		arrElem := v.estack.Pop()
 
-		val := itemElem.value
-		if t, ok := itemElem.value.(*StructItem); ok {
-			val = t.Clone()
-		}
+		val := cloneIfStruct(itemElem.value)
 
 		switch t := arrElem.value.(type) {
 		case *ArrayItem:
@@ -705,11 +707,11 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 		v.estack.PushVal(l)
 
 	case PICKITEM:
-		var (
-			key   = v.estack.Pop()
-			obj   = v.estack.Pop()
-			index = int(key.BigInt().Int64())
-		)
+		key := v.estack.Pop()
+		validateMapKey(key)
+
+		obj := v.estack.Pop()
+		index := int(key.BigInt().Int64())
 
 		switch t := obj.value.(type) {
 		// Struct and Array items have their underlying value as []StackItem.
@@ -720,6 +722,12 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 			}
 			item := arr[index]
 			v.estack.PushVal(item)
+		case *MapItem:
+			if !t.Has(key.value) {
+				panic("invalid key")
+			}
+			k := toMapKey(key.value)
+			v.estack.Push(&Element{value: t.value[k]})
 		default:
 			arr := obj.Bytes()
 			if index < 0 || index >= len(arr) {
@@ -730,21 +738,24 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 		}
 
 	case SETITEM:
-		var (
-			item  = v.estack.Pop().value
-			key   = v.estack.Pop()
-			obj   = v.estack.Pop()
-			index = int(key.BigInt().Int64())
-		)
+		item := v.estack.Pop().value
+		key := v.estack.Pop()
+		validateMapKey(key)
+
+		obj := v.estack.Pop()
 
 		switch t := obj.value.(type) {
 		// Struct and Array items have their underlying value as []StackItem.
 		case *ArrayItem, *StructItem:
 			arr := t.Value().([]StackItem)
+			index := int(key.BigInt().Int64())
 			if index < 0 || index >= len(arr) {
 				panic("SETITEM: invalid index")
 			}
 			arr[index] = item
+		case *MapItem:
+			t.Add(key.value, item)
+
 		default:
 			panic(fmt.Sprintf("SETITEM: invalid item type %s", t))
 		}
@@ -757,23 +768,31 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 			}
 		}
 	case REMOVE:
-		key := int(v.estack.Pop().BigInt().Int64())
+		key := v.estack.Pop()
+		validateMapKey(key)
+
 		elem := v.estack.Pop()
 		switch t := elem.value.(type) {
 		case *ArrayItem:
 			a := t.value
-			if key < 0 || key >= len(a) {
+			k := int(key.BigInt().Int64())
+			if k < 0 || k >= len(a) {
 				panic("REMOVE: invalid index")
 			}
-			a = append(a[:key], a[key+1:]...)
+			a = append(a[:k], a[k+1:]...)
 			t.value = a
 		case *StructItem:
 			a := t.value
-			if key < 0 || key >= len(a) {
+			k := int(key.BigInt().Int64())
+			if k < 0 || k >= len(a) {
 				panic("REMOVE: invalid index")
 			}
-			a = append(a[:key], a[key+1:]...)
+			a = append(a[:k], a[k+1:]...)
 			t.value = a
+		case *MapItem:
+			m := t.value
+			k := toMapKey(key.value)
+			delete(m, k)
 		default:
 			panic("REMOVE: invalid type")
 		}
@@ -785,10 +804,10 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 		switch t := elem.value.Value().(type) {
 		case []StackItem:
 			v.estack.PushVal(len(t))
-		case []uint8:
+		case map[interface{}]StackItem:
 			v.estack.PushVal(len(t))
 		default:
-			panic("ARRAYSIZE: item not of type []StackItem")
+			v.estack.PushVal(len(elem.Bytes()))
 		}
 
 	case SIZE:
@@ -912,8 +931,71 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 		}
 		v.estack.PushVal(sigok)
 
-	case NEWMAP, HASKEY, KEYS, VALUES:
-		panic("unimplemented")
+	case NEWMAP:
+		v.estack.Push(&Element{value: NewMapItem()})
+
+	case KEYS:
+		item := v.estack.Pop()
+		if item == nil {
+			panic("no argument")
+		}
+
+		m, ok := item.value.(*MapItem)
+		if !ok {
+			panic("not a Map")
+		}
+
+		arr := make([]StackItem, 0, len(m.value))
+		for k := range m.value {
+			arr = append(arr, makeStackItem(k))
+		}
+		v.estack.PushVal(arr)
+
+	case VALUES:
+		item := v.estack.Pop()
+		if item == nil {
+			panic("no argument")
+		}
+
+		var arr []StackItem
+		switch t := item.value.(type) {
+		case *ArrayItem, *StructItem:
+			src := t.Value().([]StackItem)
+			arr = make([]StackItem, len(src))
+			for i := range src {
+				arr[i] = cloneIfStruct(src[i])
+			}
+		case *MapItem:
+			arr = make([]StackItem, 0, len(t.value))
+			for k := range t.value {
+				arr = append(arr, cloneIfStruct(t.value[k]))
+			}
+		default:
+			panic("not a Map, Array or Struct")
+		}
+
+		v.estack.PushVal(arr)
+
+	case HASKEY:
+		key := v.estack.Pop()
+		validateMapKey(key)
+
+		c := v.estack.Pop()
+		if c == nil {
+			panic("no value found")
+		}
+		switch t := c.value.(type) {
+		case *ArrayItem, *StructItem:
+			index := key.BigInt().Int64()
+			if index < 0 {
+				panic("negative index")
+			}
+			v.estack.PushVal(index < int64(len(c.Array())))
+		case *MapItem:
+			v.estack.PushVal(t.Has(key.value))
+		default:
+			panic("wrong collection type")
+		}
 
 	// Cryptographic operations.
 	case SHA1:
@@ -950,6 +1032,15 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 	}
 }
 
+func cloneIfStruct(item StackItem) StackItem {
+	switch it := item.(type) {
+	case *StructItem:
+		return it.Clone()
+	default:
+		return it
+	}
+}
+
 func makeArrayOfFalses(n int) []StackItem {
 	items := make([]StackItem, n)
 	for i := range items {
@@ -958,6 +1049,16 @@ func makeArrayOfFalses(n int) []StackItem {
 	return items
 }
 
+func validateMapKey(key *Element) {
+	if key == nil {
+		panic("no key found")
+	}
+	switch key.value.(type) {
+	case *ArrayItem, *StructItem, *MapItem:
+		panic("key can't be a collection")
+	}
+}
+
 func init() {
 	log.SetPrefix("NEO-GO-VM > ")
 	log.SetFlags(0)

pkg/vm/vm_test.go

@@ -446,6 +446,27 @@ func TestEQUALArrayFalse(t *testing.T) {
 	assert.Equal(t, &BoolItem{false}, vm.estack.Pop().value)
 }
 
+func TestEQUALMapTrue(t *testing.T) {
+	prog := makeProgram(DUP, EQUAL)
+	vm := load(prog)
+	vm.estack.Push(&Element{value: NewMapItem()})
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, &BoolItem{true}, vm.estack.Pop().value)
+}
+
+func TestEQUALMapFalse(t *testing.T) {
+	prog := makeProgram(EQUAL)
+	vm := load(prog)
+	vm.estack.Push(&Element{value: NewMapItem()})
+	vm.estack.Push(&Element{value: NewMapItem()})
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, &BoolItem{false}, vm.estack.Pop().value)
+}
+
 func TestNumEqual(t *testing.T) {
 	prog := makeProgram(NUMEQUAL)
 	vm := load(prog)
@@ -646,6 +667,39 @@ func TestPICKITEMByteArray(t *testing.T) {
 	assert.Equal(t, makeStackItem(2), vm.estack.Pop().value)
 }
 
+func TestPICKITEMMap(t *testing.T) {
+	prog := makeProgram(PICKITEM)
+	vm := load(prog)
+
+	m := NewMapItem()
+	m.Add(makeStackItem(5), makeStackItem(3))
+	vm.estack.Push(&Element{value: m})
+	vm.estack.PushVal(makeStackItem(5))
+
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem(3), vm.estack.Pop().value)
+}
+
+func TestSETITEMMap(t *testing.T) {
+	prog := makeProgram(SETITEM, PICKITEM)
+	vm := load(prog)
+
+	m := NewMapItem()
+	m.Add(makeStackItem(5), makeStackItem(3))
+	vm.estack.Push(&Element{value: m})
+	vm.estack.PushVal(5)
+	vm.estack.Push(&Element{value: m})
+	vm.estack.PushVal(5)
+	vm.estack.PushVal([]byte{0, 1})
+
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem([]byte{0, 1}), vm.estack.Pop().value)
+}
+
 func TestSIZENoArgument(t *testing.T) {
 	prog := makeProgram(SIZE)
 	vm := load(prog)
@@ -673,6 +727,207 @@ func TestSIZEBool(t *testing.T) {
 	assert.Equal(t, makeStackItem(1), vm.estack.Pop().value)
 }
 
+func TestARRAYSIZEArray(t *testing.T) {
+	prog := makeProgram(ARRAYSIZE)
+	vm := load(prog)
+	vm.estack.PushVal([]StackItem{
+		makeStackItem(1),
+		makeStackItem([]byte{}),
+	})
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem(2), vm.estack.Pop().value)
+}
+
+func TestARRAYSIZEMap(t *testing.T) {
+	prog := makeProgram(ARRAYSIZE)
+	vm := load(prog)
+
+	m := NewMapItem()
+	m.Add(makeStackItem(5), makeStackItem(6))
+	m.Add(makeStackItem([]byte{0, 1}), makeStackItem(6))
+	vm.estack.Push(&Element{value: m})
+
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem(2), vm.estack.Pop().value)
+}
+
+func TestKEYSMap(t *testing.T) {
+	prog := makeProgram(KEYS)
+	vm := load(prog)
+
+	m := NewMapItem()
+	m.Add(makeStackItem(5), makeStackItem(6))
+	m.Add(makeStackItem([]byte{0, 1}), makeStackItem(6))
+	vm.estack.Push(&Element{value: m})
+
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+
+	top := vm.estack.Pop().value.(*ArrayItem)
+	assert.Equal(t, 2, len(top.value))
+	assert.Contains(t, top.value, makeStackItem(5))
+	assert.Contains(t, top.value, makeStackItem([]byte{0, 1}))
+}
+
+func TestKEYSNoArgument(t *testing.T) {
+	prog := makeProgram(KEYS)
+	vm := load(prog)
+	vm.Run()
+	assert.Equal(t, true, vm.HasFailed())
+}
+
+func TestKEYSWrongType(t *testing.T) {
+	prog := makeProgram(KEYS)
+	vm := load(prog)
+	vm.estack.PushVal([]StackItem{})
+	vm.Run()
+	assert.Equal(t, true, vm.HasFailed())
+}
+
+func TestVALUESMap(t *testing.T) {
+	prog := makeProgram(VALUES)
+	vm := load(prog)
+
+	m := NewMapItem()
+	m.Add(makeStackItem(5), makeStackItem([]byte{2, 3}))
+	m.Add(makeStackItem([]byte{0, 1}), makeStackItem([]StackItem{}))
+	vm.estack.Push(&Element{value: m})
+
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+
+	top := vm.estack.Pop().value.(*ArrayItem)
+	assert.Equal(t, 2, len(top.value))
+	assert.Contains(t, top.value, makeStackItem([]byte{2, 3}))
+	assert.Contains(t, top.value, makeStackItem([]StackItem{}))
+}
+
+func TestVALUESArray(t *testing.T) {
+	prog := makeProgram(VALUES)
+	vm := load(prog)
+	vm.estack.PushVal([]StackItem{makeStackItem(4)})
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, &ArrayItem{[]StackItem{makeStackItem(4)}}, vm.estack.Pop().value)
+}
+
+func TestVALUESNoArgument(t *testing.T) {
+	prog := makeProgram(VALUES)
+	vm := load(prog)
+	vm.Run()
+	assert.Equal(t, true, vm.HasFailed())
+}
+
+func TestVALUESWrongType(t *testing.T) {
+	prog := makeProgram(VALUES)
+	vm := load(prog)
+	vm.estack.PushVal(5)
+	vm.Run()
+	assert.Equal(t, true, vm.HasFailed())
+}
+
+func TestHASKEYArrayTrue(t *testing.T) {
+	prog := makeProgram(PUSH5, NEWARRAY, PUSH4, HASKEY)
+	vm := load(prog)
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem(true), vm.estack.Pop().value)
+}
+
+func TestHASKEYArrayFalse(t *testing.T) {
+	prog := makeProgram(PUSH5, NEWARRAY, PUSH5, HASKEY)
+	vm := load(prog)
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem(false), vm.estack.Pop().value)
+}
+
+func TestHASKEYStructTrue(t *testing.T) {
+	prog := makeProgram(PUSH5, NEWSTRUCT, PUSH4, HASKEY)
+	vm := load(prog)
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem(true), vm.estack.Pop().value)
+}
+
+func TestHASKEYStructFalse(t *testing.T) {
+	prog := makeProgram(PUSH5, NEWSTRUCT, PUSH5, HASKEY)
+	vm := load(prog)
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem(false), vm.estack.Pop().value)
+}
+
+func TestHASKEYMapTrue(t *testing.T) {
+	prog := makeProgram(HASKEY)
+	vm := load(prog)
+	m := NewMapItem()
+	m.Add(makeStackItem(5), makeStackItem(6))
+	vm.estack.Push(&Element{value: m})
+	vm.estack.PushVal(5)
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem(true), vm.estack.Pop().value)
+}
+
+func TestHASKEYMapFalse(t *testing.T) {
+	prog := makeProgram(HASKEY)
+	vm := load(prog)
+	m := NewMapItem()
+	m.Add(makeStackItem(5), makeStackItem(6))
+	vm.estack.Push(&Element{value: m})
+	vm.estack.PushVal(6)
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem(false), vm.estack.Pop().value)
+}
+
+func TestHASKEYNoArguments(t *testing.T) {
+	prog := makeProgram(HASKEY)
+	vm := load(prog)
+	vm.Run()
+	assert.Equal(t, true, vm.HasFailed())
+}
+
+func TestHASKEY1Argument(t *testing.T) {
+	prog := makeProgram(HASKEY)
+	vm := load(prog)
+	vm.estack.PushVal(1)
+	vm.Run()
+	assert.Equal(t, true, vm.HasFailed())
+}
+
+func TestHASKEYWrongKeyType(t *testing.T) {
+	prog := makeProgram(HASKEY)
+	vm := load(prog)
+	vm.estack.PushVal([]StackItem{})
+	vm.estack.PushVal([]StackItem{})
+	vm.Run()
+	assert.Equal(t, true, vm.HasFailed())
+}
+
+func TestHASKEYWrongCollectionType(t *testing.T) {
+	prog := makeProgram(HASKEY)
+	vm := load(prog)
+	vm.estack.PushVal(1)
+	vm.estack.PushVal(2)
+	vm.Run()
+	assert.Equal(t, true, vm.HasFailed())
+}
+
 func TestSIGNNoArgument(t *testing.T) {
 	prog := makeProgram(SIGN)
 	vm := load(prog)
@@ -1490,6 +1745,23 @@ func TestREMOVEGood(t *testing.T) {
 	assert.Equal(t, makeStackItem(1), vm.estack.Pop().value)
 }
 
+func TestREMOVEMap(t *testing.T) {
+	prog := makeProgram(REMOVE, PUSH5, HASKEY)
+	vm := load(prog)
+
+	m := NewMapItem()
+	m.Add(makeStackItem(5), makeStackItem(3))
+	m.Add(makeStackItem([]byte{0, 1}), makeStackItem([]byte{2, 3}))
+	vm.estack.Push(&Element{value: m})
+	vm.estack.Push(&Element{value: m})
+	vm.estack.PushVal(makeStackItem(5))
+
+	vm.Run()
+	assert.Equal(t, false, vm.HasFailed())
+	assert.Equal(t, 1, vm.estack.Len())
+	assert.Equal(t, makeStackItem(false), vm.estack.Pop().value)
+}
+
 func TestCHECKSIGNoArgs(t *testing.T) {
 	prog := makeProgram(CHECKSIG)
 	vm := load(prog)