diff --git a/Cargo.toml b/Cargo.toml index 716635a..7785aa3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ sha2 = "0.8.1" sleep-parser = "0.8.0" sparse-bitfield = "0.11.0" tree-index = "0.6.0" +bitfield-rle = "0.1.1" [dev-dependencies] quickcheck = "0.9.2" diff --git a/src/bitfield/mod.rs b/src/bitfield/mod.rs index 833cc47..367c319 100644 --- a/src/bitfield/mod.rs +++ b/src/bitfield/mod.rs @@ -229,6 +229,38 @@ impl Bitfield { } } + // TODO: use the index to speed this up *a lot* + /// https://github.com/mafintosh/hypercore/blob/06f3a1f573cb74ee8cfab2742455318fbf7cc3a2/lib/bitfield.js#L111-L126 + pub fn compress(&self, start: usize, length: usize) -> std::io::Result> { + // On Node versions this fields might not be present on the want/request message + // When both start and length are not present (!0 in node is false), return all data bytes encoded + if start == 0 && length == 0 { + return Ok(bitfield_rle::encode(&self.data.to_bytes()?)); + } + + use std::io::{Cursor, Write}; + let mut buf = Cursor::new(Vec::with_capacity(length)); + + let page_size = self.data.page_size() as f64; + let mut p = start as f64 / page_size / 8.0; + let end = p + length as f64 / page_size / 8.0; + let offset = p * page_size; + + while p < end { + let index = p as usize; + let page = self.data.pages.get(index); + if let Some(page) = page { + if page.len() != 0 { + buf.set_position((p * page_size - offset) as u64); + buf.write_all(&page)?; + } + } + p += 1.0; + } + + Ok(bitfield_rle::encode(&buf.into_inner())) + } + /// Constructs an iterator from start to end pub fn iterator(&mut self) -> iterator::Iterator<'_> { let len = self.length; diff --git a/src/feed.rs b/src/feed.rs index a1433ea..34a4232 100644 --- a/src/feed.rs +++ b/src/feed.rs @@ -524,6 +524,11 @@ where }) } + /// Expose the bitfield attribute to use on during download + pub fn bitfield(&self) -> &Bitfield { + &self.bitfield + } + /// (unimplemented) Provide a range of data to download. pub fn download(&mut self, _range: Range) -> Result<()> { unimplemented!(); diff --git a/tests/bitfield.rs b/tests/bitfield.rs index bc02642..b90ee7d 100644 --- a/tests/bitfield.rs +++ b/tests/bitfield.rs @@ -174,3 +174,25 @@ fn bitfield_dedup() { assert!(!b.get(8 * 1024)); assert!(b.get(16 * 1024)); } + +#[test] +fn bitfield_compress() { + let mut b = Bitfield::new(); + assert_eq!(b.compress(0, 0).unwrap(), vec![0]); + + b.set(1, true); + assert_eq!(b.compress(0, 0).unwrap(), vec![2, 64, 253, 31]); + + b.set(1_424_244, true); + assert_eq!( + b.compress(0, 0).unwrap(), + vec![2, 64, 181, 187, 43, 2, 8, 197, 4] + ); + assert_eq!(b.compress(0, 1).unwrap(), vec![2, 64, 253, 31]); + assert_eq!( + b.compress(1_424_244, 1).unwrap(), + vec![185, 27, 2, 8, 197, 4] + ); + + assert_eq!(b.compress(1_424_244_000, 1).unwrap(), vec![0]); +} diff --git a/tests/model.rs b/tests/model.rs index b7e62fc..5e71999 100644 --- a/tests/model.rs +++ b/tests/model.rs @@ -1,12 +1,7 @@ -#[macro_use] -extern crate quickcheck; -extern crate hypercore; -extern crate rand; - mod common; use common::create_feed; -use quickcheck::{Arbitrary, Gen}; +use quickcheck::{quickcheck, Arbitrary, Gen}; use rand::seq::SliceRandom; use rand::Rng; use std::u8;