FLAC image extractor tweak

crates/pile-value/src/extract/item/flac.rs

@@ -14,50 +14,79 @@ use crate::{
 
 pub struct FlacImagesExtractor {
 	item: Item,
+	cached_count: OnceLock<usize>,
 }
 
 impl FlacImagesExtractor {
 	pub fn new(item: &Item) -> Self {
-		Self { item: item.clone() }
+		Self {
+			item: item.clone(),
+			cached_count: OnceLock::new(),
+		}
 	}
 
-	async fn get_images(&self) -> Result<Vec<PileValue>, std::io::Error> {
+	async fn get_count(&self) -> Result<usize, std::io::Error> {
 		let reader = SyncReadBridge::new_current(self.item.read().await?);
-		let raw_images = tokio::task::spawn_blocking(move || {
+		let count = tokio::task::spawn_blocking(move || {
 			let reader = FlacReader::new(BufReader::new(reader));
-			let mut images: Vec<(Mime, Vec<u8>)> = Vec::new();
+			let mut count = 0usize;
 			for block in reader {
 				match block.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))? {
-					FlacBlock::Picture(picture) => {
-						images.push((picture.mime, picture.img_data));
-					}
 					FlacBlock::AudioFrame(_) => break,
+					FlacBlock::Picture(_) => count += 1,
 					_ => {}
 				}
 			}
-			Ok::<_, std::io::Error>(images)
+			Ok::<_, std::io::Error>(count)
 		})
 		.await
 		.map_err(std::io::Error::other)??;
 
-		Ok(raw_images
+		return Ok(count);
-			.into_iter()
-			.map(|(mime, data)| PileValue::Blob {
-				mime,
-				bytes: Arc::new(data),
-			})
-			.collect())
 	}
 }
 
 #[async_trait::async_trait]
 impl ListExtractor for FlacImagesExtractor {
-	async fn get<'a>(&'a self, idx: usize) -> Result<Option<PileValue>, std::io::Error> {
+	async fn get<'a>(&'a self, mut idx: usize) -> Result<Option<PileValue>, std::io::Error> {
-		Ok(self.get_images().await?.into_iter().nth(idx))
+		let reader = SyncReadBridge::new_current(self.item.read().await?);
+		let image = tokio::task::spawn_blocking(move || {
+			let reader = FlacReader::new(BufReader::new(reader));
+			let mut out: Option<(Mime, Vec<u8>)> = None;
+			'blocks: for block in reader {
+				match block.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))? {
+					FlacBlock::AudioFrame(_) => break,
+					FlacBlock::Picture(picture) => {
+						if idx > 0 {
+							idx -= 1;
+							continue;
+						}
+
+						out = Some((picture.mime, picture.img_data));
+						break 'blocks;
+					}
+
+					_ => {}
+				}
+			}
+			Ok::<_, std::io::Error>(out)
+		})
+		.await
+		.map_err(std::io::Error::other)??;
+
+		Ok(image.map(|(mime, data)| PileValue::Blob {
+			mime,
+			bytes: Arc::new(data),
+		}))
 	}
 
 	async fn len(&self) -> Result<usize, std::io::Error> {
-		Ok(self.get_images().await?.len())
+		if let Some(x) = self.cached_count.get() {
+			return Ok(*x);
+		}
+
+		let count = self.get_count().await?;
+		return Ok(*self.cached_count.get_or_init(|| count));
 	}
 }