Files
tubular/src/thumbnail.rs
T

348 lines
12 KiB
Rust

use std::{
io::Cursor,
sync::mpsc::{self, Receiver, TryRecvError},
thread,
};
use image::{ImageFormat, ImageReader, Limits};
use ratatui::{
Frame,
layout::{Alignment, Rect, Size},
widgets::{Paragraph, Wrap},
};
use ratatui_image::{
FilterType, Resize, StatefulImage,
picker::Picker,
thread::{ResizeRequest, ResizeResponse, ThreadProtocol},
};
use crate::youtube::ThumbnailJob;
struct ThumbnailImage {
id: String,
protocol: ThreadProtocol,
replies: Receiver<Result<ResizeResponse, String>>,
complete: bool,
}
pub struct Thumbnail {
picker: Picker,
selected: Option<String>,
fetch: Option<ThumbnailJob>,
decoder: Option<Receiver<Result<ThumbnailImage, String>>>,
image: Option<ThumbnailImage>,
// ponytail: one completed thumbnail plus an in-flight replacement, not a video cache.
cached: Option<ThumbnailImage>,
error: Option<String>,
}
impl Thumbnail {
pub fn new() -> Self {
Self {
picker: Picker::halfblocks(),
selected: None,
fetch: None,
decoder: None,
image: None,
cached: None,
error: None,
}
}
/// Call once in raw mode/alternate screen, before the event loop reads stdin.
pub fn query_terminal(&mut self) {
self.picker = Picker::from_query_stdio().unwrap_or_else(|_| Picker::halfblocks());
let font = self.picker.font_size();
if font.width == 0 || font.height == 0 {
self.picker = Picker::halfblocks();
}
}
pub fn select(&mut self, id: &str) {
if self.selected.as_deref() == Some(id) {
return;
}
self.cancel();
self.selected = Some(id.to_owned());
if self.cached.as_ref().is_some_and(|image| image.id == id) {
self.image = self.cached.take();
return;
}
match ThumbnailJob::start(id) {
Ok(job) => self.fetch = Some(job),
Err(error) => self.error = Some(format!("Could not load thumbnail: {error}")),
}
}
pub fn cancel(&mut self) {
self.fetch = None;
// Each decode has its own receiver: dropping it invalidates even a queued completion.
self.decoder = None;
self.selected = None;
self.error = None;
if let Some(image) = self.image.take()
&& image.complete
{
self.cached = Some(image);
}
}
pub fn is_loading(&self) -> bool {
self.selected.is_some()
&& self.error.is_none()
&& !self
.image
.as_ref()
.is_some_and(|image| image.complete && image.protocol.protocol_type().is_some())
}
pub fn draw(&mut self, frame: &mut Frame, area: Rect) {
self.poll();
if area.width == 0 || area.height == 0 || self.selected.is_none() {
return;
}
let mut message = self.error.as_deref().unwrap_or("Loading thumbnail...");
if let Some(image) = self.image.as_mut() {
let resize = Resize::Scale(Some(FilterType::Triangle));
if let Some(size) = image.protocol.size_for(resize.clone(), area.into()) {
frame.render_stateful_widget(
StatefulImage::new().resize(resize),
right_aligned(area, size),
&mut image.protocol,
);
if image.protocol.protocol_type().is_some() {
return;
}
}
message = "Preparing thumbnail...";
} else if self.decoder.is_some() {
message = "Decoding thumbnail...";
}
frame.render_widget(
Paragraph::new(message)
.alignment(Alignment::Center)
.wrap(Wrap { trim: true }),
Rect::new(
area.x,
area.y + area.height / 2,
area.width,
area.height - area.height / 2,
),
);
}
fn poll(&mut self) {
if let Some(result) = self.fetch.as_mut().and_then(ThumbnailJob::poll) {
self.fetch = None;
match result {
Ok(bytes) => {
let picker = self.picker.clone();
let id = self.selected.clone().expect("A fetch has a selected video");
let (tx, rx) = mpsc::channel();
match thread::Builder::new()
.name("thumbnail-decode".into())
.spawn(move || {
let _ = tx.send(prepare(id, bytes, picker));
}) {
Ok(_) => self.decoder = Some(rx),
Err(error) => {
self.error = Some(format!("Could not decode thumbnail: {error}"));
}
}
}
Err(error) => self.error = Some(format!("Could not load thumbnail: {error}")),
}
}
if let Some(result) = self.decoder.as_ref().map(Receiver::try_recv) {
match result {
Ok(result) => {
self.decoder = None;
match result {
Ok(image) => self.image = Some(image),
Err(error) => self.error = Some(error),
}
}
Err(TryRecvError::Disconnected) => {
self.decoder = None;
self.error = Some("Thumbnail decoder stopped before completing.".into());
}
Err(TryRecvError::Empty) => {}
}
}
if let Some(image) = self.image.as_mut() {
match image.replies.try_recv() {
Ok(Ok(response)) => {
// ThreadProtocol rejects obsolete requests; the widget checks the current
// area again before rendering, so rapid resizes cannot draw an old size.
if image.protocol.update_resized_protocol(response) {
image.complete = true;
self.cached = None;
}
}
Ok(Err(error)) => {
self.image = None;
self.error = Some(format!("Could not render thumbnail: {error}"));
}
Err(TryRecvError::Disconnected) => {
self.image = None;
self.error = Some("Thumbnail renderer stopped before completing.".into());
}
Err(TryRecvError::Empty) => {}
}
}
}
}
fn prepare(id: String, bytes: Vec<u8>, picker: Picker) -> Result<ThumbnailImage, String> {
let image = decode(bytes)?;
let protocol = picker.new_resize_protocol(image);
let (requests_tx, requests_rx) = mpsc::channel::<ResizeRequest>();
let (replies_tx, replies_rx) = mpsc::channel();
thread::Builder::new()
.name("thumbnail-resize".into())
.spawn(move || {
for request in requests_rx {
if replies_tx
.send(request.resize_encode().map_err(|error| error.to_string()))
.is_err()
{
break;
}
}
})
.map_err(|error| format!("Could not prepare thumbnail renderer: {error}"))?;
Ok(ThumbnailImage {
id,
protocol: ThreadProtocol::new(requests_tx, Some(protocol)),
replies: replies_rx,
complete: false,
})
}
fn decode(bytes: Vec<u8>) -> Result<image::DynamicImage, String> {
if bytes.len() > 2 * 1024 * 1024 {
return Err("Thumbnail exceeds the 2 MiB download limit.".into());
}
let mut reader = ImageReader::with_format(Cursor::new(bytes), ImageFormat::Jpeg);
let mut limits = Limits::default();
limits.max_image_width = Some(4096);
limits.max_image_height = Some(4096);
limits.max_alloc = Some(16 * 1024 * 1024);
reader.limits(limits);
reader
.decode()
.map_err(|error| format!("Could not decode thumbnail JPEG: {error}"))
}
fn right_aligned(area: Rect, size: Size) -> Rect {
let width = size.width.min(area.width);
let height = size.height.min(area.height);
Rect::new(
area.x + area.width - width,
area.y + (area.height - height) / 2,
width,
height,
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rejects_corrupt_and_oversized_jpeg_input() {
assert!(decode(b"not a JPEG".to_vec()).is_err());
assert!(decode(vec![0; 2 * 1024 * 1024 + 1]).is_err());
}
#[test]
fn cancellation_discards_queued_decode_completion() {
let mut thumbnail = Thumbnail::new();
let (tx, rx) = mpsc::channel();
let (requests_tx, _requests_rx) = mpsc::channel();
let (_replies_tx, replies_rx) = mpsc::channel();
let protocol = thumbnail
.picker
.new_resize_protocol(image::DynamicImage::new_rgb8(1, 1));
thumbnail.selected = Some("old-video".into());
thumbnail.decoder = Some(rx);
assert!(
tx.send(Ok(ThumbnailImage {
id: "old-video".into(),
protocol: ThreadProtocol::new(requests_tx, Some(protocol)),
replies: replies_rx,
complete: false,
}))
.is_ok()
);
thumbnail.cancel();
thumbnail.selected = Some("new-video".into());
thumbnail.poll();
assert!(thumbnail.image.is_none());
assert!(thumbnail.error.is_none());
}
#[test]
fn small_thumbnail_fills_popup_proportionally_and_resizes() {
use ratatui::{Terminal, backend::TestBackend, style::Color};
use std::time::Duration;
let mut thumbnail = Thumbnail::new();
let font = thumbnail.picker.font_size();
let image = image::DynamicImage::ImageRgb8(image::RgbImage::from_pixel(
48,
36,
image::Rgb([240, 40, 70]),
));
let (requests_tx, requests_rx) = mpsc::channel::<ResizeRequest>();
let (replies_tx, replies_rx) = mpsc::channel();
thumbnail.selected = Some("thumbnail".into());
thumbnail.image = Some(ThumbnailImage {
id: "thumbnail".into(),
protocol: ThreadProtocol::new(
requests_tx,
Some(thumbnail.picker.new_resize_protocol(image)),
),
replies: replies_rx,
complete: false,
});
let mut terminal = Terminal::new(TestBackend::new(100, 40)).unwrap();
for area in [Rect::new(3, 2, 80, 18), Rect::new(7, 4, 30, 8)] {
terminal.draw(|frame| thumbnail.draw(frame, area)).unwrap();
assert!(thumbnail.is_loading());
let response = requests_rx
.recv_timeout(Duration::from_secs(2))
.unwrap()
.resize_encode()
.unwrap();
replies_tx.send(Ok(response)).unwrap();
terminal.draw(|frame| thumbnail.draw(frame, area)).unwrap();
assert!(!thumbnail.is_loading());
let buffer = terminal.backend().buffer();
let pixels: Vec<_> = (0..40)
.flat_map(|y| (0..100).map(move |x| (x, y)))
.filter(|&(x, y)| {
let cell = &buffer[(x, y)];
matches!(cell.fg, Color::Rgb(..)) || matches!(cell.bg, Color::Rgb(..))
})
.collect();
let left = pixels.iter().map(|&(x, _)| x).min().unwrap();
let right = pixels.iter().map(|&(x, _)| x).max().unwrap() + 1;
let top = pixels.iter().map(|&(_, y)| y).min().unwrap();
let bottom = pixels.iter().map(|&(_, y)| y).max().unwrap() + 1;
assert!(left >= area.x && right <= area.right());
assert_eq!(right, area.right()); // Keep the image against the metadata-facing edge.
assert!(top >= area.y && bottom <= area.bottom());
assert!(right - left == area.width || bottom - top == area.height);
assert!((top - area.y).abs_diff(area.bottom() - bottom) <= 1);
let width = u32::from(right - left) * u32::from(font.width);
let height = u32::from(bottom - top) * u32::from(font.height);
assert!(
width.abs_diff(height * 4 / 3)
<= u32::from(font.width) + u32::from(font.height) * 4 / 3
);
}
}
}