Merge pull request #177 from CannibalVox/master

add Word method with tests

bitset.go

@@ -205,6 +205,26 @@ func (b *BitSet) Test(i uint) bool {
 	return b.set[i>>log2WordSize]&(1<<wordsIndex(i)) != 0
 }
 
+// GetWord64AtBit retrieves bits i through i+63 as a single uint64 value
+func (b *BitSet) GetWord64AtBit(i uint) uint64 {
+	firstWordIndex := int(i >> log2WordSize)
+	subWordIndex := wordsIndex(i)
+
+	// The word that the index falls within, shifted so the index is at bit 0
+	var firstWord, secondWord uint64
+	if firstWordIndex < len(b.set) {
+		firstWord = b.set[firstWordIndex] >> subWordIndex
+	}
+
+	// The next word, masked to only include the necessary bits and shifted to cover the
+	// top of the word
+	if (firstWordIndex + 1) < len(b.set) {
+		secondWord = b.set[firstWordIndex+1] << uint64(wordSize-subWordIndex)
+	}
+
+	return firstWord | secondWord
+}
+
 // Set bit i to 1, the capacity of the bitset is automatically
 // increased accordingly.
 // Warning: using a very large value for 'i'

bitset_test.go

@@ -2035,3 +2035,55 @@ func TestShiftRight(t *testing.T) {
 	test("full shift", 89)
 	test("remove all", 242)
 }
+
+func TestWord(t *testing.T) {
+	data := []uint64{0x0bfd85fc01af96dd, 0x3fe212a7eae11414, 0x7aa412221245dee1, 0x557092c1711306d5}
+	testCases := map[string]struct {
+		index    uint
+		expected uint64
+	}{
+		"first word": {
+			index:    0,
+			expected: 0x0bfd85fc01af96dd,
+		},
+		"third word": {
+			index:    128,
+			expected: 0x7aa412221245dee1,
+		},
+		"off the edge": {
+			index:    256,
+			expected: 0,
+		},
+		"way off the edge": {
+			index:    6346235235,
+			expected: 0,
+		},
+		"split between two words": {
+			index:    96,
+			expected: 0x1245dee13fe212a7,
+		},
+		"partly off edge": {
+			index:    254,
+			expected: 1,
+		},
+		"skip one nibble": {
+			index:    4,
+			expected: 0x40bfd85fc01af96d,
+		},
+		"slightly offset results": {
+			index:    65,
+			expected: 0x9ff10953f5708a0a,
+		},
+	}
+
+	for name, testCase := range testCases {
+		t.Run(name, func(t *testing.T) {
+			bitSet := From(data)
+			output := bitSet.GetWord64AtBit(testCase.index)
+
+			if output != testCase.expected {
+				t.Errorf("Word should have returned %d for input %d, but returned %d", testCase.expected, testCase.index, output)
+			}
+		})
+	}
+}