jamt 0.1.28

use codec::Decode;
use corevm_host::fs::{BlockRef, Dir, Error, MainBlock, NodeKind, NodeReader, ReadBlock};
use jam_types::{Hash, ServiceId};
use std::collections::VecDeque;

/// Preimage parameters returned by [`PreimageIter`].
pub struct PreimageParams {
	pub service_id: ServiceId,
	pub hash: Hash,
	pub length: Option<u32>,
}

/// Recursively traverses file system nodes and returns preimage service ID, hash and optionally
/// length for every file block.
///
/// This iterator only reads main blocks and should be more efficient than reading files in full.
pub struct PreimageIter<R: ReadBlock> {
	queue: VecDeque<BlockRef>,
	preimages: VecDeque<PreimageParams>,
	reader: R,
}

impl<R: ReadBlock> PreimageIter<R> {
	pub fn new(main_block_ref: BlockRef, reader: R) -> Self {
		let mut queue = VecDeque::new();
		queue.push_back(main_block_ref);
		Self { queue, preimages: VecDeque::new(), reader }
	}
}

impl<R: ReadBlock> Iterator for PreimageIter<R> {
	type Item = Result<PreimageParams, Error>;

	fn next(&mut self) -> Option<Self::Item> {
		macro_rules! check {
			($body: expr) => {
				match $body {
					Ok(ret) => ret,
					Err(e) => return Some(Err(e)),
				}
			};
		}
		if let Some(preimage) = self.preimages.pop_front() {
			return Some(Ok(preimage));
		}
		let block_ref = self.queue.pop_front()?;
		let block = check!(self.reader.read_block(&block_ref).map_err(|_| Error::Io));
		let block_len = block.len();
		let main_block = check!(MainBlock::decode(block).map_err(|_| Error::Io));
		for block_ref in main_block.block_refs() {
			self.preimages.push_back(PreimageParams {
				service_id: block_ref.service_id,
				hash: block_ref.hash.0,
				length: None,
			});
		}
		if matches!(main_block.kind(), NodeKind::Dir) {
			let mut file = NodeReader::from_main_block(main_block, &mut self.reader);
			let dir = check!(Dir::decode(&mut file).map_err(|_| Error::Io));
			for (_name, block_ref) in dir.0.into_iter() {
				self.queue.push_back(block_ref);
			}
		}
		Some(Ok(PreimageParams {
			service_id: block_ref.service_id,
			hash: block_ref.hash.0,
			length: Some(block_len as u32),
		}))
	}
}