Set/Clear

sbromberger · Nov 5, 2022 · 095528b · 095528b
1 parent b01912e
commit 095528b
Show file tree

Hide file tree

Showing 3 changed files with 148 additions and 13 deletions.
diff --git a/abitvec.go b/abitvec.go
@@ -83,3 +83,31 @@ func (bv ABitVec) Get(k uint64) (bool, error) {
 	bucket, bit := offset(k)
 	return bv.buckets[bucket]&bit != 0, nil
 }
+
+// Set sets bit `k` to true unconditionally.
+func (bv ABitVec) Set(k uint64) error {
+	if k >= bv.capacity {
+		return errors.New("Attempt to access element beyond vector bounds")
+	}
+	bucket, bit := offset(k)
+	old := atomic.LoadUint64(&bv.buckets[bucket])
+retry:
+	if atomic.CompareAndSwapUint64(&bv.buckets[bucket], old, old|bit) {
+		return nil
+	}
+	goto retry
+}
+
+// Clear sets bit `k` to false unconditionally.
+func (bv ABitVec) Clear(k uint64) error {
+	if k >= bv.capacity {
+		return errors.New("Attempt to access element beyond vector bounds")
+	}
+	bucket, bit := offset(k)
+	old := atomic.LoadUint64(&bv.buckets[bucket])
+retry:
+	if atomic.CompareAndSwapUint64(&bv.buckets[bucket], old, old&(allSet^bit)) {
+		return nil
+	}
+	goto retry
+}
diff --git a/bitvec.go b/bitvec.go
@@ -1,12 +1,15 @@
 // Package bitvec is bit-vector with atomic and non-atomic access
 package bitvec
 
