Skip to content

Commit 1da6e7c

Browse files
s1cks1ck
committed
Implement clear at index
1 parent 1fd4af2 commit 1da6e7c

File tree

2 files changed

+68
-57
lines changed

2 files changed

+68
-57
lines changed

‎core/src/main/java/org/neo4j/gds/core/utils/paged/HugeAtomicPagedBitSet.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,32 @@ public void forEachSetBit(LongConsumer consumer) {
172172
}
173173
}
174174

175+
public void clear(long index) {
176+
long longIndex = index >>> 6;
177+
int pageIndex = HugeArrays.pageIndex(longIndex, pageShift);
178+
int wordIndex = HugeArrays.indexInPage(longIndex, pageMask);
179+
int bitIndex = (int) (index & BIT_MASK);
180+
181+
var page = getPage(pageIndex);
182+
long bitMask = ~(1L << bitIndex);
183+
184+
long oldWord = page.get(wordIndex);
185+
while (true) {
186+
long newWord = oldWord & bitMask;
187+
if (newWord == oldWord) {
188+
// already cleared
189+
return;
190+
}
191+
long currentWord = page.compareAndExchange(wordIndex, oldWord, newWord);
192+
if (currentWord == oldWord) {
193+
// CAS successful
194+
return;
195+
}
196+
// CAS unsuccessful, try again
197+
oldWord = currentWord;
198+
}
199+
}
200+
175201
public long capacity() {
176202
return pages.get().length() * (1L << pageShift);
177203
}

‎core/src/test/java/org/neo4j/gds/core/utils/paged/HugeAtomicPagedBitSetTest.java‎

