From 293455fef9ec8bc5bfcfbd38a12825fa236be105 Mon Sep 17 00:00:00 2001
From: Nils Homer <nilshomer@gmail.com>
Date: Sat, 13 Jun 2026 16:41:32 -0700
Subject: [PATCH] fix(upstream): stop FixDupsIter replaying the final training
 row

FixDupsIter's terminal branch returned `self.last_item.take()`, which
re-emitted the already-yielded final row. `RMITrainingData::iter()` therefore
produced len()+1 items, so every active-path `LinearModel` leaf/routing fit
double-weighted its last point in `slr`. Return `None` at exhaustion instead.

Add regression tests pinning one-item-per-row iteration and duplicate-key
collapsing, document FixDupsIter's contract, and drop the now-stale
replay-workaround note in the lower-bound test. Inherited from
learnedsystems/RMI; also fixed upstream in learnedsystems/RMI#15.

Refs #3
---
 prmi/src/upstream/models/mod.rs               | 46 ++++++++++++++++++-
 .../upstream/train/lower_bound_correction.rs  |  3 +-
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/prmi/src/upstream/models/mod.rs b/prmi/src/upstream/models/mod.rs
index 538077c..95f6b19 100644
--- a/prmi/src/upstream/models/mod.rs
+++ b/prmi/src/upstream/models/mod.rs
@@ -186,6 +186,10 @@ impl<K: TrainingKey> RMITrainingDataIteratorProvider for Vec<(K, usize)> {
     }
 }
 
+/// Iterator adapter over `(key, offset)` pairs that collapses each run of
+/// duplicate keys onto the offset of the run's *first* occurrence (lower-bound
+/// semantics for the CDF). It must yield exactly as many items as the
+/// underlying iterator — one per input pair.
 struct FixDupsIter<K, T: Iterator<Item = (K, usize)>> {
     iter: T,
     last_item: Option<(K, usize)>,
@@ -225,7 +229,12 @@ where
                         return Some(nxt);
                     }
                 }
-                None => self.last_item.take(),
+                // Underlying iterator exhausted: stop. Previously this returned
+                // `self.last_item.take()`, which re-emitted the already-yielded
+                // final row, so `iter()` produced len()+1 items and every
+                // `slr` leaf/routing fit (and the CDF trainers) double-weighted
+                // the last point.
+                None => None,
             },
         }
     }
@@ -848,3 +857,38 @@ mod training_key_epsilon_tests {
         assert_eq!(5u32.plus_epsilon(), 6);
     }
 }
+
+#[cfg(test)]
+mod fix_dups_iter_tests {
+    use super::*;
+
+    /// `iter()` must yield exactly one item per input pair. A prior bug in
+    /// `FixDupsIter`'s terminal branch re-emitted the final row, producing
+    /// len()+1 items and double-weighting the last point in every `slr` fit.
+    #[test]
+    fn iter_yields_one_item_per_input_row() {
+        let pairs: Vec<(u64, usize)> = vec![(0, 0), (1, 1), (3, 2), (100, 3)];
+        let data = RMITrainingData::<u64>::new(Box::new(pairs.clone()));
+
+        let iterated: Vec<(u64, usize)> = data.iter().collect();
+        assert_eq!(
+            iterated, pairs,
+            "iter() must reproduce the input one-to-one, with no trailing replay"
+        );
+    }
+
+    /// A run of duplicate keys collapses onto the first occurrence's offset, and
+    /// still yields exactly `len()` items (one per input row, no replay).
+    #[test]
+    fn iter_collapses_duplicate_keys_to_first_offset() {
+        let pairs: Vec<(u64, usize)> = vec![(5, 0), (5, 1), (5, 2), (7, 3)];
+        let data = RMITrainingData::<u64>::new(Box::new(pairs));
+
+        let iterated: Vec<(u64, usize)> = data.iter().collect();
+        assert_eq!(
+            iterated,
+            vec![(5, 0), (5, 0), (5, 0), (7, 3)],
+            "duplicate keys map to the first occurrence's offset; length is unchanged"
+        );
+    }
+}
diff --git a/prmi/src/upstream/train/lower_bound_correction.rs b/prmi/src/upstream/train/lower_bound_correction.rs
index 01687a9..5a7a4df 100644
--- a/prmi/src/upstream/train/lower_bound_correction.rs
+++ b/prmi/src/upstream/train/lower_bound_correction.rs
@@ -390,8 +390,7 @@ mod tests {
     ///
     /// All iterated entries fall in one run, so `longest_run(0)` must equal the
     /// number of entries the trainer iterates. (We derive the expected count
-    /// from `data.iter()` rather than hard-coding it, since `FixDupsIter` adds a
-    /// trailing replay of the last element.)
+    /// from `data.iter()` rather than hard-coding it.)
     #[test]
     fn single_run_reports_full_length() {
         let pairs: Vec<(u64, usize)> = vec![(5u64, 0), (5u64, 1), (5u64, 2)];