loro_kv_store/
block.rs

1use std::{
2    fmt::Debug,
3    io::Write,
4    ops::{Bound, Range},
5    sync::Arc,
6};
7
8use bytes::{Buf, Bytes};
9use loro_common::{LoroError, LoroResult};
10use once_cell::sync::OnceCell;
11
12use crate::{
13    compress::{compress, decompress, CompressionType},
14    iter::KvIterator,
15    sstable::{get_common_prefix_len_and_strip, SIZE_OF_U32, XXH_SEED},
16};
17
18use super::sstable::{SIZE_OF_U16, SIZE_OF_U8};
19
20const MAX_NORMAL_BLOCK_DATA_LEN: usize = u16::MAX as usize;
21const MAX_NORMAL_BLOCK_ENTRIES: usize = u16::MAX as usize;
22
23#[derive(Debug, Clone)]
24pub struct LargeValueBlock {
25    // without checksum
26    pub value_bytes: Bytes,
27    pub encoded_bytes: OnceCell<(Bytes, CompressionType)>,
28    pub key: Bytes,
29}
30
31impl LargeValueBlock {
32    /// ┌──────────────────────────┐
33    /// │Large Block               │
34    /// │┌ ─ ─ ─ ┬ ─ ─ ─ ─ ─ ─ ─ ─ │
35    /// │  value   Block Checksum ││
36    /// ││ bytes │      u32        │
37    /// │ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘│
38    /// └──────────────────────────┘
39    fn encode(&self, w: &mut Vec<u8>, mut compression_type: CompressionType) -> CompressionType {
40        if let Some((bytes, encoded_compression_type)) = self.encoded_bytes.get() {
41            if encoded_compression_type == &compression_type {
42                w.extend_from_slice(bytes);
43                return compression_type;
44            }
45        }
46
47        let origin_len = w.len();
48        compress(w, &self.value_bytes, compression_type);
49        if !compression_type.is_none() && w.len() - origin_len > self.value_bytes.len() {
50            w.truncate(origin_len);
51            compress(w, &self.value_bytes, CompressionType::None);
52            ensure_cov::notify_cov("kv_store::block::LargeValueBlock::encode::compress_fallback");
53            compression_type = CompressionType::None;
54        }
55        let checksum = xxhash_rust::xxh32::xxh32(&w[origin_len..], XXH_SEED);
56        w.write_all(&checksum.to_le_bytes()).unwrap();
57        compression_type
58    }
59
60    fn decode(bytes: Bytes, key: Bytes, compression_type: CompressionType) -> LoroResult<Self> {
61        let mut value_bytes = vec![];
62        decompress(
63            &mut value_bytes,
64            bytes.slice(..bytes.len() - SIZE_OF_U32),
65            compression_type,
66        )?;
67        Ok(LargeValueBlock {
68            value_bytes: Bytes::from(value_bytes),
69            encoded_bytes: OnceCell::with_value((bytes, compression_type)),
70            key,
71        })
72    }
73}
74
75#[derive(Debug, Clone)]
76pub struct NormalBlock {
77    pub data: Bytes,
78    pub encoded_data: OnceCell<(Bytes, CompressionType)>,
79    pub first_key: Bytes,
80    pub offsets: Vec<u16>,
81}
82
83impl NormalBlock {
84    /// ┌────────────────────────────────────────────────────────────────────────────────────────┐
85    /// │Block                                                                                   │
86    /// │┌ ─ ─ ─ ─ ─ ─ ─ ┬ ─ ─ ─┌ ─ ─ ─ ─ ─ ─ ─ ┬ ─ ─ ─ ─┌ ─ ─ ─┌ ─ ─ ─ ┬ ─ ─ ─ ─┌ ─ ─ ─ ─ ─ ─ ─ │
87    /// │ Key Value Chunk  ...  │Key Value Chunk  offset │ ...  │ offset  kv len │Block Checksum││
88    /// ││     bytes     │      │     bytes     │  u16   │      │  u16  │  u16   │     u32       │
89    /// │ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘─ ─ ─ ┘─ ─ ─ ─ ─ ─ ─ ─ ┘─ ─ ─ ─ ─ ─ ─ ┘│
90    /// └────────────────────────────────────────────────────────────────────────────────────────┘
91    ///
92    /// The block body may be compressed then we calculate its checksum (the checksum is not compressed).
93    fn encode(&self, w: &mut Vec<u8>, mut compression_type: CompressionType) -> CompressionType {
94        if let Some((encoded_data, encoded_compression_type)) = self.encoded_data.get() {
95            if encoded_compression_type == &compression_type {
96                w.extend_from_slice(encoded_data);
97                return compression_type;
98            }
99        }
100
101        let origin_len = w.len();
102        let mut buf = self.data.to_vec();
103        for offset in &self.offsets {
104            buf.extend_from_slice(&offset.to_le_bytes());
105        }
106        buf.extend_from_slice(&(self.offsets.len() as u16).to_le_bytes());
107        compress(w, &buf, compression_type);
108        if !compression_type.is_none() && w.len() - origin_len > buf.len() {
109            w.truncate(origin_len);
110            compress(w, &buf, CompressionType::None);
111            ensure_cov::notify_cov("kv_store::block::NormalBlock::encode::compress_fallback");
112            compression_type = CompressionType::None;
113        }
114        let checksum = xxhash_rust::xxh32::xxh32(&w[origin_len..], XXH_SEED);
115        w.extend_from_slice(&checksum.to_le_bytes());
116        compression_type
117    }
118
119    fn decode(
120        raw_block_and_check: Bytes,
121        first_key: Bytes,
122        compression_type: CompressionType,
123    ) -> LoroResult<NormalBlock> {
124        if raw_block_and_check.len() < SIZE_OF_U32 {
125            return Err(LoroError::DecodeError("Invalid bytes".into()));
126        }
127
128        let buf = raw_block_and_check.slice(..raw_block_and_check.len() - SIZE_OF_U32);
129        let mut data = vec![];
130        decompress(&mut data, buf, compression_type)?;
131        if data.len() < SIZE_OF_U16 {
132            return Err(LoroError::DecodeError("Invalid bytes".into()));
133        }
134
135        let offsets_len = (&data[data.len() - SIZE_OF_U16..]).get_u16_le() as usize;
136        if offsets_len == 0 {
137            return Err(LoroError::DecodeError("Invalid bytes".into()));
138        }
139
140        let offsets_bytes_len = SIZE_OF_U16
141            .checked_mul(offsets_len + 1)
142            .ok_or_else(|| LoroError::DecodeError("Invalid bytes".into()))?;
143        if data.len() < offsets_bytes_len {
144            return Err(LoroError::DecodeError("Invalid bytes".into()));
145        }
146
147        let data_end = data.len() - offsets_bytes_len;
148        if data_end > u16::MAX as usize {
149            return Err(LoroError::DecodeError("Invalid bytes".into()));
150        }
151
152        let offsets = &data[data_end..data.len() - SIZE_OF_U16];
153        let offsets: Vec<u16> = offsets
154            .chunks(SIZE_OF_U16)
155            .map(|mut chunk| chunk.get_u16_le())
156            .collect();
157        Self::validate_decoded_data(&data[..data_end], &offsets, &first_key)?;
158        Ok(NormalBlock {
159            data: Bytes::copy_from_slice(&data[..data_end]),
160            encoded_data: OnceCell::with_value((raw_block_and_check, compression_type)),
161            offsets,
162            first_key,
163        })
164    }
165
166    fn validate_decoded_data(data: &[u8], offsets: &[u16], first_key: &[u8]) -> LoroResult<()> {
167        if offsets.first().copied() != Some(0) {
168            return Err(LoroError::DecodeError("Invalid bytes".into()));
169        }
170
171        let mut prev_key: Option<Vec<u8>> = None;
172        let mut prev_offset = 0usize;
173        for (idx, offset) in offsets.iter().map(|x| *x as usize).enumerate() {
174            let offset_end = offsets
175                .get(idx + 1)
176                .map_or(data.len(), |next| *next as usize);
177            if offset < prev_offset || offset > offset_end || offset_end > data.len() {
178                return Err(LoroError::DecodeError("Invalid bytes".into()));
179            }
180
181            let key = if idx == 0 {
182                first_key.to_vec()
183            } else {
184                let header_end = offset
185                    .checked_add(SIZE_OF_U8 + SIZE_OF_U16)
186                    .ok_or_else(|| LoroError::DecodeError("Invalid bytes".into()))?;
187                if header_end > offset_end {
188                    return Err(LoroError::DecodeError("Invalid bytes".into()));
189                }
190
191                let common_prefix_len = data[offset] as usize;
192                if common_prefix_len > first_key.len() {
193                    return Err(LoroError::DecodeError("Invalid bytes".into()));
194                }
195
196                let key_suffix_len =
197                    u16::from_le_bytes(data[offset + SIZE_OF_U8..header_end].try_into().unwrap())
198                        as usize;
199                let key_end = header_end
200                    .checked_add(key_suffix_len)
201                    .ok_or_else(|| LoroError::DecodeError("Invalid bytes".into()))?;
202                if key_end > offset_end {
203                    return Err(LoroError::DecodeError("Invalid bytes".into()));
204                }
205
206                let mut key = Vec::with_capacity(common_prefix_len + key_suffix_len);
207                key.extend_from_slice(&first_key[..common_prefix_len]);
208                key.extend_from_slice(&data[header_end..key_end]);
209                key
210            };
211
212            if key.is_empty()
213                || prev_key
214                    .as_ref()
215                    .is_some_and(|prev_key| prev_key.as_slice() >= key.as_slice())
216            {
217                return Err(LoroError::DecodeError("Invalid bytes".into()));
218            }
219
220            prev_offset = offset;
221            prev_key = Some(key);
222        }
223
224        Ok(())
225    }
226}
227
228#[derive(Debug, Clone)]
229pub enum Block {
230    Normal(NormalBlock),
231    Large(LargeValueBlock),
232}
233
234impl Block {
235    pub fn is_large(&self) -> bool {
236        matches!(self, Block::Large(_))
237    }
238
239    pub fn data(&self) -> Bytes {
240        match self {
241            Block::Normal(block) => block.data.clone(),
242            Block::Large(block) => block.value_bytes.clone(),
243        }
244    }
245
246    pub fn first_key(&self) -> Bytes {
247        match self {
248            Block::Normal(block) => block.first_key.clone(),
249            Block::Large(block) => block.key.clone(),
250        }
251    }
252
253    pub fn last_key(&self) -> Bytes {
254        match self {
255            Block::Normal(block) => {
256                if block.offsets.len() == 1 {
257                    return block.first_key.clone();
258                }
259
260                let offset = *block.offsets.last().unwrap() as usize;
261                let mut bytes = &block.data[offset..];
262                let common_prefix_len = bytes.get_u8() as usize;
263                let key_suffix_len = bytes.get_u16_le() as usize;
264                let mut last_key = Vec::with_capacity(common_prefix_len + key_suffix_len);
265                last_key.extend_from_slice(&block.first_key[..common_prefix_len]);
266                last_key.extend_from_slice(&bytes[..key_suffix_len]);
267                last_key.into()
268            }
269            Block::Large(block) => block.key.clone(),
270        }
271    }
272
273    pub fn encode(&self, w: &mut Vec<u8>, compression_type: CompressionType) -> CompressionType {
274        match self {
275            Block::Normal(block) => block.encode(w, compression_type),
276            Block::Large(block) => block.encode(w, compression_type),
277        }
278    }
279
280    pub(crate) fn try_decode(
281        raw_block_and_check: Bytes,
282        is_large: bool,
283        key: Bytes,
284        compression_type: CompressionType,
285    ) -> LoroResult<Self> {
286        if key.is_empty() {
287            return Err(LoroError::DecodeError("Invalid bytes".into()));
288        }
289
290        if is_large {
291            return LargeValueBlock::decode(raw_block_and_check, key, compression_type)
292                .map(Block::Large);
293        }
294        NormalBlock::decode(raw_block_and_check, key, compression_type).map(Block::Normal)
295    }
296
297    pub fn decode(
298        raw_block_and_check: Bytes,
299        is_large: bool,
300        key: Bytes,
301        compression_type: CompressionType,
302    ) -> Self {
303        // The caller is responsible for validating SSTable integrity before lazy block reads.
304        Self::try_decode(raw_block_and_check, is_large, key, compression_type)
305            .expect("validated SSTable block should decode")
306    }
307
308    pub fn len(&self) -> usize {
309        match self {
310            Block::Normal(block) => block.offsets.len(),
311            Block::Large(_) => 1,
312        }
313    }
314
315    pub fn is_empty(&self) -> bool {
316        match self {
317            Block::Normal(block) => block.offsets.is_empty(),
318            Block::Large(_) => false,
319        }
320    }
321}
322
323#[derive(Debug)]
324pub struct BlockBuilder {
325    data: Vec<u8>,
326    offsets: Vec<u16>,
327    block_size: usize,
328    // for key compression
329    first_key: Bytes,
330    is_large: bool,
331}
332
333impl BlockBuilder {
334    pub fn new(block_size: usize) -> Self {
335        Self {
336            data: Vec::new(),
337            offsets: Vec::new(),
338            block_size,
339            first_key: Bytes::new(),
340            is_large: false,
341        }
342    }
343
344    pub fn estimated_size(&self) -> usize {
345        if self.is_large {
346            self.data.len()
347        } else {
348            // key-value pairs number
349            SIZE_OF_U16 +
350            // offsets
351            self.offsets.len() * SIZE_OF_U16 +
352            // key-value pairs data
353            self.data.len() +
354            // checksum
355            SIZE_OF_U32
356        }
357    }
358
359    pub fn is_empty(&self) -> bool {
360        !self.is_large && self.offsets.is_empty()
361    }
362
363    /// Add a key-value pair to the block.
364    /// Returns true if the key-value pair is added successfully, false the block is full.
365    ///
366    /// ┌─────────────────────────────────────────────────────┐
367    /// │  Key Value Chunk                                    │
368    /// │┌ ─ ─ ─ ─ ─ ─ ─ ─ ┬ ─ ─ ─ ─ ─ ─ ─┌ ─ ─ ─ ─ ─┬ ─ ─ ─ ┐│
369    /// │ common prefix len key suffix len│key suffix│ value ││
370    /// ││       u8        │     u16      │  bytes   │ bytes ││
371    /// │ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘─ ─ ─ ─ ─ ┘ ─ ─ ─ ┘│
372    /// └─────────────────────────────────────────────────────┘
373    ///
374    pub fn add(&mut self, key: &[u8], value: &[u8]) -> bool {
375        if key.is_empty() {
376            return false;
377        }
378
379        debug_assert!(!key.is_empty(), "key cannot be empty");
380        if self.first_key.is_empty() {
381            if value.len() > self.block_size || value.len() > MAX_NORMAL_BLOCK_DATA_LEN {
382                self.data.extend_from_slice(value);
383                self.is_large = true;
384                self.first_key = Bytes::copy_from_slice(key);
385                return true;
386            }
387
388            self.first_key = Bytes::copy_from_slice(key);
389            self.offsets.push(self.data.len() as u16);
390            self.data.extend_from_slice(value);
391            return true;
392        }
393
394        if self.offsets.len() >= MAX_NORMAL_BLOCK_ENTRIES {
395            return false;
396        }
397
398        let (common, suffix) = get_common_prefix_len_and_strip(key, &self.first_key);
399        let key_len = suffix.len();
400        let Some(next_data_len) = self
401            .data
402            .len()
403            .checked_add(SIZE_OF_U8 + SIZE_OF_U16)
404            .and_then(|len| len.checked_add(key_len))
405            .and_then(|len| len.checked_add(value.len()))
406        else {
407            return false;
408        };
409        if next_data_len > MAX_NORMAL_BLOCK_DATA_LEN {
410            return false;
411        }
412
413        // whether the block is full
414        let Some(estimated_size) = self
415            .estimated_size()
416            .checked_add(key_len)
417            .and_then(|len| len.checked_add(value.len()))
418            .and_then(|len| len.checked_add(SIZE_OF_U8 + SIZE_OF_U16))
419        else {
420            return false;
421        };
422        if estimated_size > self.block_size {
423            return false;
424        }
425
426        self.offsets.push(self.data.len() as u16);
427        self.data.push(common);
428        self.data.extend_from_slice(&(key_len as u16).to_le_bytes());
429        self.data.extend_from_slice(suffix);
430        self.data.extend_from_slice(value);
431        true
432    }
433
434    pub fn build(self) -> Block {
435        if self.is_large {
436            return Block::Large(LargeValueBlock {
437                value_bytes: Bytes::from(self.data),
438                key: self.first_key,
439                encoded_bytes: OnceCell::new(),
440            });
441        }
442        debug_assert!(!self.offsets.is_empty(), "block is empty");
443        Block::Normal(NormalBlock {
444            data: Bytes::from(self.data),
445            offsets: self.offsets,
446            first_key: self.first_key,
447            encoded_data: OnceCell::new(),
448        })
449    }
450}
451
452/// Block iterator
453///
454/// If the key is empty, it means the iterator is invalid.
455#[derive(Clone)]
456pub struct BlockIter {
457    block: Arc<Block>,
458    next_key: Bytes,
459    next_value_range: Range<usize>,
460    prev_key: Bytes,
461    prev_value_range: Range<usize>,
462    next_idx: usize,
463    prev_idx: isize,
464    first_key: Bytes,
465}
466
467impl Debug for BlockIter {
468    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
469        f.debug_struct("BlockIter")
470            .field("is_large", &self.block.is_large())
471            .field("next_key", &self.next_key)
472            .field("next_value_range", &self.next_value_range)
473            .field("prev_key", &self.prev_key)
474            .field("prev_value_range", &self.prev_value_range)
475            .field("next_idx", &self.next_idx)
476            .field("prev_idx", &self.prev_idx)
477            .field("first_key", &Bytes::copy_from_slice(&self.first_key))
478            .finish()
479    }
480}
481
482impl BlockIter {
483    pub fn new(block: Arc<Block>) -> Self {
484        let prev_idx = block.len() as isize - 1;
485        let mut iter = Self {
486            first_key: block.first_key(),
487            block,
488            next_key: Bytes::new(),
489            next_value_range: 0..0,
490            prev_key: Bytes::new(),
491            prev_value_range: 0..0,
492            next_idx: 0,
493            prev_idx,
494        };
495        iter.seek_to_idx(0);
496        iter.back_to_idx(prev_idx);
497        iter
498    }
499
500    pub fn new_seek_to_key(block: Arc<Block>, key: &[u8]) -> Self {
501        let prev_idx = block.len() as isize - 1;
502        let mut iter = Self {
503            first_key: block.first_key(),
504            block,
505            next_key: Bytes::new(),
506            next_value_range: 0..0,
507            prev_key: Bytes::new(),
508            prev_value_range: 0..0,
509            next_idx: 0,
510            prev_idx,
511        };
512        iter.seek_to_key(key);
513        iter.back_to_idx(prev_idx);
514        iter
515    }
516
517    pub fn new_back_to_key(block: Arc<Block>, key: &[u8]) -> Self {
518        let prev_idx = block.len() as isize - 1;
519        let mut iter = Self {
520            first_key: block.first_key(),
521            block,
522            next_key: Bytes::new(),
523            next_value_range: 0..0,
524            prev_key: Bytes::new(),
525            prev_value_range: 0..0,
526            next_idx: 0,
527            prev_idx,
528        };
529        iter.seek_to_idx(0);
530        iter.back_to_key(key);
531        iter
532    }
533
534    pub fn new_scan(block: Arc<Block>, start: Bound<&[u8]>, end: Bound<&[u8]>) -> Self {
535        let mut iter = match start {
536            Bound::Included(key) => Self::new_seek_to_key(block, key),
537            Bound::Excluded(key) => {
538                let mut iter = Self::new_seek_to_key(block, key);
539                while iter.has_next() && iter.peek_next_curr_key().unwrap() == key {
540                    iter.next();
541                }
542                iter
543            }
544            Bound::Unbounded => Self::new(block),
545        };
546        match end {
547            Bound::Included(key) => {
548                iter.back_to_key(key);
549            }
550            Bound::Excluded(key) => {
551                iter.back_to_key(key);
552                while iter.has_next_back() && iter.peek_back_curr_key().unwrap() == key {
553                    iter.next_back();
554                }
555            }
556            Bound::Unbounded => {}
557        }
558        iter
559    }
560
561    pub fn peek_next_curr_key(&self) -> Option<Bytes> {
562        if self.has_next() {
563            Some(Bytes::copy_from_slice(&self.next_key))
564        } else {
565            None
566        }
567    }
568
569    pub fn peek_next_curr_value(&self) -> Option<Bytes> {
570        if self.has_next() {
571            Some(self.block.data().slice(self.next_value_range.clone()))
572        } else {
573            None
574        }
575    }
576
577    pub fn has_next(&self) -> bool {
578        !self.next_key.is_empty() && self.next_idx as isize <= self.prev_idx
579    }
580
581    pub fn peek_back_curr_key(&self) -> Option<Bytes> {
582        if self.has_next_back() {
583            Some(Bytes::copy_from_slice(&self.prev_key))
584        } else {
585            None
586        }
587    }
588
589    pub fn peek_back_curr_value(&self) -> Option<Bytes> {
590        if self.has_next_back() {
591            Some(self.block.data().slice(self.prev_value_range.clone()))
592        } else {
593            None
594        }
595    }
596
597    pub fn has_next_back(&self) -> bool {
598        !self.prev_key.is_empty() && self.next_idx as isize <= self.prev_idx
599    }
600
601    pub fn next(&mut self) {
602        self.next_idx += 1;
603        if self.next_idx as isize > self.prev_idx {
604            self.next_key.clear();
605            self.next_value_range = 0..0;
606            return;
607        }
608        self.seek_to_idx(self.next_idx);
609    }
610
611    pub fn next_back(&mut self) {
612        self.prev_idx -= 1;
613        if self.prev_idx < 0 || self.prev_idx < (self.next_idx as isize) {
614            self.prev_key.clear();
615            self.prev_value_range = 0..0;
616            return;
617        }
618        self.back_to_idx(self.prev_idx);
619    }
620
621    pub fn seek_to_key(&mut self, key: &[u8]) {
622        match self.block.as_ref() {
623            Block::Normal(block) => {
624                let mut left = 0;
625                let mut right = block.offsets.len();
626                while left < right {
627                    let mid = left + (right - left) / 2;
628                    self.seek_to_idx(mid);
629                    debug_assert!(self.has_next());
630                    if self.next_key == key {
631                        return;
632                    }
633                    if self.next_key < key {
634                        left = mid + 1;
635                    } else {
636                        right = mid;
637                    }
638                }
639                self.seek_to_idx(left);
640            }
641            Block::Large(block) => {
642                if key > block.key {
643                    self.seek_to_idx(1);
644                } else {
645                    self.seek_to_idx(0);
646                }
647            }
648        }
649    }
650
651    /// MUST be called after seek_to_key()
652    pub fn back_to_key(&mut self, key: &[u8]) {
653        match self.block.as_ref() {
654            Block::Normal(block) => {
655                let mut left = self.next_idx;
656                let mut right = block.offsets.len();
657                while left < right {
658                    let mid = left + (right - left) / 2;
659                    self.back_to_idx(mid as isize);
660                    // prev idx <= next idx
661                    if !self.has_next_back() {
662                        return;
663                    }
664                    debug_assert!(self.has_next_back());
665                    if self.prev_key > key {
666                        right = mid;
667                    } else {
668                        left = mid + 1;
669                    }
670                }
671                self.back_to_idx(left as isize - 1);
672            }
673            Block::Large(block) => {
674                if key < block.key {
675                    self.back_to_idx(-1);
676                } else {
677                    self.back_to_idx(0);
678                }
679            }
680        }
681    }
682
683    fn seek_to_idx(&mut self, idx: usize) {
684        match self.block.as_ref() {
685            Block::Normal(block) => {
686                if idx >= block.offsets.len() {
687                    self.next_key.clear();
688                    self.next_value_range = 0..0;
689                    self.next_idx = idx;
690                    return;
691                }
692                let offset = block.offsets[idx] as usize;
693                self.seek_to_offset(
694                    offset,
695                    *block
696                        .offsets
697                        .get(idx + 1)
698                        .unwrap_or(&(block.data.len() as u16)) as usize,
699                    idx == 0,
700                );
701                self.next_idx = idx;
702            }
703            Block::Large(block) => {
704                if idx > 0 {
705                    self.next_key.clear();
706                    self.next_value_range = 0..0;
707                    self.next_idx = idx;
708                    return;
709                }
710                self.next_key = block.key.clone();
711                self.next_value_range = 0..block.value_bytes.len();
712                self.next_idx = idx;
713            }
714        }
715    }
716
717    fn back_to_idx(&mut self, idx: isize) {
718        match self.block.as_ref() {
719            Block::Normal(block) => {
720                if idx < 0 {
721                    self.prev_key.clear();
722                    self.prev_value_range = 0..0;
723                    self.prev_idx = idx;
724                    return;
725                }
726                let offset = block.offsets[idx as usize] as usize;
727                self.back_to_offset(
728                    offset,
729                    *block
730                        .offsets
731                        .get(idx as usize + 1)
732                        .unwrap_or(&(block.data.len() as u16)) as usize,
733                    idx == 0,
734                );
735                self.prev_idx = idx;
736            }
737            Block::Large(block) => {
738                if idx < 0 {
739                    self.prev_key.clear();
740                    self.prev_value_range = 0..0;
741                    self.prev_idx = idx;
742                    return;
743                }
744                self.prev_key = block.key.clone();
745                self.prev_value_range = 0..block.value_bytes.len();
746                self.prev_idx = idx;
747            }
748        }
749    }
750
751    fn seek_to_offset(&mut self, offset: usize, offset_end: usize, is_first: bool) {
752        match self.block.as_ref() {
753            Block::Normal(block) => {
754                if is_first {
755                    self.next_key = self.first_key.clone();
756                    self.next_value_range = offset..offset_end;
757                    return;
758                }
759                let mut rest = &block.data[offset..];
760                let common_prefix_len = rest.get_u8() as usize;
761                let key_suffix_len = rest.get_u16_le() as usize;
762                let mut next_key = Vec::with_capacity(common_prefix_len + key_suffix_len);
763                next_key.extend_from_slice(&self.first_key[..common_prefix_len]);
764                next_key.extend_from_slice(&rest[..key_suffix_len]);
765                self.next_key = next_key.into();
766                let value_start = offset + SIZE_OF_U8 + SIZE_OF_U16 + key_suffix_len;
767                self.next_value_range = value_start..offset_end;
768            }
769            Block::Large(_) => {
770                unreachable!()
771            }
772        }
773    }
774
775    fn back_to_offset(&mut self, offset: usize, offset_end: usize, is_first: bool) {
776        match self.block.as_ref() {
777            Block::Normal(block) => {
778                if is_first {
779                    self.prev_key = self.first_key.clone();
780                    self.prev_value_range = offset..offset_end;
781                    return;
782                }
783                let mut rest = &block.data[offset..];
784                let common_prefix_len = rest.get_u8() as usize;
785                let key_suffix_len = rest.get_u16_le() as usize;
786                let mut prev_key = Vec::with_capacity(common_prefix_len + key_suffix_len);
787                prev_key.extend_from_slice(&self.first_key[..common_prefix_len]);
788                prev_key.extend_from_slice(&rest[..key_suffix_len]);
789                self.prev_key = prev_key.into();
790                let value_start = offset + SIZE_OF_U8 + SIZE_OF_U16 + key_suffix_len;
791                self.prev_value_range = value_start..offset_end;
792            }
793            Block::Large(_) => {
794                unreachable!()
795            }
796        }
797    }
798
799    pub fn peek_block(&self) -> &Arc<Block> {
800        &self.block
801    }
802
803    pub fn finish(&mut self) {
804        self.next_key.clear();
805        self.next_value_range = 0..0;
806        self.prev_key.clear();
807        self.prev_value_range = 0..0;
808    }
809}
810
811impl KvIterator for BlockIter {
812    fn peek_next_key(&self) -> Option<Bytes> {
813        self.peek_next_curr_key()
814    }
815
816    fn peek_next_value(&self) -> Option<Bytes> {
817        self.peek_next_curr_value()
818    }
819
820    fn next_(&mut self) {
821        self.next();
822    }
823
824    fn has_next(&self) -> bool {
825        self.has_next()
826    }
827
828    fn peek_next_back_key(&self) -> Option<Bytes> {
829        self.peek_back_curr_key()
830    }
831
832    fn peek_next_back_value(&self) -> Option<Bytes> {
833        self.peek_back_curr_value()
834    }
835
836    fn next_back_(&mut self) {
837        self.next_back();
838    }
839
840    fn has_next_back(&self) -> bool {
841        self.has_next_back()
842    }
843}
844
845impl Iterator for BlockIter {
846    type Item = (Bytes, Bytes);
847
848    fn next(&mut self) -> Option<Self::Item> {
849        if !self.has_next() {
850            return None;
851        }
852        let key = self.peek_next_curr_key().unwrap();
853        let value = self.peek_next_curr_value().unwrap();
854        self.next();
855        Some((key, value))
856    }
857}
858
859impl DoubleEndedIterator for BlockIter {
860    fn next_back(&mut self) -> Option<Self::Item> {
861        if !self.has_next_back() {
862            return None;
863        }
864        let key = self.peek_back_curr_key().unwrap();
865        let value = self.peek_back_curr_value().unwrap();
866        self.next_back();
867        Some((key, value))
868    }
869}
loro_kv_store/block.rs

loro_kv_store/
block.rs
