Skip to content

Commit afe456d

Browse files
s1cks1ck
committed
Implement getAndSet and refactor tests
1 parent 68bbac0 commit afe456d

File tree

2 files changed

+95
-24
lines changed

2 files changed

+95
-24
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
@@ -105,6 +105,32 @@ public boolean get(long index) {
105105
return (page.get(wordIndex) & bitMask) != 0;
106106
}
107107

108+
public boolean getAndSet(long index) {
109+
long longIndex = index >>> 6;
110+
int pageIndex = HugeArrays.pageIndex(longIndex, pageShift);
111+
int wordIndex = HugeArrays.indexInPage(longIndex, pageMask);
112+
int bitIndex = (int) (index & BIT_MASK);
113+
114+
var page = getPage(pageIndex);
115+
long bitMask = 1L << bitIndex;
116+
117+
long oldWord = page.get(wordIndex);
118+
while (true) {
119+
long newWord = oldWord | bitMask;
120+
if (newWord == oldWord) {
121+
// already set
122+
return true;
123+
}
124+
long currentWord = page.compareAndExchange(wordIndex, oldWord, newWord);
125+
if (currentWord == oldWord) {
126+
// CAS successful
127+
return false;
128+
}
129+
// CAS unsuccessful, try again
130+
oldWord = currentWord;
131+
}
132+
}
133+
108134
public long cardinality() {
109135
final Pages pages = this.pages.get();
110136
final long pageCount = pages.length();

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

Lines changed: 69 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -20,49 +20,94 @@
2020
package org.neo4j.gds.core.utils.paged;
2121

2222
import org.junit.jupiter.api.Test;
23+
import org.junit.jupiter.params.ParameterizedTest;
24+
import org.junit.jupiter.params.provider.Arguments;
25+
import org.junit.jupiter.params.provider.MethodSource;
2326
import org.neo4j.gds.core.concurrency.RunWithConcurrency;
2427
import org.neo4j.gds.core.utils.partition.PartitionUtils;
2528

2629
import java.util.List;
2730
import java.util.Optional;
31+
import java.util.stream.Stream;
2832

2933
import static org.assertj.core.api.Assertions.assertThat;
3034
import static org.neo4j.gds.core.utils.paged.HugeAtomicPagedBitSet.PAGE_SHIFT_BITS;
3135

3236
class HugeAtomicPagedBitSetTest {
3337

34-
@Test
35-
void testSinglePageBitSet() {
36-
var atomicBitSet = HugeAtomicPagedBitSet.create(42);
37-
assertThat(atomicBitSet.get(23)).isFalse();
38-
atomicBitSet.set(23);
39-
assertThat(atomicBitSet.get(23)).isTrue();
38+
public static Stream<Arguments> bitsets() {
39+
return Stream.of(
40+
// empty bit set
41+
Arguments.of(HugeAtomicPagedBitSet.create(0)),
42+
// bit set of 3 pages
43+
Arguments.of(HugeAtomicPagedBitSet.create(2 * (1L << PAGE_SHIFT_BITS) + 42))
44+
);
4045
}
4146

42-
@Test
43-
void testMultiPageBitSet() {
44-
long size = 2 * (1L << PAGE_SHIFT_BITS) + 42; // 3 pages
45-
46-
var atomicBitSet = HugeAtomicPagedBitSet.create(size);
47+
@ParameterizedTest
48+
@MethodSource("bitsets")
49+
void testSet(HugeAtomicPagedBitSet atomicBitSet) {
4750
// page 0
48-
assertThat(atomicBitSet.get(23)).isFalse();
49-
atomicBitSet.set(23);
50-
assertThat(atomicBitSet.get(23)).isTrue();
51+
long index = 23;
52+
assertThat(atomicBitSet.get(index)).isFalse();
53+
atomicBitSet.set(index);
54+
assertThat(atomicBitSet.get(index)).isTrue();
5155
// page 1
52-
assertThat(atomicBitSet.get((1L << PAGE_SHIFT_BITS) + 23)).isFalse();
53-
atomicBitSet.set((1L << PAGE_SHIFT_BITS) + 23);
54-
assertThat(atomicBitSet.get((1L << PAGE_SHIFT_BITS) + 23)).isTrue();
56+
index = (1L << PAGE_SHIFT_BITS) + 23;
57+
assertThat(atomicBitSet.get(index)).isFalse();
58+
atomicBitSet.set(index);
59+
assertThat(atomicBitSet.get(index)).isTrue();
5560
// page 2
56-
assertThat(atomicBitSet.get(2 * (1L << PAGE_SHIFT_BITS) + 23)).isFalse();
57-
atomicBitSet.set(2 * (1L << PAGE_SHIFT_BITS) + 23);
58-
assertThat(atomicBitSet.get(2 * (1L << PAGE_SHIFT_BITS) + 23)).isTrue();
61+
index = 2 * (1L << PAGE_SHIFT_BITS) + 23;
62+
assertThat(atomicBitSet.get(index)).isFalse();
63+
atomicBitSet.set(index);
64+
assertThat(atomicBitSet.get(index)).isTrue();
5965
}
6066

61-
@Test
62-
void testCardinality() {
63-
long size = 2 * (1L << PAGE_SHIFT_BITS) + 42; // 3 pages
67+
@ParameterizedTest
68+
@MethodSource("bitsets")
69+
void testGetAndSet(HugeAtomicPagedBitSet atomicBitSet) {
70+
// page 0
71+
// getAndSet a bit that is currently false
72+
long index = 23;
73+
assertThat(atomicBitSet.get(index)).isFalse();
74+
assertThat(atomicBitSet.getAndSet(index)).isFalse();
75+
assertThat(atomicBitSet.get(index)).isTrue();
76+
// getAndSet a bit that is currently true
77+
index = 42;
78+
atomicBitSet.set(index);
79+
assertThat(atomicBitSet.get(index)).isTrue();
80+
assertThat(atomicBitSet.getAndSet(index)).isTrue();
81+
assertThat(atomicBitSet.get(index)).isTrue();
82+
// page 1
83+
// getAndSet a bit that is currently false
84+
index = (1L << PAGE_SHIFT_BITS) + 23;
85+
assertThat(atomicBitSet.get(index)).isFalse();
86+
assertThat(atomicBitSet.getAndSet(index)).isFalse();
87+
assertThat(atomicBitSet.get(index)).isTrue();
88+
// getAndSet a bit that is currently true
89+
index = (1L << PAGE_SHIFT_BITS) + 42;
90+
atomicBitSet.set(index);
91+
assertThat(atomicBitSet.get(index)).isTrue();
92+
assertThat(atomicBitSet.getAndSet(index)).isTrue();
93+
assertThat(atomicBitSet.get(index)).isTrue();
94+
// page 2
95+
// getAndSet a bit that is currently false
96+
index = 2 * (1L << PAGE_SHIFT_BITS) + 23;
97+
assertThat(atomicBitSet.get(index)).isFalse();
98+
assertThat(atomicBitSet.getAndSet(index)).isFalse();
99+
assertThat(atomicBitSet.get(index)).isTrue();
100+
// getAndSet a bit that is currently true
101+
index = 2 * (1L << PAGE_SHIFT_BITS) + 42;
102+
atomicBitSet.set(index);
103+
assertThat(atomicBitSet.get(index)).isTrue();
104+
assertThat(atomicBitSet.getAndSet(index)).isTrue();
105+
assertThat(atomicBitSet.get(index)).isTrue();
106+
}
64107

65-
var atomicBitSet = HugeAtomicPagedBitSet.create(size);
108+
@ParameterizedTest
109+
@MethodSource("bitsets")
110+
void testCardinality(HugeAtomicPagedBitSet atomicBitSet) {
66111
assertThat(atomicBitSet.cardinality()).isEqualTo(0);
67112
// page 0
68113
atomicBitSet.set(23);

0 commit comments

Comments
 (0)