Refactor in progress

src/block.rs

@@ -1,8 +1,10 @@
 use std::fmt::Formatter;
 use std::str::FromStr;
 
-use rand::Rng;
+use bitvec::slice::BitSlice;
+use bitvec::vec::BitVec;
 use rand::prelude::SliceRandom;
+use rand::Rng;
 
 use crate::embed::EmbedAlgorithm;
 
@@ -42,7 +44,7 @@ impl BlockMode {
 
 		Self {
 			len: 1usize << len,
-			crc_len: (crc * len) as usize
+			crc_len: (crc * len) as usize,
 		}
 	}
 }
@@ -76,50 +78,238 @@ impl FromStr for BlockMode {
 
 #[derive(Debug)]
 pub struct BlockPlacement<'a> {
+	algorithm: &'a EmbedAlgorithm,
 	data: &'a mut [u8],
 	block_size: usize,
 	blocks: Vec<usize>,
 }
 
-impl<'a> BlockPlacement<'a>
-{
+impl<'a> BlockPlacement<'a> {
 	// Attempts to create a new block placement
 	//
 	// # Errors
 	//
 	// Will fail if the data is too small to hold all the blocks
-	pub fn new<R>(data: &'a mut [u8], block_size: usize, algorithm: &EmbedAlgorithm, embed_size: usize, rng: &mut R) -> Result<Self, String>
-    	where R: Rng + ?Sized
+	pub fn new<R>(
+		data: &'a mut [u8],
+		algorithm: &'a EmbedAlgorithm,
+		block_size: usize,
+		embed_size: usize,
+		rng: &mut R,
+	) -> Result<Self, String>
+	where
+		R: Rng + ?Sized,
 	{
 		// Total size of the embed (crc included)
 		let embedded_size = algorithm.embedded_size(embed_size);
 
+		// Maximum number of blocks
+		let max_blocks = data.len() / block_size;
+
 		// Number of blocks
 		let blocks_num = (embedded_size as f64 / block_size as f64).ceil() as usize;
 
-
-		// Safe distance for spacing the blocks equally
-		let safe_spacing = data.len() / blocks_num;
-		if safe_spacing*blocks_num < embedded_size {
-			return Err(format!("Failed to determine a safe spacing size: {safe_spacing} < {}", embedded_size as f64 / blocks_num as f64))
+		if blocks_num > max_blocks {
+			return Err(format!(
+				"Too many blocks required: {blocks_num}, maximum: {max_blocks}"
+			));
 		}
 
 		// Blocks in the resulting data
-		let mut blocks = Vec::with_capacity(blocks_num);
-		for i in 0 .. blocks_num {
-			// Choose a random position within [0, safe_spacing[ for the block
-			let pos = rng.gen_range(i*safe_spacing..(i+1)*safe_spacing);
-
-			blocks.push(pos);
-		}
+		let mut blocks = (0..max_blocks).collect::<Vec<_>>();
 
 		// Shuffle the block order
 		blocks.shuffle(rng);
 
+		// Only keep the first blocks_num blocks
+		blocks.resize(blocks_num, 0);
+
 		Ok(Self {
+			algorithm,
 			data,
 			block_size,
-			blocks
+			blocks,
 		})
 	}
+
+	pub fn write_embed(&mut self) 
+}
+
+// Iterator over blocks in the resulting image
+pub struct BlockPlacementIterator<'a> {
+	algorithm: &'a EmbedAlgorithm,
+	data: &'a [u8],
+	block_size: usize,
+
+	// Block index
+	index: usize,
+	// Position of the blocks
+	blocks: Vec<usize>,
+}
+
+impl<'a> BlockPlacementIterator<'a> {
+	pub fn new<R: Rng + ?Sized>(
+		algorithm: &'a EmbedAlgorithm,
+		data: &'a [u8],
+		block_size: usize,
+		rng: &mut R,
+	) -> Self {
+		// Maximum number of blocks
+		let max_blocks = data.len() / block_size;
+
+		// Blocks
+		let mut blocks = (0..max_blocks).collect::<Vec<_>>();
+
+		// Shuffle the block order
+		blocks.shuffle(rng);
+
+		Self {
+			algorithm,
+			data,
+			block_size,
+			index: 0,
+			blocks,
+		}
+	}
+}
+
+impl<'a> Iterator for BlockPlacementIterator<'a> {
+	type Item = Block<'a>;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		if self.index == self.blocks.len() {
+			return None;
+		}
+
+		let pos = self.blocks[self.index] * self.block_size;
+		let slice = &self.data[pos..pos + self.block_size];
+		self.index += 1;
+
+		Some(Block(self.algorithm, slice))
+	}
+}
+
+// Block of data in the resulting image
+#[derive(Debug)]
+pub struct Block<'a>(&'a EmbedAlgorithm, &'a [u8]);
+
+// Iterator to read embedded data inside a block
+pub struct BlockIterator<'a> {
+	// Block of the iterator
+	block: &'a Block<'a>,
+
+	// Byte position in [`data`]
+	index: usize,
+
+	// Remainder, i.e. bits that have been read and will be part of the `next` byte
+	remainder: BitVec<u8>,
+}
+
+impl<'a> BlockIterator<'a> {
+	pub fn new(block: &'a Block, previous: Option<BlockIterator>) -> Self {
+		if let Some(previous) = previous {
+			Self {
+				block,
+				index: 0,
+				remainder: previous.remainder,
+			}
+		} else {
+			Self {
+				block,
+				index: 0,
+				remainder: BitVec::with_capacity(7),
+			}
+		}
+	}
+}
+
+impl<'a> Iterator for BlockIterator<'a> {
+	type Item = u8;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		let mut byte = 0u8;
+
+		// Read remainder
+		let mut bit_idx = 0;
+		for bit in &self.remainder {
+			byte |= (*bit as u8) << bit_idx;
+			bit_idx += 1;
+		}
+		self.remainder.clear();
+
+		match self.block.0 {
+			EmbedAlgorithm::Lo(bits) => {
+				while bit_idx < 8 {
+					// End of data
+					if self.index == self.block.1.len() {
+						return None;
+					}
+
+					// Read next byte
+					let next = self.block.1[self.index];
+					self.index += 1;
+
+					for i in 0..*bits {
+						if bit_idx < 8 {
+							// Prepend bit to result
+							byte |= ((next >> i) & 0b1) << bit_idx;
+						} else {
+							// Append bit to remainder
+							self.remainder.push((next >> i) & 0b1 == 0b1)
+						}
+
+						bit_idx += 1;
+					}
+				}
+
+				Some(byte)
+			}
+		}
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use rand::SeedableRng;
+	use rand_chacha::ChaCha8Rng;
+
+	use super::*;
+
+	#[test]
+	fn block_iterator() {
+		let algorithm = EmbedAlgorithm::Lo(3);
+		let data = vec![
+			0b10111000, 0b11111001, 0b01101010, 0b00111011, 0b11011100, 0b11100110, 0b01100111,
+			0b01100000,
+		];
+
+		let block = Block(&algorithm, &data);
+		let mut it = BlockIterator::new(&block, None);
+
+		assert_eq!(it.next(), Some(0b10_001_000));
+		assert_eq!(it.next(), Some(0b0_100_011_0));
+		assert_eq!(it.next(), Some(0b000_111_11));
+	}
+
+	#[test]
+	fn blockplacement_iterator() {
+		let algorithm = EmbedAlgorithm::Lo(4);
+		let mut data = vec![
+			0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
+			0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b,
+			0x1c, 0x1d, 0x1e, 0x1f,
+		];
+
+		let mut rand = ChaCha8Rng::from_seed([1u8; 32]);
+		let mut it = BlockPlacementIterator::new::<_>(&algorithm, &data, 4, &mut rand);
+
+		let mut rand = ChaCha8Rng::from_seed([1u8; 32]);
+		let mut positions = (0..8).collect::<Vec<_>>();
+		positions.shuffle(&mut rand);
+
+		for i in 0..8 {
+			let block = it.next().unwrap();
+			assert_eq!(block.1[0] / 4, positions[i]);
+		}
+	}
 }

