ipfrs_transport/
range_request.rs

1//! Partial block requests (range queries)
2//!
3//! Support for requesting byte ranges from blocks, useful for:
4//! - Partial tensor loading
5//! - Sparse block access
6//! - Efficient streaming of large blocks
7
8use ipfrs_core::Cid;
9use serde::{Deserialize, Deserializer, Serialize, Serializer};
10use std::ops::Range;
11use thiserror::Error;
12
13/// Serialize CID as string
14fn serialize_cid<S>(cid: &Cid, serializer: S) -> Result<S::Ok, S::Error>
15where
16    S: Serializer,
17{
18    serializer.serialize_str(&cid.to_string())
19}
20
21/// Deserialize CID from string
22fn deserialize_cid<'de, D>(deserializer: D) -> Result<Cid, D::Error>
23where
24    D: Deserializer<'de>,
25{
26    let s = String::deserialize(deserializer)?;
27    s.parse().map_err(serde::de::Error::custom)
28}
29
30/// Error types for range requests
31#[derive(Error, Debug)]
32pub enum RangeError {
33    #[error("Invalid range: {0}")]
34    InvalidRange(String),
35    #[error("Range out of bounds: requested {requested}, available {available}")]
36    OutOfBounds { requested: u64, available: u64 },
37    #[error("Block not found: {0}")]
38    BlockNotFound(Cid),
39    #[error("Unsatisfiable range")]
40    Unsatisfiable,
41}
42
43/// Byte range specification
44#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
45pub enum ByteRange {
46    /// Request bytes from offset to end
47    FromTo { start: u64, end: u64 },
48    /// Request bytes from offset to end of block
49    From(u64),
50    /// Request last N bytes
51    Suffix(u64),
52    /// Request entire block
53    All,
54}
55
56impl ByteRange {
57    /// Create a range from start to end (inclusive)
58    pub fn from_to(start: u64, end: u64) -> Result<Self, RangeError> {
59        if start > end {
60            return Err(RangeError::InvalidRange(format!(
61                "start ({}) > end ({})",
62                start, end
63            )));
64        }
65        Ok(ByteRange::FromTo { start, end })
66    }
67
68    /// Create a range from offset to end
69    pub fn from(start: u64) -> Self {
70        ByteRange::From(start)
71    }
72
73    /// Create a range for the last N bytes
74    pub fn suffix(count: u64) -> Self {
75        ByteRange::Suffix(count)
76    }
77
78    /// Convert to concrete byte range given total size
79    pub fn to_range(&self, total_size: u64) -> Result<Range<u64>, RangeError> {
80        match self {
81            ByteRange::FromTo { start, end } => {
82                if *end >= total_size {
83                    return Err(RangeError::OutOfBounds {
84                        requested: *end,
85                        available: total_size,
86                    });
87                }
88                Ok(*start..*end + 1)
89            }
90            ByteRange::From(start) => {
91                if *start >= total_size {
92                    return Err(RangeError::OutOfBounds {
93                        requested: *start,
94                        available: total_size,
95                    });
96                }
97                Ok(*start..total_size)
98            }
99            ByteRange::Suffix(count) => {
100                if *count > total_size {
101                    Ok(0..total_size)
102                } else {
103                    Ok(total_size - count..total_size)
104                }
105            }
106            ByteRange::All => Ok(0..total_size),
107        }
108    }
109
110    /// Check if this range overlaps with another
111    pub fn overlaps(&self, other: &ByteRange, total_size: u64) -> bool {
112        if let (Ok(r1), Ok(r2)) = (self.to_range(total_size), other.to_range(total_size)) {
113            r1.start < r2.end && r2.start < r1.end
114        } else {
115            false
116        }
117    }
118
119    /// Merge two ranges if they overlap or are adjacent
120    pub fn merge(&self, other: &ByteRange, total_size: u64) -> Option<ByteRange> {
121        if let (Ok(r1), Ok(r2)) = (self.to_range(total_size), other.to_range(total_size)) {
122            if r1.start <= r2.end && r2.start <= r1.end {
123                let start = r1.start.min(r2.start);
124                let end = (r1.end - 1).max(r2.end - 1);
125                Some(ByteRange::FromTo { start, end })
126            } else {
127                None
128            }
129        } else {
130            None
131        }
132    }
133
134    /// Get the size of this range
135    pub fn size(&self, total_size: u64) -> u64 {
136        self.to_range(total_size)
137            .map(|r| r.end - r.start)
138            .unwrap_or(0)
139    }
140}
141
142/// Range request for a specific block
143#[derive(Debug, Clone, Serialize, Deserialize)]
144pub struct RangeRequest {
145    /// CID of the block
146    #[serde(serialize_with = "serialize_cid", deserialize_with = "deserialize_cid")]
147    pub cid: Cid,
148    /// Byte range to request
149    pub range: ByteRange,
150    /// Priority (higher = more important)
151    pub priority: i32,
152}
153
154impl RangeRequest {
155    /// Create a new range request
156    pub fn new(cid: Cid, range: ByteRange) -> Self {
157        Self {
158            cid,
159            range,
160            priority: 0,
161        }
162    }
163
164    /// Create a range request with priority
165    pub fn with_priority(cid: Cid, range: ByteRange, priority: i32) -> Self {
166        Self {
167            cid,
168            range,
169            priority,
170        }
171    }
172}
173
174/// Range response containing partial block data
175#[derive(Debug, Clone, Serialize, Deserialize)]
176pub struct RangeResponse {
177    /// CID of the block
178    #[serde(serialize_with = "serialize_cid", deserialize_with = "deserialize_cid")]
179    pub cid: Cid,
180    /// Byte range of this response
181    pub range: Range<u64>,
182    /// Partial block data
183    pub data: Vec<u8>,
184    /// Total size of the complete block
185    pub total_size: u64,
186}
187
188impl RangeResponse {
189    /// Create a new range response
190    pub fn new(cid: Cid, range: Range<u64>, data: Vec<u8>, total_size: u64) -> Self {
191        Self {
192            cid,
193            range,
194            data,
195            total_size,
196        }
197    }
198
199    /// Check if this response satisfies a request
200    pub fn satisfies(&self, request: &RangeRequest) -> bool {
201        if self.cid != request.cid {
202            return false;
203        }
204        if let Ok(req_range) = request.range.to_range(self.total_size) {
205            self.range.start <= req_range.start && self.range.end >= req_range.end
206        } else {
207            false
208        }
209    }
210
211    /// Extract data for a specific sub-range
212    pub fn extract_range(&self, range: &Range<u64>) -> Result<Vec<u8>, RangeError> {
213        if range.start < self.range.start || range.end > self.range.end {
214            return Err(RangeError::OutOfBounds {
215                requested: range.end,
216                available: self.range.end,
217            });
218        }
219
220        let offset = (range.start - self.range.start) as usize;
221        let len = (range.end - range.start) as usize;
222
223        if offset + len > self.data.len() {
224            return Err(RangeError::OutOfBounds {
225                requested: (offset + len) as u64,
226                available: self.data.len() as u64,
227            });
228        }
229
230        Ok(self.data[offset..offset + len].to_vec())
231    }
232}
233
234/// Manager for assembling partial blocks from range responses
235pub struct RangeAssembler {
236    /// CID of the block being assembled
237    cid: Cid,
238    /// Total size of the block
239    total_size: u64,
240    /// Received ranges and their data
241    received: Vec<(Range<u64>, Vec<u8>)>,
242}
243
244impl RangeAssembler {
245    /// Create a new range assembler
246    pub fn new(cid: Cid, total_size: u64) -> Self {
247        Self {
248            cid,
249            total_size,
250            received: Vec::new(),
251        }
252    }
253
254    /// Add a range response
255    pub fn add_range(&mut self, response: RangeResponse) -> Result<(), RangeError> {
256        if response.cid != self.cid {
257            return Err(RangeError::InvalidRange("CID mismatch".to_string()));
258        }
259
260        if response.total_size != self.total_size {
261            return Err(RangeError::InvalidRange("Total size mismatch".to_string()));
262        }
263
264        self.received.push((response.range, response.data));
265        Ok(())
266    }
267
268    /// Check if the block is complete
269    pub fn is_complete(&self) -> bool {
270        let mut covered = vec![false; self.total_size as usize];
271
272        for (range, _) in &self.received {
273            for i in range.start..range.end {
274                if (i as usize) < covered.len() {
275                    covered[i as usize] = true;
276                }
277            }
278        }
279
280        covered.iter().all(|&x| x)
281    }
282
283    /// Get missing ranges
284    pub fn missing_ranges(&self) -> Vec<Range<u64>> {
285        let mut covered = vec![false; self.total_size as usize];
286
287        for (range, _) in &self.received {
288            for i in range.start..range.end {
289                if (i as usize) < covered.len() {
290                    covered[i as usize] = true;
291                }
292            }
293        }
294
295        let mut missing = Vec::new();
296        let mut start = None;
297
298        for (i, &is_covered) in covered.iter().enumerate() {
299            if !is_covered && start.is_none() {
300                start = Some(i as u64);
301            } else if is_covered && start.is_some() {
302                missing.push(start.expect("just checked start.is_some()")..i as u64);
303                start = None;
304            }
305        }
306
307        if let Some(s) = start {
308            missing.push(s..self.total_size);
309        }
310
311        missing
312    }
313
314    /// Assemble the complete block
315    pub fn assemble(&self) -> Result<Vec<u8>, RangeError> {
316        if !self.is_complete() {
317            return Err(RangeError::InvalidRange("Block incomplete".to_string()));
318        }
319
320        let mut data = vec![0u8; self.total_size as usize];
321
322        for (range, chunk) in &self.received {
323            let start = range.start as usize;
324            let end = range.end as usize;
325            let len = end - start;
326
327            if chunk.len() != len {
328                return Err(RangeError::InvalidRange("Chunk size mismatch".to_string()));
329            }
330
331            data[start..end].copy_from_slice(chunk);
332        }
333
334        Ok(data)
335    }
336
337    /// Get completion percentage
338    pub fn completion_percentage(&self) -> f64 {
339        let mut covered = vec![false; self.total_size as usize];
340
341        for (range, _) in &self.received {
342            for i in range.start..range.end {
343                if (i as usize) < covered.len() {
344                    covered[i as usize] = true;
345                }
346            }
347        }
348
349        let covered_count = covered.iter().filter(|&&x| x).count();
350        (covered_count as f64 / self.total_size as f64) * 100.0
351    }
352}
353
354#[cfg(test)]
355mod tests {
356    use super::*;
357
358    fn test_cid() -> Cid {
359        "bafybeigdyrzt5sfp7udm7hu76uh7y26nf3efuylqabf3oclgtqy55fbzdi"
360            .parse()
361            .expect("test: valid CID string")
362    }
363
364    #[test]
365    fn test_byte_range_from_to() {
366        let range = ByteRange::from_to(0, 99).expect("test: create byte range");
367        assert_eq!(
368            range.to_range(1000).expect("test: convert to range"),
369            0..100
370        );
371    }
372
373    #[test]
374    fn test_byte_range_from() {
375        let range = ByteRange::from(500);
376        assert_eq!(
377            range.to_range(1000).expect("test: convert to range"),
378            500..1000
379        );
380    }
381
382    #[test]
383    fn test_byte_range_suffix() {
384        let range = ByteRange::suffix(100);
385        assert_eq!(
386            range.to_range(1000).expect("test: convert to range"),
387            900..1000
388        );
389    }
390
391    #[test]
392    fn test_byte_range_all() {
393        let range = ByteRange::All;
394        assert_eq!(
395            range.to_range(1000).expect("test: convert to range"),
396            0..1000
397        );
398    }
399
400    #[test]
401    fn test_byte_range_out_of_bounds() {
402        let range = ByteRange::from_to(0, 1500).expect("test: create byte range");
403        assert!(range.to_range(1000).is_err());
404    }
405
406    #[test]
407    fn test_byte_range_invalid() {
408        assert!(ByteRange::from_to(100, 50).is_err());
409    }
410
411    #[test]
412    fn test_byte_range_overlaps() {
413        let range1 = ByteRange::from_to(0, 99).expect("test: create byte range");
414        let range2 = ByteRange::from_to(50, 149).expect("test: create byte range");
415        assert!(range1.overlaps(&range2, 1000));
416
417        let range3 = ByteRange::from_to(200, 299).expect("test: create byte range");
418        assert!(!range1.overlaps(&range3, 1000));
419    }
420
421    #[test]
422    fn test_byte_range_merge() {
423        let range1 = ByteRange::from_to(0, 99).expect("test: create byte range");
424        let range2 = ByteRange::from_to(50, 149).expect("test: create byte range");
425        let merged = range1.merge(&range2, 1000).expect("test: merge ranges");
426        assert_eq!(
427            merged.to_range(1000).expect("test: convert to range"),
428            0..150
429        );
430    }
431
432    #[test]
433    fn test_byte_range_size() {
434        let range = ByteRange::from_to(100, 199).expect("test: create byte range");
435        assert_eq!(range.size(1000), 100);
436    }
437
438    #[test]
439    fn test_range_request() {
440        let cid = test_cid();
441        let range = ByteRange::from_to(0, 99).expect("test: create byte range");
442        let req = RangeRequest::new(cid, range);
443        assert_eq!(req.priority, 0);
444
445        let req2 = RangeRequest::with_priority(cid, range, 10);
446        assert_eq!(req2.priority, 10);
447    }
448
449    #[test]
450    fn test_range_response_satisfies() {
451        let cid = test_cid();
452        let range = ByteRange::from_to(0, 99).expect("test: create byte range");
453        let req = RangeRequest::new(cid, range);
454
455        let response = RangeResponse::new(cid, 0..100, vec![0u8; 100], 1000);
456        assert!(response.satisfies(&req));
457
458        let response2 = RangeResponse::new(cid, 50..150, vec![0u8; 100], 1000);
459        assert!(!response2.satisfies(&req));
460    }
461
462    #[test]
463    fn test_range_response_extract() {
464        let cid = test_cid();
465        let data = (0..100).collect::<Vec<u8>>();
466        let response = RangeResponse::new(cid, 0..100, data.clone(), 1000);
467
468        let extracted = response
469            .extract_range(&(10..20))
470            .expect("test: extract range");
471        assert_eq!(extracted, &data[10..20]);
472    }
473
474    #[test]
475    fn test_range_assembler() {
476        let cid = test_cid();
477        let mut assembler = RangeAssembler::new(cid, 100);
478
479        assert!(!assembler.is_complete());
480        assert_eq!(assembler.completion_percentage(), 0.0);
481
482        let resp1 = RangeResponse::new(cid, 0..50, vec![1u8; 50], 100);
483        assembler
484            .add_range(resp1)
485            .expect("test: add range to assembler");
486        assert_eq!(assembler.completion_percentage(), 50.0);
487
488        let resp2 = RangeResponse::new(cid, 50..100, vec![2u8; 50], 100);
489        assembler
490            .add_range(resp2)
491            .expect("test: add range to assembler");
492        assert!(assembler.is_complete());
493        assert_eq!(assembler.completion_percentage(), 100.0);
494
495        let data = assembler.assemble().expect("test: assemble ranges");
496        assert_eq!(data.len(), 100);
497        assert_eq!(&data[0..50], &vec![1u8; 50][..]);
498        assert_eq!(&data[50..100], &vec![2u8; 50][..]);
499    }
500
501    #[test]
502    fn test_range_assembler_missing_ranges() {
503        let cid = test_cid();
504        let mut assembler = RangeAssembler::new(cid, 100);
505
506        let resp1 = RangeResponse::new(cid, 0..25, vec![0u8; 25], 100);
507        assembler
508            .add_range(resp1)
509            .expect("test: add range to assembler");
510
511        let resp2 = RangeResponse::new(cid, 75..100, vec![0u8; 25], 100);
512        assembler
513            .add_range(resp2)
514            .expect("test: add range to assembler");
515
516        let missing = assembler.missing_ranges();
517        assert_eq!(missing, vec![25..75]);
518    }
519
520    #[test]
521    fn test_range_assembler_overlapping() {
522        let cid = test_cid();
523        let mut assembler = RangeAssembler::new(cid, 100);
524
525        let resp1 = RangeResponse::new(cid, 0..60, vec![1u8; 60], 100);
526        assembler
527            .add_range(resp1)
528            .expect("test: add range to assembler");
529
530        let resp2 = RangeResponse::new(cid, 40..100, vec![2u8; 60], 100);
531        assembler
532            .add_range(resp2)
533            .expect("test: add range to assembler");
534
535        assert!(assembler.is_complete());
536    }
537
538    #[test]
539    fn test_range_assembler_incomplete() {
540        let cid = test_cid();
541        let mut assembler = RangeAssembler::new(cid, 100);
542
543        let resp = RangeResponse::new(cid, 0..50, vec![0u8; 50], 100);
544        assembler
545            .add_range(resp)
546            .expect("test: add range to assembler");
547
548        assert!(assembler.assemble().is_err());
549    }
550}
ipfrs_transport/range_request.rs

ipfrs_transport/
range_request.rs