Lines changed: 42 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -50,18 +50,12 @@ static Stream<Arguments> bitsets() {
5050
@ParameterizedTest
5151
@MethodSource("bitsets")
5252
void testSet(HugeAtomicPagedBitSet bitSet) {
53-
// page 0
54-
long index = 23;
55-
assertThat(bitSet.get(index)).isFalse();
56-
bitSet.set(index);
57-
assertThat(bitSet.get(index)).isTrue();
58-
// page 1
59-
index = (1L << PAGE_SHIFT_BITS) + 23;
60-
assertThat(bitSet.get(index)).isFalse();
61-
bitSet.set(index);
62-
assertThat(bitSet.get(index)).isTrue();
63-
// page 2
64-
index = 2 * (1L << PAGE_SHIFT_BITS) + 23;
53+
testSetAssert(bitSet, 23); // page 0
54+
testSetAssert(bitSet, (1L << PAGE_SHIFT_BITS) + 23); // page 1
55+
testSetAssert(bitSet, 2 * (1L << PAGE_SHIFT_BITS) + 23); // page 2
56+
}
57+
58+
private static void testSetAssert(HugeAtomicPagedBitSet bitSet, long index) {
6559
assertThat(bitSet.get(index)).isFalse();
6660
bitSet.set(index);
6761
assertThat(bitSet.get(index)).isTrue();
@@ -70,76 +64,67 @@ void testSet(HugeAtomicPagedBitSet bitSet) {
7064
@ParameterizedTest
7165
@MethodSource("bitsets")
7266
void testGetAndSet(HugeAtomicPagedBitSet bitSet) {
73-
// page 0
74-
// getAndSet a bit that is currently false
75-
long index = 23;
76-
assertThat(bitSet.get(index)).isFalse();
77-
assertThat(bitSet.getAndSet(index)).isFalse();
78-
assertThat(bitSet.get(index)).isTrue();
79-
// getAndSet a bit that is currently true
80-
index = 42;
81-
bitSet.set(index);
82-
assertThat(bitSet.get(index)).isTrue();
83-
assertThat(bitSet.getAndSet(index)).isTrue();
84-
assertThat(bitSet.get(index)).isTrue();
85-
// page 1
86-
// getAndSet a bit that is currently false
87-
index = (1L << PAGE_SHIFT_BITS) + 23;
88-
assertThat(bitSet.get(index)).isFalse();
89-
assertThat(bitSet.getAndSet(index)).isFalse();
90-
assertThat(bitSet.get(index)).isTrue();
91-
// getAndSet a bit that is currently true
92-
index = (1L << PAGE_SHIFT_BITS) + 42;
93-
bitSet.set(index);
94-
assertThat(bitSet.get(index)).isTrue();
95-
assertThat(bitSet.getAndSet(index)).isTrue();
96-
assertThat(bitSet.get(index)).isTrue();
97-
// page 2
67+
testGetAndSetAssert(bitSet, 23, 42); // page 0
68+
testGetAndSetAssert(bitSet, (1L << PAGE_SHIFT_BITS) + 23, (1L << PAGE_SHIFT_BITS) + 42); // page 1
69+
testGetAndSetAssert(bitSet, 2 * (1L << PAGE_SHIFT_BITS) + 23, 2 * (1L << PAGE_SHIFT_BITS) + 42); // page 2
70+
}
71+
72+
private static void testGetAndSetAssert(HugeAtomicPagedBitSet bitSet, long idx0, long idx1) {
9873
// getAndSet a bit that is currently false
99-
index = 2 * (1L << PAGE_SHIFT_BITS) + 23;
100-
assertThat(bitSet.get(index)).isFalse();
101-
assertThat(bitSet.getAndSet(index)).isFalse();
102-
assertThat(bitSet.get(index)).isTrue();
74+
assertThat(bitSet.get(idx0)).isFalse();
75+
assertThat(bitSet.getAndSet(idx0)).isFalse();
76+
assertThat(bitSet.get(idx0)).isTrue();
10377
// getAndSet a bit that is currently true
104-
index = 2 * (1L << PAGE_SHIFT_BITS) + 42;
105-
bitSet.set(index);
106-
assertThat(bitSet.get(index)).isTrue();
107-
assertThat(bitSet.getAndSet(index)).isTrue();
108-
assertThat(bitSet.get(index)).isTrue();
78+
bitSet.set(idx1);
79+
assertThat(bitSet.get(idx1)).isTrue();
80+
assertThat(bitSet.getAndSet(idx1)).isTrue();
81+
assertThat(bitSet.get(idx1)).isTrue();
10982
}
11083

11184
@ParameterizedTest
11285
@MethodSource("bitsets")
11386
void testCardinality(HugeAtomicPagedBitSet bitSet) {
11487
assertThat(bitSet.cardinality()).isEqualTo(0);
115-
// page 0
116-
bitSet.set(23);
88+
bitSet.set(23); // page 0
11789
assertThat(bitSet.cardinality()).isEqualTo(1);
118-
// page 1
119-
bitSet.set((1L << PAGE_SHIFT_BITS) + 23);
90+
bitSet.set((1L << PAGE_SHIFT_BITS) + 23); // page 1
12091
assertThat(bitSet.cardinality()).isEqualTo(2);
121-
// page 2
122-
bitSet.set(2 * (1L << PAGE_SHIFT_BITS) + 23);
92+
bitSet.set(2 * (1L << PAGE_SHIFT_BITS) + 23); // page 2
12393
assertThat(bitSet.cardinality()).isEqualTo(3);
12494
}
12595

96+
@ParameterizedTest
97+
@MethodSource("bitsets")
98+
void testClearAtIndex(HugeAtomicPagedBitSet bitSet) {
99+
assertThat(bitSet.cardinality()).isEqualTo(0);
100+
testClearAtIndexAssert(bitSet, 23); // page 0
101+
testClearAtIndexAssert(bitSet, (1L << PAGE_SHIFT_BITS) + 23); // page 1
102+
testClearAtIndexAssert(bitSet, 2 * (1L << PAGE_SHIFT_BITS) + 23); // page 2
103+
}
104+
105+
private static void testClearAtIndexAssert(HugeAtomicPagedBitSet bitSet, long index) {
106+
bitSet.set(index);
107+
assertThat(bitSet.cardinality()).isEqualTo(1);
108+
bitSet.clear(index);
109+
assertThat(bitSet.cardinality()).isEqualTo(0);
110+
}
111+
126112
@ParameterizedTest
127113
@MethodSource("bitsets")
128114
void testForEachSetBit(HugeAtomicPagedBitSet bitSet) {
129115
var rng = ThreadLocalRandom.current();
130116
long bound = 42 * (1L << PAGE_SHIFT_BITS) + 42;
131117

132-
var expected = IntStream
133-
.range(0, 10_000)
118+
var expectedSetBits = IntStream.range(0, 10_000)
134119
.mapToLong(__ -> rng.nextLong(0, bound))
135120
.boxed()
136121
.peek(bitSet::set)
137122
.collect(Collectors.toSet());
138123

139-
var actual = new HashSet<Long>();
140-
bitSet.forEachSetBit(actual::add);
124+
var actualSetBits = new HashSet<Long>();
125+
bitSet.forEachSetBit(actualSetBits::add);
141126

142-
assertThat(actual).isEqualTo(expected);
127+
assertThat(actualSetBits).isEqualTo(expectedSetBits);
143128
}
144129

145130
@ParameterizedTest

0 commit comments

Comments
 (0)