Merge pull request #2659 from nspcc-dev/optimize-refcount

vm: optimize refcount

pkg/vm/fuzz_test.go

@@ -16,6 +16,12 @@ var fuzzSeedValidScripts = [][]byte{
 	makeProgram(opcode.PUSH2, opcode.NEWARRAY, opcode.DUP, opcode.PUSH0, opcode.PUSH1, opcode.SETITEM, opcode.VALUES),
 	append([]byte{byte(opcode.PUSHDATA1), 10}, randomBytes(10)...),
 	append([]byte{byte(opcode.PUSHDATA1), 100}, randomBytes(100)...),
+	// Simplified version of fuzzer output from #2659.
+	{byte(opcode.CALL), 3, byte(opcode.ASSERT),
+		byte(opcode.CALL), 3, byte(opcode.ASSERT),
+		byte(opcode.DEPTH), byte(opcode.PACKSTRUCT), byte(opcode.DUP),
+		byte(opcode.UNPACK), byte(opcode.PACKSTRUCT), byte(opcode.POPITEM),
+		byte(opcode.DEPTH)},
 }
 
 func FuzzIsScriptCorrect(f *testing.F) {

pkg/vm/ref_counter.go

@@ -7,15 +7,6 @@ import (
 // refCounter represents a reference counter for the VM.
 type refCounter int
 
-type (
-	rcInc interface {
-		IncRC() int
-	}
-	rcDec interface {
-		DecRC() int
-	}
-)
-
 func newRefCounter() *refCounter {
 	return new(refCounter)
 }
@@ -27,20 +18,26 @@ func (r *refCounter) Add(item stackitem.Item) {
 	}
 	*r++
 
-	irc, ok := item.(rcInc)
-	if !ok || irc.IncRC() > 1 {
-		return
-	}
 	switch t := item.(type) {
-	case *stackitem.Array, *stackitem.Struct:
-		for _, it := range item.Value().([]stackitem.Item) {
-			r.Add(it)
+	case *stackitem.Array:
+		if t.IncRC() == 1 {
+			for _, it := range t.Value().([]stackitem.Item) {
+				r.Add(it)
+			}
+		}
+	case *stackitem.Struct:
+		if t.IncRC() == 1 {
+			for _, it := range t.Value().([]stackitem.Item) {
+				r.Add(it)
+			}
 		}
 	case *stackitem.Map:
-		elems := t.Value().([]stackitem.MapElement)
-		for i := range elems {
-			r.Add(elems[i].Key)
-			r.Add(elems[i].Value)
+		if t.IncRC() == 1 {
+			elems := t.Value().([]stackitem.MapElement)
+			for i := range elems {
+				r.Add(elems[i].Key)
+				r.Add(elems[i].Value)
+			}
 		}
 	}
 }
@@ -52,20 +49,26 @@ func (r *refCounter) Remove(item stackitem.Item) {
 	}
 	*r--
 
-	irc, ok := item.(rcDec)
-	if !ok || irc.DecRC() > 0 {
-		return
-	}
 	switch t := item.(type) {
-	case *stackitem.Array, *stackitem.Struct:
-		for _, it := range item.Value().([]stackitem.Item) {
-			r.Remove(it)
+	case *stackitem.Array:
+		if t.DecRC() == 0 {
+			for _, it := range t.Value().([]stackitem.Item) {
+				r.Remove(it)
+			}
+		}
+	case *stackitem.Struct:
+		if t.DecRC() == 0 {
+			for _, it := range t.Value().([]stackitem.Item) {
+				r.Remove(it)
+			}
 		}
 	case *stackitem.Map:
-		elems := t.Value().([]stackitem.MapElement)
-		for i := range elems {
-			r.Remove(elems[i].Key)
-			r.Remove(elems[i].Value)
+		if t.DecRC() == 0 {
+			elems := t.Value().([]stackitem.MapElement)
+			for i := range elems {
+				r.Remove(elems[i].Key)
+				r.Remove(elems[i].Value)
+			}
 		}
 	}
 }

pkg/vm/ref_counter_test.go

@@ -55,3 +55,14 @@ func BenchmarkRefCounter_Add(b *testing.B) {
 		rc.Add(a)
 	}
 }
+
+func BenchmarkRefCounter_AddRemove(b *testing.B) {
+	a := stackitem.NewArray([]stackitem.Item{})
+	rc := newRefCounter()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		rc.Add(a)
+		rc.Remove(a)
+	}
+}