diff --git a/tlsn/tlsn-core/src/span/mod.rs b/tlsn/tlsn-core/src/span/mod.rs index cc1efcca29..1411d653e4 100644 --- a/tlsn/tlsn-core/src/span/mod.rs +++ b/tlsn/tlsn-core/src/span/mod.rs @@ -44,7 +44,7 @@ pub fn invert_ranges( ranges: Vec>, len: usize, ) -> Result>, SpanError> { - for range in ranges.iter() { + for (k, range) in ranges.iter().enumerate() { // Check that there is no invalid or empty range if range.start >= range.end { return Err(SpanError::InvalidRange); @@ -58,7 +58,8 @@ pub fn invert_ranges( // Check that ranges are not overlapping if ranges .iter() - .any(|r| r.start < range.end && r.end > range.start) + .enumerate() + .any(|(l, r)| k != l && r.start < range.end && r.end > range.start) { return Err(SpanError::InvalidRange); } @@ -70,18 +71,21 @@ pub fn invert_ranges( for range in ranges.iter() { let inv = inverted .iter_mut() - .find(|inv| range.start >= inv.start) - .unwrap(); + .find(|inv| range.start >= inv.start && range.end <= inv.end) + .expect("Should have found range to invert"); - let original_len = inv.end; + let original_end = inv.end; inv.end = range.start; inverted.push(Range { start: range.end, - end: original_len, + end: original_end, }); } + // Remove empty ranges + inverted.retain(|r| r.start != r.end); + Ok(inverted) } @@ -96,3 +100,60 @@ pub enum SpanError { #[error("Custom error: {0}")] Custom(String), } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_invert_ranges_errors() { + let empty_range = Range { start: 0, end: 0 }; + let invalid_range = Range { start: 2, end: 1 }; + let out_of_bounds = Range { start: 4, end: 11 }; + + let ranges = vec![empty_range, invalid_range, out_of_bounds]; + + for range in ranges { + assert!(invert_ranges(vec![range], 10).is_err()); + } + } + + #[test] + fn test_invert_ranges_overlapping() { + let overlapping1 = vec![Range { start: 2, end: 5 }, Range { start: 4, end: 7 }]; + let overlapping2 = vec![Range { start: 2, end: 5 }, Range { start: 1, end: 4 }]; + let overlapping3 = vec![Range { start: 2, end: 5 }, Range { start: 3, end: 4 }]; + let overlapping4 = vec![Range { start: 2, end: 5 }, Range { start: 2, end: 5 }]; + + // this should not be an error + let ok1 = vec![Range { start: 2, end: 5 }, Range { start: 5, end: 8 }]; + let ok2 = vec![Range { start: 2, end: 5 }, Range { start: 7, end: 10 }]; + + let overlap = vec![overlapping1, overlapping2, overlapping3, overlapping4]; + let ok = vec![ok1, ok2]; + + for range in overlap { + assert!(invert_ranges(range, 10).is_err()); + } + + for range in ok { + assert!(invert_ranges(range, 10).is_ok()); + } + } + + #[test] + fn test_invert_ranges() { + let len = 20; + + let ranges = vec![ + Range { start: 0, end: 5 }, + Range { start: 5, end: 10 }, + Range { start: 12, end: 16 }, + Range { start: 18, end: 20 }, + ]; + + let expected = vec![Range { start: 10, end: 12 }, Range { start: 16, end: 18 }]; + + assert_eq!(invert_ranges(ranges, len).unwrap(), expected); + } +}