src/embed.rs

@@ -1,12 +1,13 @@
 use std::fmt::Formatter;
 use std::str::FromStr;
 
-use bitvec::prelude::*;
+use bitvec::slice::BitSlice;
 use bitvec::vec::BitVec;
 
 use crate::block::BlockMode;
 use crate::image::ImageInfo;
 
+#[derive(Debug)]
 pub enum EmbedAlgorithm {
 	Lo(u8),
 }
@@ -15,9 +16,7 @@ impl EmbedAlgorithm {
 	/// Get the size of the data (in bytes) once embedded by the algorithm
 	pub fn embedded_size(&self, size: usize) -> usize {
 		match self {
-			EmbedAlgorithm::Lo(bits) => {
-				((size * 8) as f64 / *bits as f64).ceil() as usize
-			}
+			EmbedAlgorithm::Lo(bits) => ((size * 8) as f64 / *bits as f64).ceil() as usize,
 		}
 	}
 
@@ -26,31 +25,39 @@ impl EmbedAlgorithm {
 
 		match self {
 			EmbedAlgorithm::Lo(bits) => {
-				(((blockmode.len - blockmode.crc_len)*blocks_num) as f64 * (*bits as f64) / 8f64).floor() as usize
+				(((blockmode.len - blockmode.crc_len) * blocks_num) as f64 * (*bits as f64) / 8f64)
+					.floor() as usize
 			}
 		}
 	}
 
-	pub fn next_block(&self, original_data: &mut [u8], mut data_pos: usize, embed_data: &BitVec<u8>, mut embed_offset: usize, blockmode: &BlockMode) -> (usize, usize) {
+	pub fn next_block(
+		&self,
+		original_data: &mut [u8],
+		mut data_pos: usize,
+		embed_data: &BitVec<u8>,
+		mut embed_offset: usize,
+		blockmode: &BlockMode,
+	) -> (usize, usize) {
 		match self {
 			EmbedAlgorithm::Lo(bits) => {
 				let mask = (1 << bits) - 1;
 
-				fn bits_to_byte(slice: &BitSlice<u8>, bits: u8) -> u8
-				{
+				fn bits_to_byte(slice: &BitSlice<u8>, bits: u8) -> u8 {
 					let mut result: u8 = 0;
-					for i in 0..bits
-					{
+					for i in 0..bits {
 						result |= (slice[i as usize] as u8) << i;
 					}
 					result
 				}
 
 				let start = embed_offset;
-				while embed_offset-start < (blockmode.len-blockmode.crc_len)*8
-				{
+				while embed_offset - start < (blockmode.len - blockmode.crc_len) * 8 {
 					let hi = std::cmp::min(*bits as usize, embed_data.len() - embed_offset);
-					let embed = bits_to_byte(embed_data.get(embed_offset..embed_offset+hi).unwrap(), hi as u8);
+					let embed = bits_to_byte(
+						embed_data.get(embed_offset..embed_offset + hi).unwrap(),
+						hi as u8,
+					);
 
 					original_data[data_pos] &= !mask;
 					original_data[data_pos] |= embed;
@@ -66,20 +73,23 @@ impl EmbedAlgorithm {
 		(data_pos, embed_offset)
 	}
 
-	pub fn read_block(&self, encoded_data: &[u8], mut data_pos: usize, incoming: &mut BitVec<u8>, blockmode: &BlockMode) -> usize {
+	pub fn read_block(
+		&self,
+		encoded_data: &[u8],
+		mut data_pos: usize,
+		incoming: &mut BitVec<u8>,
+		blockmode: &BlockMode,
+	) -> usize {
 		match self {
 			EmbedAlgorithm::Lo(bits) => {
-				fn push(vec: &mut BitVec<u8>, bits: u8, b: u8)
-				{
-					for i in 0..bits
-					{
+				fn push(vec: &mut BitVec<u8>, bits: u8, b: u8) {
+					for i in 0..bits {
 						vec.push((b >> i) & 0b1 == 0b1)
 					}
 				}
 
 				let start = incoming.len();
-				while incoming.len()-start < (blockmode.len-blockmode.crc_len)*8
-				{
+				while incoming.len() - start < (blockmode.len - blockmode.crc_len) * 8 {
 					push(incoming, *bits, encoded_data[data_pos]);
 					data_pos += 1;
 				}

src/header.rs

@@ -1,13 +1,49 @@
-use bitvec::slice::BitSlice;
+use bitvec::{slice::BitSlice, vec::BitVec};
+use crc::Crc;
 
 use crate::block::BlockMode;
 
-pub struct Header {
-	pub blockmode: BlockMode,
-	pub comment: Option<String>,
+#[repr(u16)]
+#[allow(non_camel_case_types)]
+#[derive(Debug, Clone, Copy)]
+pub enum Version {
+	VERSION_1,
 }
 
+pub trait Encode {
+	// Encode the data
+	fn encode(&self, vec: &mut Vec<u8>);
+}
+
+//pub trait Decode {
+//	fn decode(incoming: &mut EmbedIterator) -> (usize, Self);
+//}
+
+#[derive(Debug)]
+pub struct Header {
+	version: Version,
+	block_size: usize,
+	data_len: u32,
+	data_crc: u32,
+	comment: Option<String>,
+}
+
+
 impl Header {
+	pub fn new(version: Version, block_size: usize, data: &[u8], comment: Option<String>) -> Self {
+		assert_eq!((data.len() as u32) as usize, data.len());
+		assert_eq!(1 << usize::trailing_zeros(block_size), block_size);
+		assert!(comment.as_ref().map_or(0, |c| c.len()) < u16::MAX as usize);
+
+		Self {
+			version,
+			block_size,
+			data_len: data.len() as u32,
+			data_crc: Crc::<u32>::new(&crc::CRC_32_CKSUM).checksum(data),
+			comment,
+		}
+	}
+	/*
 	pub fn to_data(&self, version: u16, embed_len: u32) -> Vec<u8> {
 		let mut header = vec![];
 
@@ -55,4 +91,30 @@ impl Header {
 
 		(version, blockmode, len, comment_len)
 	}
+	*/
+}
+
+impl Encode for Header {
+    fn encode(&self, vec: &mut Vec<u8>) {
+		// Version
+		vec.extend_from_slice((self.version as u16).to_le_bytes().as_slice());
+
+		// Block size
+		vec.push((usize::trailing_zeros(self.block_size) as u8).to_le());
+
+		// Data Len
+		vec.extend_from_slice(self.data_len.to_le_bytes().as_slice());
+
+		// Data CRC
+		vec.extend_from_slice(self.data_crc.to_le_bytes().as_slice());
+
+		// Comment length
+		let comment_length = self.comment.as_ref().map_or(0u16, |c| c.len() as u16);
+		vec.extend_from_slice(comment_length.to_le_bytes().as_slice());
+
+		// Comment
+		if let Some(comment) = &self.comment {
+			vec.extend_from_slice(comment.as_bytes());
+		}
+    }
 }

src/main.rs

@@ -43,6 +43,7 @@ Public domain\n"
 	);
 }
 
+/*
 impl ImageInfo for png::OutputInfo {
 	fn width(&self) -> u32 { self.width }
 
@@ -202,9 +203,9 @@ fn encode(image: String, matches: Matches) -> Result<Vec<u8>, String> {
 		.map_err(|err| format!("Failed to read embed file `{embed_file}`: {err}"))?;
 
 	let mut rand = ChaCha8Rng::from_seed(seed);
-	//let placement = BlockPlacement::new(data.as_mut_slice(), blockmode.len, &algorithm, embed_data.len(), &mut rand)?;
+	let placement = BlockPlacement::new(data.as_mut_slice(), blockmode.len, &algorithm, embed_data.len(), &mut rand)?;
 
-	//return Ok(vec![]);
+	return Ok(vec![]);
 
 	// Get header
 	let header = Header {
@@ -295,6 +296,7 @@ fn encode(image: String, matches: Matches) -> Result<Vec<u8>, String> {
 
 	Ok(vec![])
 }
+*/
 
 fn main() -> ExitCode {
 	let args: Vec<String> = env::args().collect();
@@ -331,6 +333,7 @@ fn main() -> ExitCode {
 		return ExitCode::FAILURE;
 	}
 
+	/*
 	let input = matches.free[0].clone();
 
 	if matches.opt_present("z") {
@@ -350,6 +353,7 @@ fn main() -> ExitCode {
 			}
 		}
 	}
+	*/
 
 	ExitCode::SUCCESS
 }