Refactoring collections/orderedintset and benchmarks for it

Docker-DCO-1.1-Signed-off-by: Alexandr Morozov <lk4d4math@gmail.com> (github: LK4D4)

pkg/collections/orderedintset.go

@@ -1,12 +1,13 @@
 package collections
 
 import (
+	"sort"
 	"sync"
 )
 
 // OrderedIntSet is a thread-safe sorted set and a stack.
 type OrderedIntSet struct {
-	sync.RWMutex
+	sync.Mutex
 	set []int
 }
 
@@ -15,29 +16,22 @@ func NewOrderedIntSet() *OrderedIntSet {
 	return &OrderedIntSet{}
 }
 
-// Push takes a string and adds it to the set. If the elem aready exists, it has no effect.
+// Push takes an int and adds it to the set. If the elem aready exists, it has no effect.
 func (s *OrderedIntSet) Push(elem int) {
-	s.RLock()
-	for _, e := range s.set {
-		if e == elem {
-			s.RUnlock()
-			return
-		}
-	}
-	s.RUnlock()
-
 	s.Lock()
+	if len(s.set) == 0 {
+		s.set = append(s.set, elem)
+		s.Unlock()
+		return
+	}
 
 	// Make sure the list is always sorted
-	for i, e := range s.set {
-		if elem < e {
-			s.set = append(s.set[:i], append([]int{elem}, s.set[i:]...)...)
-			s.Unlock()
-			return
-		}
+	i := sort.SearchInts(s.set, elem)
+	if i < len(s.set) && s.set[i] == elem {
+		s.Unlock()
+		return
 	}
-	// If we reach here, then elem is the biggest elem of the list.
-	s.set = append(s.set, elem)
+	s.set = append(s.set[:i], append([]int{elem}, s.set[i:]...)...)
 	s.Unlock()
 }
 
@@ -46,28 +40,26 @@ func (s *OrderedIntSet) Pop() int {
 	return s.PopFront()
 }
 
-// Pop returns the first elemen from the list and removes it.
+// Pop returns the first element from the list and removes it.
 // If the list is empty, it returns 0
 func (s *OrderedIntSet) PopFront() int {
-	s.RLock()
-
-	for i, e := range s.set {
-		ret := e
-		s.RUnlock()
-		s.Lock()
-		s.set = append(s.set[:i], s.set[i+1:]...)
+	s.Lock()
+	if len(s.set) == 0 {
 		s.Unlock()
-		return ret
+		return 0
 	}
-	s.RUnlock()
-
-	return 0
+	ret := s.set[0]
+	s.set = s.set[1:]
+	s.Unlock()
+	return ret
 }
 
 // PullBack retrieve the last element of the list.
 // The element is not removed.
 // If the list is empty, an empty element is returned.
 func (s *OrderedIntSet) PullBack() int {
+	s.Lock()
+	defer s.Unlock()
 	if len(s.set) == 0 {
 		return 0
 	}
@@ -76,21 +68,28 @@ func (s *OrderedIntSet) PullBack() int {
 
 // Exists checks if the given element present in the list.
 func (s *OrderedIntSet) Exists(elem int) bool {
-	for _, e := range s.set {
-		if e == elem {
-			return true
-		}
+	s.Lock()
+	if len(s.set) == 0 {
+		s.Unlock()
+		return false
 	}
-	return false
+	i := sort.SearchInts(s.set, elem)
+	res := i < len(s.set) && s.set[i] == elem
+	s.Unlock()
+	return res
 }
 
 // Remove removes an element from the list.
 // If the element is not found, it has no effect.
 func (s *OrderedIntSet) Remove(elem int) {
-	for i, e := range s.set {
-		if e == elem {
-			s.set = append(s.set[:i], s.set[i+1:]...)
-			return
-		}
+	s.Lock()
+	if len(s.set) == 0 {
+		s.Unlock()
+		return
 	}
+	i := sort.SearchInts(s.set, elem)
+	if i < len(s.set) && s.set[i] == elem {
+		s.set = append(s.set[:i], s.set[i+1:]...)
+	}
+	s.Unlock()
 }

pkg/collections/orderedintset_test.go

@@ -0,0 +1,71 @@
+package collections
+
+import (
+	"math/rand"
+	"testing"
+)
+
+func BenchmarkPush(b *testing.B) {
+	var testSet []int
+	for i := 0; i < 1000; i++ {
+		testSet = append(testSet, rand.Int())
+	}
+	s := NewOrderedIntSet()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for _, elem := range testSet {
+			s.Push(elem)
+		}
+	}
+}
+
+func BenchmarkPop(b *testing.B) {
+	var testSet []int
+	for i := 0; i < 1000; i++ {
+		testSet = append(testSet, rand.Int())
+	}
+	s := NewOrderedIntSet()
+	for _, elem := range testSet {
+		s.Push(elem)
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for j := 0; j < 1000; j++ {
+			s.Pop()
+		}
+	}
+}
+
+func BenchmarkExist(b *testing.B) {
+	var testSet []int
+	for i := 0; i < 1000; i++ {
+		testSet = append(testSet, rand.Intn(2000))
+	}
+	s := NewOrderedIntSet()
+	for _, elem := range testSet {
+		s.Push(elem)
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for j := 0; j < 1000; j++ {
+			s.Exists(j)
+		}
+	}
+}
+
+func BenchmarkRemove(b *testing.B) {
+	var testSet []int
+	for i := 0; i < 1000; i++ {
+		testSet = append(testSet, rand.Intn(2000))
+	}
+	s := NewOrderedIntSet()
+	for _, elem := range testSet {
+		s.Push(elem)
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for j := 0; j < 1000; j++ {
+			s.Remove(j)
+		}
+	}
+}