diff --git a/bitset.go b/bitset.go index c860986..cdf91c3 100644 --- a/bitset.go +++ b/bitset.go @@ -205,6 +205,26 @@ func (b *BitSet) Test(i uint) bool { return b.set[i>>log2WordSize]&(1<> 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' diff --git a/bitset_test.go b/bitset_test.go index 4114bec..ca21d8b 100644 --- a/bitset_test.go +++ b/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) + } + }) + } +}