Matt Jeanneret | cab955f | 2019-04-10 15:45:57 -0400 | [diff] [blame] | 1 | package bitmap |
| 2 | |
| 3 | import ( |
| 4 | "sync/atomic" |
| 5 | "unsafe" |
| 6 | ) |
| 7 | |
| 8 | var oobPanic = "SetAtomic not allowed on a bitmapSlice of cap() < 4" |
| 9 | |
| 10 | // SetAtomic is similar to Set except that it performs the operation atomically. |
| 11 | func SetAtomic(bitmap []byte, targetBit int, targetValue bool) { |
| 12 | ov := (*[1]uint32)(unsafe.Pointer(&bitmap[targetBit/32]))[:] |
| 13 | SetAtomicUint32(ov, targetBit%32, targetValue) |
| 14 | } |
| 15 | |
| 16 | // SetAtomic is similar to Set except that it performs the operation atomically. |
| 17 | // It needs a bitmapSlice where the capacity is at least 4 bytes. |
| 18 | func _SetAtomic(bitmapSlice []byte, targetBit int, targetValue bool) { |
| 19 | targetByteIndex := targetBit / 8 |
| 20 | targetBitIndex := targetBit % 8 |
| 21 | targetOffset := 0 |
| 22 | |
| 23 | // SetAtomic needs to modify 4 bytes of data so we panic when the slice |
| 24 | // doesn't have a capacity of at least 4 bytes. |
| 25 | if cap(bitmapSlice) < 4 { |
| 26 | panic(oobPanic) |
| 27 | } |
| 28 | |
| 29 | // Calculate the Offset of the targetByte inside the 4-byte atomic batch. |
| 30 | // This is needed to ensure that atomic operations can happen as long as |
| 31 | // the bitmapSlice equals 4 bytes or more. |
| 32 | if cap(bitmapSlice) < targetByteIndex+3 { |
| 33 | targetOffset = cap(bitmapSlice) - targetByteIndex |
| 34 | } |
| 35 | |
| 36 | // This gets a pointer to the memory of 4 bytes inside the bitmapSlice. |
| 37 | // It stores this pointer as an *uint32 so that it can be used to |
| 38 | // execute sync.atomic operations. |
| 39 | targetBytePointer := (*uint32)(unsafe.Pointer(&bitmapSlice[targetByteIndex-targetOffset])) |
| 40 | |
| 41 | for { |
| 42 | // localValue is a copy of the uint32 value at *targetBytePointer. |
| 43 | // It's used to check whether the targetBit must be updated, |
| 44 | // and if so, to construct the new value for targetBytePointer. |
| 45 | localValue := atomic.LoadUint32(targetBytePointer) |
| 46 | |
| 47 | // This "neutralizes" the uint32 conversion by getting a pointer to the |
| 48 | // 4-byte array stored undereneath the uint32. |
| 49 | targetByteCopyPointer := (*[4]byte)(unsafe.Pointer(&localValue)) |
| 50 | |
| 51 | // Work is done when targetBit is already set to targetValue. |
| 52 | if GetBit(targetByteCopyPointer[targetOffset], targetBitIndex) == targetValue { |
| 53 | return |
| 54 | } |
| 55 | |
| 56 | // Modify the targetBit and update memory so that the targetBit is the only bit |
| 57 | // that has been modified in the batch. |
| 58 | referenceValue := localValue |
| 59 | SetBitRef(&targetByteCopyPointer[targetOffset], targetBitIndex, targetValue) |
| 60 | if atomic.CompareAndSwapUint32(targetBytePointer, referenceValue, localValue) { |
| 61 | break |
| 62 | } |
| 63 | } |
| 64 | } |