Akash Reddy Kankanala | 92dfdf8 | 2025-03-23 22:07:09 +0530 | [diff] [blame^] | 1 | // Copyright 2024 The Go Authors. All rights reserved. |
| 2 | // Use of this source code is governed by a BSD-style |
| 3 | // license that can be found in the LICENSE file. |
| 4 | |
| 5 | package impl |
| 6 | |
| 7 | import ( |
| 8 | "sync/atomic" |
| 9 | "unsafe" |
| 10 | ) |
| 11 | |
| 12 | // presenceSize represents the size of a presence set, which should be the largest index of the set+1 |
| 13 | type presenceSize uint32 |
| 14 | |
| 15 | // presence is the internal representation of the bitmap array in a generated protobuf |
| 16 | type presence struct { |
| 17 | // This is a pointer to the beginning of an array of uint32 |
| 18 | P unsafe.Pointer |
| 19 | } |
| 20 | |
| 21 | func (p presence) toElem(num uint32) (ret *uint32) { |
| 22 | const ( |
| 23 | bitsPerByte = 8 |
| 24 | siz = unsafe.Sizeof(*ret) |
| 25 | ) |
| 26 | // p.P points to an array of uint32, num is the bit in this array that the |
| 27 | // caller wants to check/manipulate. Calculate the index in the array that |
| 28 | // contains this specific bit. E.g.: 76 / 32 = 2 (integer division). |
| 29 | offset := uintptr(num) / (siz * bitsPerByte) * siz |
| 30 | return (*uint32)(unsafe.Pointer(uintptr(p.P) + offset)) |
| 31 | } |
| 32 | |
| 33 | // Present checks for the presence of a specific field number in a presence set. |
| 34 | func (p presence) Present(num uint32) bool { |
| 35 | if p.P == nil { |
| 36 | return false |
| 37 | } |
| 38 | return Export{}.Present(p.toElem(num), num) |
| 39 | } |
| 40 | |
| 41 | // SetPresent adds presence for a specific field number in a presence set. |
| 42 | func (p presence) SetPresent(num uint32, size presenceSize) { |
| 43 | Export{}.SetPresent(p.toElem(num), num, uint32(size)) |
| 44 | } |
| 45 | |
| 46 | // SetPresentUnatomic adds presence for a specific field number in a presence set without using |
| 47 | // atomic operations. Only to be called during unmarshaling. |
| 48 | func (p presence) SetPresentUnatomic(num uint32, size presenceSize) { |
| 49 | Export{}.SetPresentNonAtomic(p.toElem(num), num, uint32(size)) |
| 50 | } |
| 51 | |
| 52 | // ClearPresent removes presence for a specific field number in a presence set. |
| 53 | func (p presence) ClearPresent(num uint32) { |
| 54 | Export{}.ClearPresent(p.toElem(num), num) |
| 55 | } |
| 56 | |
| 57 | // LoadPresenceCache (together with PresentInCache) allows for a |
| 58 | // cached version of checking for presence without re-reading the word |
| 59 | // for every field. It is optimized for efficiency and assumes no |
| 60 | // simltaneous mutation of the presence set (or at least does not have |
| 61 | // a problem with simultaneous mutation giving inconsistent results). |
| 62 | func (p presence) LoadPresenceCache() (current uint32) { |
| 63 | if p.P == nil { |
| 64 | return 0 |
| 65 | } |
| 66 | return atomic.LoadUint32((*uint32)(p.P)) |
| 67 | } |
| 68 | |
| 69 | // PresentInCache reads presence from a cached word in the presence |
| 70 | // bitmap. It caches up a new word if the bit is outside the |
| 71 | // word. This is for really fast iteration through bitmaps in cases |
| 72 | // where we either know that the bitmap will not be altered, or we |
| 73 | // don't care about inconsistencies caused by simultaneous writes. |
| 74 | func (p presence) PresentInCache(num uint32, cachedElement *uint32, current *uint32) bool { |
| 75 | if num/32 != *cachedElement { |
| 76 | o := uintptr(num/32) * unsafe.Sizeof(uint32(0)) |
| 77 | q := (*uint32)(unsafe.Pointer(uintptr(p.P) + o)) |
| 78 | *current = atomic.LoadUint32(q) |
| 79 | *cachedElement = num / 32 |
| 80 | } |
| 81 | return (*current & (1 << (num % 32))) > 0 |
| 82 | } |
| 83 | |
| 84 | // AnyPresent checks if any field is marked as present in the bitmap. |
| 85 | func (p presence) AnyPresent(size presenceSize) bool { |
| 86 | n := uintptr((size + 31) / 32) |
| 87 | for j := uintptr(0); j < n; j++ { |
| 88 | o := j * unsafe.Sizeof(uint32(0)) |
| 89 | q := (*uint32)(unsafe.Pointer(uintptr(p.P) + o)) |
| 90 | b := atomic.LoadUint32(q) |
| 91 | if b > 0 { |
| 92 | return true |
| 93 | } |
| 94 | } |
| 95 | return false |
| 96 | } |
| 97 | |
| 98 | // toRaceDetectData finds the preceding RaceDetectHookData in a |
| 99 | // message by using pointer arithmetic. As the type of the presence |
| 100 | // set (bitmap) varies with the number of fields in the protobuf, we |
| 101 | // can not have a struct type containing the array and the |
| 102 | // RaceDetectHookData. instead the RaceDetectHookData is placed |
| 103 | // immediately before the bitmap array, and we find it by walking |
| 104 | // backwards in the struct. |
| 105 | // |
| 106 | // This method is only called from the race-detect version of the code, |
| 107 | // so RaceDetectHookData is never an empty struct. |
| 108 | func (p presence) toRaceDetectData() *RaceDetectHookData { |
| 109 | var template struct { |
| 110 | d RaceDetectHookData |
| 111 | a [1]uint32 |
| 112 | } |
| 113 | o := (uintptr(unsafe.Pointer(&template.a)) - uintptr(unsafe.Pointer(&template.d))) |
| 114 | return (*RaceDetectHookData)(unsafe.Pointer(uintptr(p.P) - o)) |
| 115 | } |
| 116 | |
| 117 | func atomicLoadShadowPresence(p **[]byte) *[]byte { |
| 118 | return (*[]byte)(atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(p)))) |
| 119 | } |
| 120 | func atomicStoreShadowPresence(p **[]byte, v *[]byte) { |
| 121 | atomic.CompareAndSwapPointer((*unsafe.Pointer)(unsafe.Pointer(p)), nil, unsafe.Pointer(v)) |
| 122 | } |
| 123 | |
| 124 | // findPointerToRaceDetectData finds the preceding RaceDetectHookData |
| 125 | // in a message by using pointer arithmetic. For the methods called |
| 126 | // directy from generated code, we don't have a pointer to the |
| 127 | // beginning of the presence set, but a pointer inside the array. As |
| 128 | // we know the index of the bit we're manipulating (num), we can |
| 129 | // calculate which element of the array ptr is pointing to. With that |
| 130 | // information we find the preceding RaceDetectHookData and can |
| 131 | // manipulate the shadow bitmap. |
| 132 | // |
| 133 | // This method is only called from the race-detect version of the |
| 134 | // code, so RaceDetectHookData is never an empty struct. |
| 135 | func findPointerToRaceDetectData(ptr *uint32, num uint32) *RaceDetectHookData { |
| 136 | var template struct { |
| 137 | d RaceDetectHookData |
| 138 | a [1]uint32 |
| 139 | } |
| 140 | o := (uintptr(unsafe.Pointer(&template.a)) - uintptr(unsafe.Pointer(&template.d))) + uintptr(num/32)*unsafe.Sizeof(uint32(0)) |
| 141 | return (*RaceDetectHookData)(unsafe.Pointer(uintptr(unsafe.Pointer(ptr)) - o)) |
| 142 | } |