1use std::collections::VecDeque;
33
34const CHUNK_TARGET: usize = 128;
37const SPLIT_AT: usize = CHUNK_TARGET * 2;
40
41#[derive(Clone, Debug)]
42struct Chunk {
43 heights: Vec<u32>,
45 measured: Vec<bool>,
48 sum: u64,
50}
51
52impl Chunk {
53 fn new() -> Chunk {
54 Chunk {
55 heights: Vec::with_capacity(CHUNK_TARGET),
56 measured: Vec::with_capacity(CHUNK_TARGET),
57 sum: 0,
58 }
59 }
60
61 fn len(&self) -> usize {
62 self.heights.len()
63 }
64
65 fn recompute(&mut self) {
66 self.sum = self.heights.iter().map(|h| u64::from(*h)).sum();
67 }
68}
69
70#[derive(Clone, Copy, PartialEq, Eq, Debug)]
72pub struct Hit {
73 pub row: usize,
75 pub offset: u32,
77}
78
79#[derive(Debug, Default)]
81pub struct HeightIndex {
82 chunks: VecDeque<Chunk>,
83 prefix: Vec<u64>,
86 prefix_valid: usize,
87 len: usize,
88}
89
90impl HeightIndex {
91 pub fn new() -> HeightIndex {
92 HeightIndex::default()
93 }
94
95 pub fn len(&self) -> usize {
97 self.len
98 }
99
100 pub fn is_empty(&self) -> bool {
101 self.len == 0
102 }
103
104 pub fn total_height(&mut self) -> u64 {
106 self.repair_prefix();
107 self.prefix.last().copied().unwrap_or(0)
108 }
109
110 pub fn height_at(&self, row: usize) -> u32 {
112 match self.locate_chunk(row) {
113 Some((ci, within)) => self.chunks[ci].heights[within],
114 None => 0,
115 }
116 }
117
118 pub fn is_measured(&self, row: usize) -> bool {
120 match self.locate_chunk(row) {
121 Some((ci, within)) => self.chunks[ci].measured[within],
122 None => false,
123 }
124 }
125
126 pub fn unmeasured_count(&self) -> usize {
131 self.chunks
132 .iter()
133 .map(|c| c.measured.iter().filter(|m| !**m).count())
134 .sum()
135 }
136
137 pub fn push_back(&mut self, height: u32, measured: bool) {
139 let need_new = match self.chunks.back() {
140 None => true,
141 Some(c) => c.len() >= CHUNK_TARGET,
142 };
143 if need_new {
144 self.chunks.push_back(Chunk::new());
145 }
146 let ci = self.chunks.len() - 1;
147 let c = &mut self.chunks[ci];
148 c.heights.push(height);
149 c.measured.push(measured);
150 c.sum += u64::from(height);
151 self.len += 1;
152 self.invalidate_from(ci);
153 }
154
155 pub fn push_front(&mut self, height: u32, measured: bool) {
157 let need_new = match self.chunks.front() {
158 None => true,
159 Some(c) => c.len() >= CHUNK_TARGET,
160 };
161 if need_new {
162 self.chunks.push_front(Chunk::new());
163 }
164 let c = &mut self.chunks[0];
165 c.heights.insert(0, height);
166 c.measured.insert(0, true & measured);
167 c.sum += u64::from(height);
168 self.len += 1;
169 self.invalidate_from(0);
170 }
171
172 pub fn insert(&mut self, row: usize, height: u32, measured: bool) {
180 if row >= self.len {
181 self.push_back(height, measured);
182 return;
183 }
184 let (ci, within) = match self.locate_chunk(row) {
185 Some(x) => x,
186 None => {
187 self.push_back(height, measured);
188 return;
189 }
190 };
191 {
192 let c = &mut self.chunks[ci];
193 c.heights.insert(within, height);
194 c.measured.insert(within, measured);
195 c.sum += u64::from(height);
196 }
197 self.len += 1;
198 if self.chunks[ci].len() >= SPLIT_AT {
199 self.split_chunk(ci);
200 }
201 self.invalidate_from(ci);
202 }
203
204 pub fn remove(&mut self, row: usize) -> Option<u32> {
206 let (ci, within) = self.locate_chunk(row)?;
207 let h = {
208 let c = &mut self.chunks[ci];
209 let h = c.heights.remove(within);
210 c.measured.remove(within);
211 c.sum -= u64::from(h);
212 h
213 };
214 self.len -= 1;
215 if self.chunks[ci].len() == 0 && self.chunks.len() > 1 {
216 self.chunks.remove(ci);
217 }
218 self.invalidate_from(ci.min(self.chunks.len().saturating_sub(1)));
219 Some(h)
220 }
221
222 pub fn drain_front(&mut self, n: usize) {
224 let mut left = n.min(self.len);
225 while left > 0 {
226 let front_len = match self.chunks.front() {
227 Some(c) => c.len(),
228 None => break,
229 };
230 if front_len <= left {
231 let c = self.chunks.pop_front().expect("front exists");
232 left -= front_len;
233 self.len -= front_len;
234 let _ = c;
235 } else {
236 let c = self.chunks.front_mut().expect("front exists");
237 c.heights.drain(..left);
238 c.measured.drain(..left);
239 c.recompute();
240 self.len -= left;
241 left = 0;
242 }
243 }
244 self.invalidate_from(0);
245 }
246
247 pub fn set_height(&mut self, row: usize, height: u32, measured: bool) {
250 let Some((ci, within)) = self.locate_chunk(row) else {
251 return;
252 };
253 let c = &mut self.chunks[ci];
254 let old = c.heights[within];
255 if old == height && c.measured[within] == measured {
256 return;
257 }
258 c.sum = c.sum - u64::from(old) + u64::from(height);
259 c.heights[within] = height;
260 c.measured[within] = measured;
261 self.invalidate_from(ci);
262 }
263
264 pub fn invalidate_all_measurements(&mut self) {
273 for c in &mut self.chunks {
274 c.measured.fill(false);
275 }
276 }
277
278 pub fn offset_of(&mut self, row: usize) -> u64 {
280 if row == 0 || self.len == 0 {
281 return 0;
282 }
283 self.repair_prefix();
284 let mut remaining = row.min(self.len);
285 let mut acc = 0u64;
286 for (i, c) in self.chunks.iter().enumerate() {
287 if remaining >= c.len() {
288 acc = self.prefix[i + 1];
289 remaining -= c.len();
290 if remaining == 0 {
291 return acc;
292 }
293 } else {
294 for h in &c.heights[..remaining] {
295 acc += u64::from(*h);
296 }
297 return acc;
298 }
299 }
300 acc
301 }
302
303 pub fn locate(&mut self, y: u64) -> Option<Hit> {
308 if self.len == 0 {
309 return None;
310 }
311 self.repair_prefix();
312 let total = self.prefix.last().copied().unwrap_or(0);
313 if y >= total {
314 let row = self.len - 1;
315 return Some(Hit {
316 row,
317 offset: self.height_at(row),
318 });
319 }
320 let mut lo = 0usize;
322 let mut hi = self.chunks.len();
323 while lo + 1 < hi {
324 let mid = lo + (hi - lo) / 2;
325 if self.prefix[mid] <= y {
326 lo = mid;
327 } else {
328 hi = mid;
329 }
330 }
331 let ci = lo;
332 let mut acc = self.prefix[ci];
333 let start_row = self.chunk_start_row(ci);
334 for (i, h) in self.chunks[ci].heights.iter().enumerate() {
335 let h64 = u64::from(*h);
336 if acc + h64 > y {
337 return Some(Hit {
338 row: start_row + i,
339 offset: (y - acc) as u32,
340 });
341 }
342 acc += h64;
343 }
344 Some(Hit {
346 row: self.len - 1,
347 offset: 0,
348 })
349 }
350
351 fn chunk_start_row(&self, ci: usize) -> usize {
354 self.chunks.iter().take(ci).map(|c| c.len()).sum()
355 }
356
357 fn locate_chunk(&self, row: usize) -> Option<(usize, usize)> {
358 if row >= self.len {
359 return None;
360 }
361 let mut acc = 0usize;
362 for (i, c) in self.chunks.iter().enumerate() {
363 if row < acc + c.len() {
364 return Some((i, row - acc));
365 }
366 acc += c.len();
367 }
368 None
369 }
370
371 fn split_chunk(&mut self, ci: usize) {
372 let at = self.chunks[ci].len() / 2;
373 let tail_h = self.chunks[ci].heights.split_off(at);
374 let tail_m = self.chunks[ci].measured.split_off(at);
375 self.chunks[ci].recompute();
376 let mut tail = Chunk {
377 heights: tail_h,
378 measured: tail_m,
379 sum: 0,
380 };
381 tail.recompute();
382 self.chunks.insert(ci + 1, tail);
383 }
384
385 fn invalidate_from(&mut self, ci: usize) {
386 self.prefix_valid = self.prefix_valid.min(ci);
387 }
388
389 fn repair_prefix(&mut self) {
394 let n = self.chunks.len();
395 if self.prefix.len() != n + 1 {
396 self.prefix.resize(n + 1, 0);
397 self.prefix_valid = self.prefix_valid.min(n);
398 }
399 if self.prefix_valid >= n {
400 return;
401 }
402 let start = self.prefix_valid;
403 let mut acc = self.prefix[start];
404 for i in start..n {
405 acc += self.chunks[i].sum;
406 self.prefix[i + 1] = acc;
407 }
408 self.prefix_valid = n;
409 }
410}