Improve fast-path batch replay handling and harden pairindex slot allocation.

Track and roll back per-batch checkpoints when replaying a failed fast batch, and switch pairindex slot allocation to unsigned atomics to avoid negative modulo indexing after counter wraparound.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: hoffmang9

src/c_bindings/fast_wrapper.cpp

@@ -243,6 +243,9 @@ class StreamingOneWesolowskiCallback final : public WesolowskiCallback {
                 }
                 SetForm(type, data, &checkpoint);
                 process_checkpoint(pos, checkpoint, /*record_stats=*/true);
+                if (iteration >= batch_start_iteration && iteration <= batch_end_iteration) {
+                    current_batch_checkpoints.push_back(BatchCheckpoint{pos, checkpoint});
+                }
                 if (stats_enabled) {
                     checkpoint_event_total_ns += static_cast<uint64_t>(
                         std::chrono::duration_cast<std::chrono::nanoseconds>(
@@ -258,7 +261,44 @@ class StreamingOneWesolowskiCallback final : public WesolowskiCallback {
         }
     }
 
+    void OnBatchStart(uint64_t base_iteration, uint64_t batch_size) override {
+        current_batch_checkpoints.clear();
+        if (batch_size == 0) {
+            batch_start_iteration = 1;
+            batch_end_iteration = 0;
+            return;
+        }
+        batch_start_iteration = base_iteration + 1;
+        if (std::numeric_limits<uint64_t>::max() - base_iteration < batch_size) {
+            batch_end_iteration = std::numeric_limits<uint64_t>::max();
+        } else {
+            batch_end_iteration = base_iteration + batch_size;
+        }
+    }
+
+    void OnBatchReplay(uint64_t base_iteration, uint64_t batch_size) override {
+        for (const BatchCheckpoint& entry : current_batch_checkpoints) {
+            rollback_checkpoint(entry.index, entry.checkpoint);
+        }
+        OnBatchStart(base_iteration, batch_size);
+    }
+
     void process_checkpoint(uint64_t i, const form& checkpoint, bool record_stats) {
+        apply_checkpoint(i, checkpoint, record_stats);
+    }
+
+  private:
+    struct BatchCheckpoint {
+        uint64_t index;
+        form checkpoint;
+    };
+
+    void rollback_checkpoint(uint64_t i, const form& checkpoint) {
+        form inverse_checkpoint = checkpoint.inverse();
+        apply_checkpoint(i, inverse_checkpoint, /*record_stats=*/false);
+    }
+
+    void apply_checkpoint(uint64_t i, const form& checkpoint, bool record_stats) {
         const bool do_stats = stats_enabled && record_stats;
         auto started_at = std::chrono::steady_clock::time_point{};
         if (do_stats) {
@@ -359,7 +399,6 @@ class StreamingOneWesolowskiCallback final : public WesolowskiCallback {
         return out;
     }
 
-  private:
     form& bucket(uint32_t j, uint64_t b) {
         size_t idx = static_cast<size_t>(j) * (1ULL << k) + static_cast<size_t>(b);
         return buckets[idx];
@@ -391,6 +430,9 @@ class StreamingOneWesolowskiCallback final : public WesolowskiCallback {
     integer getblock_inv_2k;
     integer getblock_r;
     integer getblock_tmp;
+    uint64_t batch_start_iteration = 1;
+    uint64_t batch_end_iteration = 0;
+    std::vector<BatchCheckpoint> current_batch_checkpoints;
 
     bool stats_enabled;
     uint64_t checkpoint_total_ns = 0;

src/callback.h

@@ -73,6 +73,14 @@ class WesolowskiCallback :public INUDUPLListener {
     }
 
     virtual void OnIteration(int type, void *data, uint64_t iteration) = 0;
+    virtual void OnBatchStart(uint64_t base_iteration, uint64_t batch_size) {
+        (void)base_iteration;
+        (void)batch_size;
+    }
+    virtual void OnBatchReplay(uint64_t base_iteration, uint64_t batch_size) {
+        (void)base_iteration;
+        (void)batch_size;
+    }
 
     std::unique_ptr<form[]> forms;
     size_t forms_capacity = 0;

src/vdf.h

@@ -95,9 +95,9 @@ bool quiet_mode = false;
 // run concurrently in the same process; they must not share a pairindex.
 inline int vdf_fast_pairindex() {
 #if (defined(ARCH_X86) || defined(ARCH_X64)) && !defined(CHIA_DISABLE_ASM)
-    constexpr int kSlots = int(sizeof(master_counter) / sizeof(master_counter[0]));
-    static std::atomic<int> next_slot{0};
-    thread_local int slot = next_slot.fetch_add(1, std::memory_order_relaxed) % kSlots;
+    constexpr unsigned int kSlots = unsigned(sizeof(master_counter) / sizeof(master_counter[0]));
+    static std::atomic<unsigned int> next_slot{0};
+    thread_local int slot = int(next_slot.fetch_add(1u, std::memory_order_relaxed) % kSlots);
     return slot;
 #else
     return 0;
@@ -201,6 +201,9 @@ void repeated_square(uint64_t iterations, form f, const integer& D, const intege
         #endif
 
         uint64 batch_size=c_checkpoint_interval;
+        if (weso != NULL) {
+            weso->OnBatchStart(num_iterations, batch_size);
+        }
 
         #ifdef ENABLE_TRACK_CYCLES
             print( "track cycles enabled; results will be wrong" );
@@ -231,6 +234,9 @@ void repeated_square(uint64_t iterations, form f, const integer& D, const intege
 
         if (actual_iterations==~uint64(0)) {
             //corruption; f is unchanged. do the entire batch with the slow algorithm
+            if (weso != NULL) {
+                weso->OnBatchReplay(num_iterations, batch_size);
+            }
             repeated_square_original(*weso->vdfo, f, D, L, num_iterations, batch_size, weso);
             actual_iterations=batch_size;