fix(avm): avoid data race on shared interaction selector writes (port #24456 to v5-next)

barretenberg/cpp/src/barretenberg/vm2/tracegen/lib/lookup_builder.hpp

@@ -51,10 +51,10 @@ template <typename LookupSettings_> class IndexedLookupTraceBuilder : public Int
             }
 
             trace.set(LookupSettings::COUNTS, dst_row, trace.get(LookupSettings::COUNTS, dst_row) + 1);
-            // Set the fine grained inner selector if it's not already one.
-            if (LookupSettings::DST_SELECTOR != this->outer_dst_selector &&
-                trace.get(LookupSettings::DST_SELECTOR, dst_row) != 1) {
-                trace.set(LookupSettings::DST_SELECTOR, dst_row, 1);
+            if (LookupSettings::DST_SELECTOR != this->outer_dst_selector) {
+                // This step might write to the same cell from multiple threads, so we use atomic limbs to avoid UB.
+                // Since we are always writing a 1, the end result will be 1 even under concurrency.
+                trace.set(LookupSettings::DST_SELECTOR, dst_row, 1, /*use_atomic_limbs=*/true);
             }
         });
     }
@@ -198,10 +198,10 @@ template <typename LookupSettings> class LookupIntoDynamicTableSequential : publ
                 if (dst_selector == 1 && src_values == trace.get_multiple(LookupSettings::DST_COLUMNS, dst_row)) {
                     trace.set(LookupSettings::COUNTS, dst_row, trace.get(LookupSettings::COUNTS, dst_row) + 1);
 
-                    // Set the fine grained inner selector if it's not already one.
-                    if (LookupSettings::DST_SELECTOR != this->outer_dst_selector &&
-                        trace.get(LookupSettings::DST_SELECTOR, dst_row) != 1) {
-                        trace.set(LookupSettings::DST_SELECTOR, dst_row, 1);
+                    if (LookupSettings::DST_SELECTOR != this->outer_dst_selector) {
+                        // This step might write to the same cell from multiple threads, so we use atomic limbs to avoid
+                        // UB. Since we are always writing a 1, the end result will be 1 even under concurrency.
+                        trace.set(LookupSettings::DST_SELECTOR, dst_row, 1, /*use_atomic_limbs=*/true);
                     }
 
                     found = true;

barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.cpp

@@ -4,6 +4,7 @@
 #include <ranges>
 
 #include "barretenberg/common/assert.hpp"
+#include "barretenberg/common/compiler_hints.hpp"
 #include "barretenberg/common/log.hpp"
 #include "barretenberg/common/ref_vector.hpp"
 #include "barretenberg/vm2/common/field.hpp"
@@ -13,7 +14,19 @@ namespace bb::avm2::tracegen {
 namespace {
 
 // We need a zero value to return (a reference to) when a value is not found.
-static const FF zero = FF::zero();
+const FF zero = FF::zero();
+
+// Writes each 64-bit limb of the field with a relaxed atomic store. Each limb is naturally aligned (FF is
+// alignas(32), 4x uint64_t), so this lowers to 4 plain `movq` stores on x86-64 — no lock, no libatomic call
+// (unlike a whole-field std::atomic_ref<FF>, whose 32 bytes exceed the lock-free width). Lets set() make a
+// same-cell concurrent write data-race-free when every writer stores the same value.
+inline void store_per_limb(FF& cell, const FF& value)
+{
+    static_assert(sizeof(FF) == 4 * sizeof(uint64_t));
+    for (size_t i = 0; i < 4; ++i) {
+        std::atomic_ref<uint64_t>(cell.data[i]).store(value.data[i], std::memory_order_relaxed);
+    }
+}
 
 } // namespace
 
@@ -62,7 +75,7 @@ TraceContainer::ColumnInterval& TraceContainer::get_or_create_shard(SparseColumn
     return *expected; // CAS failure loaded the winning pointer into `expected` (acquire).
 }
 
-void TraceContainer::set(Column col, uint32_t row, const FF& value)
+void TraceContainer::set(Column col, uint32_t row, const FF& value, bool use_atomic_limbs)
 {
     auto& column_data = (*trace)[static_cast<size_t>(col)];
     const size_t shard_idx = row / INTERVAL_SIZE;
@@ -73,12 +86,22 @@ void TraceContainer::set(Column col, uint32_t row, const FF& value)
         // Lock-free: a single atomic load finds the shard (created on first write), then we write our
         // own dense cell directly. Different rows are distinct array elements, so concurrent writers of
         // this column (or even of the same shard, at a chunk boundary) never race and never serialize.
-        get_or_create_shard(column_data, shard_idx).rows[offset] = value;
+        auto& cell = get_or_create_shard(column_data, shard_idx).rows[offset];
+        if (BB_UNLIKELY(use_atomic_limbs)) {
+            store_per_limb(cell, value);
+        } else {
+            cell = value;
+        }
     } else {
-        // Zero value: clear if present. We never create a shard (clearing an absent row is a no-op).
+        // Zero value: clear if present. We never create a shard, so sparse (mostly-zero) columns are not
+        // materialized (an unset cell already reads as zero).
         ColumnInterval* shard = column_data.slots[shard_idx].load(std::memory_order_acquire);
         if (shard != nullptr) {
-            shard->rows[offset] = FF::zero();
+            if (BB_UNLIKELY(use_atomic_limbs)) {
+                store_per_limb(shard->rows[offset], zero);
+            } else {
+                shard->rows[offset] = FF::zero();
+            }
         }
     }
 }

barretenberg/cpp/src/barretenberg/vm2/tracegen/trace_container.hpp

@@ -76,7 +76,11 @@ class TraceContainer {
     // Extended version of get that works with shifted columns. More expensive.
     const FF& get_column_or_shift(ColumnAndShifts col, uint32_t row) const;
 
-    void set(Column col, uint32_t row, const FF& value);
+    // Sets the value of a cell. Thread-safe if the same cell is not written to from multiple threads.
+    // If writing to the same cell from multiple threads, use_atomic_limbs=true to use atomic limbs.
+    // This makes the write slower, but it will not be UB. However, it is also not thread-safe.
+    // Use only if you know what you are doing.
+    void set(Column col, uint32_t row, const FF& value, bool use_atomic_limbs = false);
     // Bulk setting for a given row.
     void set(uint32_t row, std::span<const std::pair<Column, FF>> values);
     // Reserve column size. Useful for precomputed columns.