use std::{ num::NonZeroUsize, time::{SystemTime, UNIX_EPOCH} }; use flume::{Receiver, SendError, Sender, TryRecvError}; use futures_util::stream::StreamExt as _; use image::DynamicImage; use kittage::{NumberOrId, action::NONZERO_ONE}; use ratatui::layout::Rect; use ratatui_image::{ Resize, picker::{Picker, ProtocolType}, protocol::Protocol }; use rayon::iter::ParallelIterator as _; use crate::{ renderer::{PageInfo, RenderError, fill_default}, skip::InterleavedAroundWithMax }; #[derive(Debug)] pub enum MaybeTransferred { NotYet(kittage::image::Image<'static>), Transferred(kittage::ImageId) } #[derive(Debug)] pub enum ConvertedImage { Generic(Protocol), Kitty { img: MaybeTransferred, cell_w: u16, cell_h: u16 } } impl ConvertedImage { #[must_use] pub fn w_h(&self) -> (u16, u16) { match self { Self::Generic(prot) => { let a = prot.area(); (a.width, a.height) } Self::Kitty { img: _, cell_w, cell_h } => (*cell_w, *cell_h) } } } pub struct ConvertedPage { pub page: ConvertedImage, pub num: usize, pub num_results: usize } pub enum ConverterMsg { NumPages(usize), GoToPage(usize), AddImg(PageInfo) } pub async fn run_conversion_loop( sender: Sender>, receiver: Receiver, picker: Picker, prerender: usize, shms_work: bool ) -> Result<(), SendError>> { let mut images = vec![]; let mut page: usize = 0; let pid = std::process::id(); fn next_page( images: &mut [Option], picker: &Picker, page: usize, iteration: &mut usize, prerender: usize, pid: u32, shms_work: bool ) -> Result, RenderError> { if images.is_empty() || *iteration >= prerender { return Ok(None); } // This kinda mimics the way the renderer alternates between going above and below the // current page (within the bounds of how many pages there are) until we've done 20 let idx_start = page.saturating_sub(prerender / 2); let idx_end = idx_start.saturating_add(prerender).min(images.len()); // If there's none to render, then why bother. let Some(idx_end) = NonZeroUsize::new(idx_end) else { return Ok(None); }; // then we go through all the indices available to us and find the first one that has an // image available to steal let Some((page_info, new_iter, page_num)) = InterleavedAroundWithMax::new(page, idx_start, idx_end) .enumerate() .take(prerender) // .skip(*iteration) .find_map(|(i_idx, p_idx)| images[p_idx].take().map(|p| (p, i_idx, p_idx))) else { return Ok(None); }; let mut dyn_img = image::load_from_memory_with_format( &page_info.img_data.pixels, image::ImageFormat::Pnm ) .map_err(|e| RenderError::Converting(format!("Can't load image: {e}")))?; match dyn_img { DynamicImage::ImageRgb8(ref mut img) => for quad in &*page_info.result_rects { img.par_enumerate_pixels_mut() .filter(|(x, y, _)| { *x > quad.ul_x && *x < quad.lr_x && *y > quad.ul_y && *y < quad.lr_y }) .for_each(|(_, _, px)| px.0[2] = px.0[2].saturating_sub(u8::MAX / 2)); }, _ => unreachable!() } let img_area = Rect { width: page_info.img_data.cell_w, height: page_info.img_data.cell_h, x: 0, y: 0 }; let txt_img = match picker.protocol_type() { ProtocolType::Kitty => { let rn = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_nanos() % 1_000_000; let mut img = if shms_work { kittage::image::Image::shm_from(dyn_img, &format!("/tdf_{pid}_{rn}_{page_num}")) .map_err(|e| { RenderError::Converting(format!("Couldn't write to shm: {e:?}")) })? } else { kittage::image::Image::from(dyn_img) }; // if ur pdf has 4 billion pages then you deserve to suffer img.num_or_id = NumberOrId::Id(NONZERO_ONE.saturating_add(page_num as u32)); ConvertedImage::Kitty { img: MaybeTransferred::NotYet(img), cell_w: page_info.img_data.cell_w, cell_h: page_info.img_data.cell_h } } _ => ConvertedImage::Generic( picker .new_protocol(dyn_img, img_area, Resize::None) .map_err(|e| { RenderError::Converting(format!( "Couldn't convert DynamicImage to ratatui image: {e}" )) })? ) }; log::debug!( "got converted page for num {} with results {:?}", page_info.page_num, page_info.result_rects ); // update the iteration to the iteration that we stole this image from *iteration = new_iter; Ok(Some(ConvertedPage { page: txt_img, num: page_info.page_num, num_results: page_info.result_rects.len() })) } fn handle_notif(msg: ConverterMsg, images: &mut Vec>, page: &mut usize) { match msg { ConverterMsg::AddImg(img) => { let page_num = img.page_num; images[page_num] = Some(img); } ConverterMsg::NumPages(n_pages) => { fill_default(images, n_pages); *page = (*page).min(n_pages - 1); } ConverterMsg::GoToPage(new_page) => *page = new_page } } 'outer: loop { let mut iteration = 0; loop { match receiver.try_recv() { Ok(msg) => { handle_notif(msg, &mut images, &mut page); continue 'outer; } Err(TryRecvError::Empty) => (), // if it's disconnected, we're done. just return. Err(TryRecvError::Disconnected) => return Ok(()) } match next_page( &mut images, &picker, page, &mut iteration, prerender, pid, shms_work ) { Ok(None) => break, Ok(Some(img)) => sender.send(Ok(img))?, Err(e) => sender.send(Err(e))? } } let Some(msg) = receiver.stream().next().await else { break; }; handle_notif(msg, &mut images, &mut page); } Ok(()) }