-import "errors"
+import (
+	"errors"
+)
 
 const (
-	nbits = 6          // 64 bits in a uint64
-	ws    = 1 << nbits // constant 64
-	mask  = ws - 1     // all ones
+	nbits  = 6          // 64 bits in a uint64
+	ws     = 1 << nbits // constant 64
+	mask   = ws - 1     // all ones
+	allSet = ^uint64(0)
 )
 
 func offset(k uint64) (bucket, bit uint64) {
@@ -52,3 +55,23 @@ func (bv BitVec) Get(k uint64) (bool, error) {
 	bucket, bit := offset(k)
 	return bv.buckets[bucket]&bit != 0, nil
 }
+
+// Set sets bit `k` to true unconditionally.
+func (bv BitVec) Set(k uint64) error {
+	if k >= bv.capacity {
+		return errors.New("Attempt to access element beyond vector bounds")
+	}
+	bucket, bit := offset(k)
+	bv.buckets[bucket] |= bit
+	return nil
+}
+
+// Clear sets bit `k` to false unconditionally.
+func (bv BitVec) Clear(k uint64) error {
+	if k >= bv.capacity {
+		return errors.New("Attempt to access element beyond vector bounds")
+	}
+	bucket, bit := offset(k)
+	bv.buckets[bucket] &= (allSet ^ bit)
+	return nil
+}
diff --git a/bitvec_test.go b/bitvec_test.go
@@ -14,7 +14,9 @@ var setSizes = []int{13, 256, 65536, 16777216}
 
 func TestBV(t *testing.T) {
 	bv := New(100)
+	bv2 := New(100)
 	abv := NewAtomic(100)
+	abv2 := NewAtomic(100)
 	t.Run("valid indices", func(t *testing.T) {
 		for _, k := range GoodSets {
 			if !bv.TrySet(k) {
@@ -23,13 +25,30 @@ func TestBV(t *testing.T) {
 			if !abv.TrySet(k) {
 				t.Errorf("abv could not set good index %d", k)
 			}
+			bv2.Set(k)
+			abv2.Set(k)
 
 			if x, _ := bv.Get(k); !x {
 				t.Errorf("bv bit %d was supposed to be set but wasn't", k)
 			}
 			if x, _ := abv.Get(k); !x {
 				t.Errorf("abv bit %d was supposed to be set but wasn't", k)
 			}
+			if x, _ := bv2.Get(k); !x {
+				t.Errorf("bv2 bit %d was supposed to be set but wasn't", k)
+			}
+			if x, _ := abv2.Get(k); !x {
+				t.Errorf("abv2 bit %d was supposed to be set but wasn't", k)
+			}
+
+			bv.Clear(k)
+			if x, _ := bv.Get(k); x {
+				t.Errorf("bv bit %d was supposed to be cleared but wasn't", k)
+			}
+			abv.Clear(k)
+			if x, _ := abv.Get(k); x {
+				t.Errorf("abv bit %d was supposed to be cleared but wasn't", k)
+			}
 		}
 	})
 	t.Run("invalid indices", func(t *testing.T) {
@@ -40,12 +59,21 @@ func TestBV(t *testing.T) {
 			if abv.TrySet(k) {
 				t.Errorf("abv successfully set bad index %d ", k)
 			}
-		}
-		if _, err := bv.Get(101); err == nil {
-			t.Errorf("bv error should be thrown for out of bounds access")
-		}
-		if _, err := abv.Get(101); err == nil {
-			t.Errorf("abv error should be thrown for out of bounds access")
+			if bv2.Set(k) == nil {
+				t.Errorf("bv2 successfully set bad index %d ", k)
+			}
+			if abv2.Set(k) == nil {
+				t.Errorf("abv2 successfully set bad index %d ", k)
+			}
+			if _, err := bv.Get(k); err == nil {
+				t.Errorf("bv error should be thrown for out of bounds access")
+			}
+			if _, err := abv.Get(k); err == nil {
+				t.Errorf("abv error should be thrown for out of bounds access")
+			}
+			if bv.Clear(k) == nil {
+				t.Errorf("bv error should be thrown for out of bounds access")
+			}
 		}
 
 	})
@@ -63,22 +91,22 @@ func TestBV(t *testing.T) {
 }
 
 var r = rand.New(rand.NewSource(99))
-var sets = r.Perm(1024 * 1024)
+var sets = r.Perm(setSizes[len(setSizes)-1])
 
 func BenchmarkSet(b *testing.B) {
 	for _, setSize := range setSizes {
 		b.Run(fmt.Sprintf("bitvec/Set size %d", setSize), func(b *testing.B) {
 			bv := New(uint64(setSize))
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				bv.TrySet(uint64(sets[i%len(sets)] % setSize))
+				bv.Set(uint64(sets[i%len(sets)] % setSize))
 			}
 		})
 		b.Run(fmt.Sprintf("abitvec/Set size %d", setSize), func(b *testing.B) {
 			abv := NewAtomic(uint64(setSize))
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				abv.TrySet(uint64(sets[i%len(sets)] % setSize))
+				abv.Set(uint64(sets[i%len(sets)] % setSize))
 			}
 		})
 		b.Run(fmt.Sprintf("slice/Set size %d", setSize), func(b *testing.B) {
@@ -89,8 +117,28 @@ func BenchmarkSet(b *testing.B) {
 			}
 		})
 	}
+	fmt.Println()
 }
 
+func BenchmarkTrySet(b *testing.B) {
+	for _, setSize := range setSizes {
+		b.Run(fmt.Sprintf("bitvec/TrySet size %d", setSize), func(b *testing.B) {
+			bv := New(uint64(setSize))
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				bv.TrySet(uint64(sets[i%len(sets)] % setSize))
+			}
+		})
+		b.Run(fmt.Sprintf("abitvec/TrySet size %d", setSize), func(b *testing.B) {
+			abv := NewAtomic(uint64(setSize))
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				abv.TrySet(uint64(sets[i%len(sets)] % setSize))
+			}
+		})
+	}
+	fmt.Println()
+}
 func BenchmarkGet(b *testing.B) {
 	for _, setSize := range setSizes {
 		b.Run(fmt.Sprintf("bitvec/Get size %d", setSize), func(b *testing.B) {
@@ -124,4 +172,40 @@ func BenchmarkGet(b *testing.B) {
 			}
 		})
 	}
+	fmt.Println()
+}
+func BenchmarkClear(b *testing.B) {
+	for _, setSize := range setSizes {
+		b.Run(fmt.Sprintf("bitvec/Clear size %d", setSize), func(b *testing.B) {
+			bv := New(uint64(setSize))
+			for i := 0; i < b.N; i++ {
+				bv.Set(uint64(sets[i%len(sets)] % setSize))
+			}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				bv.Clear(uint64(sets[i%len(sets)] % setSize))
+			}
+		})
+		b.Run(fmt.Sprintf("abitvec/Clear size %d", setSize), func(b *testing.B) {
+			abv := NewAtomic(uint64(setSize))
+			for i := 0; i < b.N; i++ {
+				abv.Set(uint64(sets[i%len(sets)] % setSize))
+			}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				abv.Clear(uint64(sets[i%len(sets)] % setSize))
+			}
+		})
+		b.Run(fmt.Sprintf("slice/Clear size %d", setSize), func(b *testing.B) {
+			slice := make([]bool, setSize)
+			for i := 0; i < b.N; i++ {
+				slice[sets[i%len(sets)]%setSize] = true
+			}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				slice[sets[i%len(sets)]%setSize] = false
+			}
+		})
+	}
+	fmt.Println()
 }