Use atomic data access for a cocurrent pointer read/write

Author: rintaro

Sources/SwiftSyntax/Syntax.swift

@@ -393,24 +393,27 @@ final class SyntaxDataArena: @unchecked Sendable {
     // The storage of the buffer address is allocated next to the SyntaxData.
     let baseAddressRef = parent.advanced(by: 1)
       .unsafeRawPointer
-      .assumingMemoryBound(to: UnsafePointer<SyntaxDataReference?>?.self)
+      .assumingMemoryBound(to: AtomicPointer.self)
 
     // If the buffer is already populated, return it.
-    if let baseAddress = baseAddressRef.pointee {
+    if let baseAddress = swiftsyntax_atomic_pointer_get(baseAddressRef)?.assumingMemoryBound(to: SyntaxDataReference?.self) {
       return SyntaxDataReferenceBuffer(UnsafeBufferPointer(start: baseAddress, count: childCount))
     }
 
     mutex.lock()
     defer { mutex.unlock() }
 
     // Recheck, maybe some other thread has populated the buffer during acquiring the lock.
-    if let baseAddress = baseAddressRef.pointee {
+    if let baseAddress = swiftsyntax_atomic_pointer_get(baseAddressRef)?.assumingMemoryBound(to: SyntaxDataReference?.self) {
       return SyntaxDataReferenceBuffer(UnsafeBufferPointer(start: baseAddress, count: childCount))
     }
 
     let buffer = createLayoutDataImpl(parent)
     // Remeber the base address of the created buffer.
-    UnsafeMutablePointer(mutating: baseAddressRef).pointee = buffer.baseAddress
+    swiftsyntax_atomic_pointer_set(
+      UnsafeMutablePointer(mutating: baseAddressRef),
+      UnsafeRawPointer(buffer.baseAddress)
+    )
 
     return SyntaxDataReferenceBuffer(buffer)
   }
@@ -460,7 +463,7 @@ final class SyntaxDataArena: @unchecked Sendable {
     return min(totalSize, 4096)
   }
 
-  /// Allocate and initialize `SyntaxData` with the trailing uninitialized buffer for the references to the children.
+  /// Allocate and initialize `SyntaxData` with the trailing pointer storage for the references to the children.
   private static func createDataImpl(
     allocator: BumpPtrAllocator,
     raw: RawSyntax,
@@ -474,7 +477,7 @@ final class SyntaxDataArena: @unchecked Sendable {
     var totalSize = MemoryLayout<SyntaxData>.stride
     if childCount != 0 {
       // Tail allocate the storage for the pointer to the lazily allocated layout data.
-      totalSize &+= MemoryLayout<UnsafePointer<SyntaxDataReference?>?>.size
+      totalSize &+= MemoryLayout<AtomicPointer>.size
     }
     let alignment = MemoryLayout<SyntaxData>.alignment
     let allocated = allocator.allocate(byteCount: totalSize, alignment: alignment).baseAddress!
@@ -492,10 +495,10 @@ final class SyntaxDataArena: @unchecked Sendable {
 
     if childCount != 0 {
       // Initlaize the tail allocated storage with 'nil'.
-      allocated
+      let layoutBufferBaseAddressRef = allocated
         .advanced(by: MemoryLayout<SyntaxData>.stride)
-        .bindMemory(to: UnsafePointer<SyntaxDataReference?>?.self, capacity: 1)
-        .initialize(to: nil)
+        .bindMemory(to: AtomicPointer.self, capacity: 1)
+      swiftsyntax_atomic_pointer_set(layoutBufferBaseAddressRef, nil)
     }
 
     return SyntaxDataReference(UnsafePointer(dataRef))

Sources/_SwiftSyntaxCShims/include/AtomicBool.h renamed to Sources/_SwiftSyntaxCShims/include/Atomics.h

@@ -38,4 +38,16 @@ static inline void swiftsyntax_atomic_bool_destroy(AtomicBool *_Nonnull atomic)
   free(atomic);
 }
 
+typedef struct {
+  _Atomic(const void *_Nullable) value;
+} AtomicPointer;
+
+static inline const void *_Nullable swiftsyntax_atomic_pointer_get(const AtomicPointer *_Nonnull atomic) {
+  return atomic->value;
+}
+
+static inline void swiftsyntax_atomic_pointer_set(AtomicPointer *_Nonnull atomic, const void *_Nullable newValue) {
+  atomic->value = newValue;
+}
+
 #endif // SWIFTSYNTAX_ATOMICBOOL_H

Sources/_SwiftSyntaxCShims/include/SwiftSyntaxCShims.h

@@ -11,7 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "_includes.h"
-#include "AtomicBool.h"
+#include "Atomics.h"
 #include "PlatformMutex.h"
 #include "swiftsyntax_errno.h"
 #include "swiftsyntax_stdio.h"

Sources/_SwiftSyntaxCShims/include/module.modulemap

@@ -1,6 +1,6 @@
 module _SwiftSyntaxCShims {
   header "_includes.h"
-  header "AtomicBool.h"
+  header "Atomics.h"
   header "PlatformMutex.h"
   header "swiftsyntax_errno.h"
   header "swiftsyntax_stdio.h"