[libc][NFC] Clean up the memory buffer handling for RPC

We do a lot of arithmetic on void pointers here, so include a helper and
make some more consistent names. Changes no functionality.

Reviewed By: JonChesterfield

Differential Revision: https://reviews.llvm.org/D150576

2 files changed, 43 insertions, 38 deletions

diff --git a/libc/src/__support/RPC/rpc.h b/libc/src/__support/RPC/rpc.h
index 7f61f2588587..540adcebd76b 100644
--- a/libc/src/__support/RPC/rpc.h
+++ b/libc/src/__support/RPC/rpc.h
@@ -115,39 +115,37 @@ template <bool InvertInbox> struct Process {
   cpp::Atomic<uint32_t> lock[DEFAULT_PORT_COUNT] = {0};
 
   /// Initialize the communication channels.
-  LIBC_INLINE void reset(uint64_t port_count, uint32_t lane_size, void *state) {
-    uint64_t p = memory_offset_primary_mailbox(port_count);
-    uint64_t s = memory_offset_secondary_mailbox(port_count);
+  LIBC_INLINE void reset(uint64_t port_count, uint32_t lane_size,
+                         void *buffer) {
     this->port_count = port_count;
     this->lane_size = lane_size;
     this->inbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(
-        static_cast<char *>(state) + (InvertInbox ? s : p));
+        advance(buffer, inbox_offset(port_count)));
     this->outbox = reinterpret_cast<cpp::Atomic<uint32_t> *>(
-        static_cast<char *>(state) + (InvertInbox ? p : s));
-    this->packet = reinterpret_cast<Packet *>(static_cast<char *>(state) +
-                                              memory_offset_buffer(port_count));
+        advance(buffer, outbox_offset(port_count)));
+    this->packet =
+        reinterpret_cast<Packet *>(advance(buffer, buffer_offset(port_count)));
   }
 
-  /// Allocate a single block of memory for use by client and server
-  /// template<size_t N>, N is generally a runtime value
-  /// struct equivalent {
-  ///   atomic<uint32_t> primary[N];
-  ///   atomic<uint32_t> secondary[N];
-  ///   Packet buffer[N];
+  /// Allocate a memory buffer sufficient to store the following equivalent
+  /// representation in memory.
+  ///
+  /// struct Equivalent {
+  ///   Atomic<uint32_t> primary[port_count];
+  ///   Atomic<uint32_t> secondary[port_count];
+  ///   Packet buffer[port_count];
   /// };
   LIBC_INLINE static uint64_t allocation_size(uint64_t port_count,
                                               uint32_t lane_size) {
-    return memory_offset_buffer(port_count) +
-           memory_allocated_buffer(port_count, lane_size);
+    return buffer_offset(port_count) + buffer_bytes(port_count, lane_size);
   }
 
   /// The length of the packet is flexible because the server needs to look up
   /// the lane size at runtime. This helper indexes at the proper offset.
   LIBC_INLINE Packet &get_packet(uint64_t index) {
-    return *reinterpret_cast<Packet *>(
-        reinterpret_cast<uint8_t *>(packet) +
-        index * align_up(sizeof(Header) + lane_size * sizeof(Buffer),
-                         alignof(Packet)));
+    return *reinterpret_cast<Packet *>(advance(
+        packet, index * align_up(sizeof(Header) + lane_size * sizeof(Buffer),
+                                 alignof(Packet))));
   }
 
   /// Inverting the bits loaded from the inbox in exactly one of the pair of
@@ -264,32 +262,34 @@ template <bool InvertInbox> struct Process {
     }
   }
 
-  /// Number of bytes allocated for mailbox or buffer
-  LIBC_INLINE static uint64_t memory_allocated_mailbox(uint64_t port_count) {
+  /// Number of bytes to allocate for an inbox or outbox.
+  LIBC_INLINE static uint64_t mailbox_bytes(uint64_t port_count) {
     return port_count * sizeof(cpp::Atomic<uint32_t>);
   }
 
-  LIBC_INLINE static uint64_t memory_allocated_buffer(uint64_t port_count,
-                                                      uint32_t lane_size) {
-#if defined(LIBC_TARGET_ARCH_IS_GPU)
-    (void)lane_size;
-    return port_count * sizeof(Packet);
-#else
-    return port_count * (sizeof(Packet) + sizeof(Buffer) * lane_size);
-#endif
+  /// Number of bytes to allocate for the buffer containing the packets.
+  LIBC_INLINE static uint64_t buffer_bytes(uint64_t port_count,
+                                           uint32_t lane_size) {
+    return is_process_gpu()
+               ? port_count * sizeof(Packet)
+               : port_count *
+                     align_up(sizeof(Header) + (lane_size * sizeof(Buffer)),
+                              alignof(Packet));
   }
 
-  /// Offset of mailbox/buffer in single allocation
-  LIBC_INLINE static uint64_t
-  memory_offset_primary_mailbox(uint64_t /*port_count*/) {
-    return 0;
+  /// Offset of the inbox in memory. This is the same as the outbox if inverted.
+  LIBC_INLINE static uint64_t inbox_offset(uint64_t port_count) {
+    return InvertInbox ? mailbox_bytes(port_count) : 0;
   }
-  LIBC_INLINE static uint64_t
-  memory_offset_secondary_mailbox(uint64_t port_count) {
-    return memory_allocated_mailbox(port_count);
+
+  /// Offset of the outbox in memory. This is the same as the inbox if inverted.
+  LIBC_INLINE static uint64_t outbox_offset(uint64_t port_count) {
+    return InvertInbox ? 0 : mailbox_bytes(port_count);
   }
-  LIBC_INLINE static uint64_t memory_offset_buffer(uint64_t port_count) {
-    return align_up(2 * memory_allocated_mailbox(port_count), alignof(Packet));
+
+  /// Offset of the buffer containing the packets after the inbox and outbox.
+  LIBC_INLINE static uint64_t buffer_offset(uint64_t port_count) {
+    return align_up(2 * mailbox_bytes(port_count), alignof(Packet));
   }
 };
 
diff --git a/libc/src/__support/RPC/rpc_util.h b/libc/src/__support/RPC/rpc_util.h
index b2b8cfeaa5c3..b9ffdaa08969 100644
--- a/libc/src/__support/RPC/rpc_util.h
+++ b/libc/src/__support/RPC/rpc_util.h
@@ -69,6 +69,11 @@ template <typename T> LIBC_INLINE const T &max(const T &x, const T &y) {
   return x < y ? y : x;
 }
 
+/// Advance the \p ptr by \p bytes.
+template <typename T, typename U> LIBC_INLINE T *advance(T ptr, U bytes) {
+  return reinterpret_cast<T *>(reinterpret_cast<uint8_t *>(ptr) + bytes);
+}
+
 } // namespace rpc
 } // namespace __llvm_libc