Separate image transmission and rendering into two steps

src/enums/icon.rs

@@ -1,10 +1,27 @@
-use crate::gfx::{compute_hash, get_rgba, render_image};
+use crate::gfx::{compute_hash, get_rgba, render_image, send_image};
 use crate::PLS;
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::{LazyLock, Mutex};
 
-static IMAGE_COUNT_MAP: LazyLock<Mutex<HashMap<u32, u8>>> =
+struct ImageData {
+	/// the ID assigned by the terminal to our image
+	///
+	/// This is different from the hash of the image data. Allowing the
+	/// terminal to choose an ID prevents new invocations of `pls` from
+	/// overwriting IDs of images that were displayed by previous
+	/// invocations.
+	id: u32,
+	/// the number of times the image has been displayed
+	///
+	/// This generates new placement IDs for the images. This is
+	/// required specifically because WezTerm has a bug where not
+	/// setting unique placement IDs overwrites placements instead of
+	/// creating new ones.
+	count: u8,
+}
+
+static IMAGE_DATA: LazyLock<Mutex<HashMap<u32, ImageData>>> =
 	LazyLock::new(|| Mutex::new(HashMap::new()));
 
 /// This enum contains the two formats of icons supported by `pls`.
@@ -79,18 +96,26 @@ impl Icon {
 				};
 
 				let size = Icon::size();
-				let id = compute_hash(&PathBuf::from(path.as_ref()), size);
+				let hash = compute_hash(&PathBuf::from(path.as_ref()), size);
 
-				let mut image_ids = IMAGE_COUNT_MAP.lock().unwrap();
-				let count = image_ids.entry(id).or_insert(0);
+				let mut image_data_store = IMAGE_DATA.lock().unwrap();
+				let data = image_data_store
+					.entry(hash)
+					.or_insert_with(|| ImageData { count: 0, id: 0 });
 
-				*count += 1;
-				let rgba_data = if *count == 1 {
-					get_rgba(id, &PathBuf::from(path.as_ref()), size)
-				} else {
-					None
-				};
-				return render_image(id, *count, size, rgba_data.as_deref());
+				data.count += 1;
+				if data.count == 1 {
+					// If the image is appearing for the first time in
+					// this session, we send it to the terminal and get
+					// an ID assigned to it.
+					match get_rgba(hash, &PathBuf::from(path.as_ref()), size) {
+						Some(rgba_data) => {
+							data.id = send_image(hash, size, &rgba_data).unwrap();
+						}
+						None => return default,
+					}
+				}
+				render_image(data.id, size, data.count)
 			}
 		}
 	}

src/exc.rs

@@ -1,6 +1,7 @@
 use crate::fmt::render;
 use std::fmt::{Display, Formatter, Result as FmtResult};
 
+#[derive(Debug)]
 pub enum Exc {
 	/// wraps all occurrences of errors in I/O operations
 	Io(std::io::Error),

src/gfx.rs

@@ -9,6 +9,7 @@
 //! * [`compute_hash`]
 //! * [`is_supported`]
 //! * [`render_image`]
+//! * [`send_image`]
 //! * [`strip_image`]
 //! * [`get_rgba`]
 
@@ -17,5 +18,5 @@ mod kitty;
 mod svg;
 
 pub use hash::compute_hash;
-pub use kitty::{is_supported, render_image, strip_image};
+pub use kitty::{is_supported, render_image, send_image, strip_image};
 pub use svg::get_rgba;

src/gfx/kitty.rs

@@ -6,9 +6,11 @@ use log::debug;
 use regex::Regex;
 use std::sync::LazyLock;
 
-static KITTY_IMAGE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\x1b_G.*?\x1b\\").unwrap());
 const CHUNK_SIZE: usize = 4096;
 
+static KITTY_IMAGE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\x1b_G.*?\x1b\\").unwrap());
+static IMAGE_ID: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"i=(?P<id>\d+)").unwrap());
+
 /// Check if the terminal supports Kitty's terminal graphics protocol.
 ///
 /// We make a Kitty request and see if the terminal responds with an OK
@@ -29,76 +31,75 @@ pub fn is_supported() -> bool {
 	false
 }
 
-/// Send the RGBA data to the terminal for immediate rendering.
+/// Send the RGBA data to the terminal and get an ID for the image.
 ///
-/// To achieve this, we must send the following control sequence:
+/// The image is sent in chunks of 4096 bytes. The last chunk has the
+/// `m` parameter set to 0. The terminal then assigns our image an ID,
+/// instead of us determining one.
 ///
-/// * f = 32 signals that data will be 32-bit RGBA
-/// * t = d signals that data will be within the control sequence
-/// * a = T signals that image should be immediately displayed
-/// * C = 1 signals that cursor should not be moved
-/// * m = 1 if more data follows, 0 if this is the last chunk
-/// * s = size for width (pixels)
-/// * v = size for height (pixels)
+/// In this stage, we do not show the image (it will be shown in a later
+/// step) so placement controls are not required.
 ///
-/// The image is sent in chunks of 4096 bytes. The last chunk has the
-/// `m` parameter set to 0.
+/// # Arguments
+///
+/// * `hash` - the hash of the image data
+/// * `size` - the size of the image, in pixels
+/// * `rgba_data` - the RGBA data to send
+pub fn send_image(hash: u32, size: u8, rgba_data: &[u8]) -> Result<u32, Exc> {
+	let mut query = String::new();
+
+	let encoded = BASE64_STANDARD.encode(rgba_data);
+	let mut iter = encoded.chars().peekable();
+
+	let first_chunk: String = iter.by_ref().take(CHUNK_SIZE).collect();
+	query.push_str(&format!(
+		"\x1b_G\
+		a=t,I={hash},s={size},v={size},t=d,f=32,m=1;\
+		{first_chunk}\
+		\x1b\\"
+	));
+
+	while iter.peek().is_some() {
+		let chunk: String = iter.by_ref().take(CHUNK_SIZE).collect();
+		query.push_str(&format!("\x1b_Gm=1;{chunk}\x1b\\"));
+	}
+
+	query.push_str("\x1b_Gm=0;\x1b\\");
+
+	let res = query_raw(&query, 200)?;
+	IMAGE_ID
+		.captures(&res)
+		.map(|cap| cap["id"].parse().unwrap())
+		.ok_or(Exc::Other(String::from("Could not extract image ID.")))
+}
+
+/// Render the image with the given ID to the screen.
+///
+/// In this stage, we do not transmit the image (it has already been
+/// done) so transmission controls are not required.
 ///
 /// The image is rendered in a way that the cursor does not move. Then
 /// we move the cursor by as many cells as the icon width (and a space).
 ///
 /// # Arguments
 ///
 /// * `id` - the unique ID of the image
-/// * `count` - the number of times this image has appeared so far
 /// * `size` - the size of the image, in pixels
-/// * `rgba_data` - the RGBA data to render
-pub fn render_image(id: u32, count: u8, size: u8, rgba_data: Option<&[u8]>) -> String {
+/// * `count` - the number of times this image has appeared so far
+pub fn render_image(id: u32, size: u8, count: u8) -> String {
 	let cell_height = PLS.window.as_ref().unwrap().cell_height();
 	let off_y = if cell_height > size {
 		(cell_height - size) / 2
 	} else {
 		0
 	};
 
-	let mut output = String::new();
-
-	// If data is provided, the image is new, so we transmit it with the
-	// control `a=t`.
-	if let Some(rgba_data) = rgba_data {
-		let encoded = BASE64_STANDARD.encode(rgba_data);
-		let mut iter = encoded.chars().peekable();
-
-		let first_chunk: String = iter.by_ref().take(CHUNK_SIZE).collect();
-
-		// TODO: By sending fresh data for existing IDs, the images
-		// shown in previous usages of `pls` disappear.
-		output.push_str(&format!(
-			"\x1b_G\
-			f=32,t=d,a=t,m=1,q=2,i={id},s={size},v={size},Y={off_y};\
-			{first_chunk}\
-			\x1b\\"
-		));
-
-		while iter.peek().is_some() {
-			let chunk: String = iter.by_ref().take(CHUNK_SIZE).collect();
-			output.push_str(&format!("\x1b_Gm=1;{chunk}\x1b\\"));
-		}
-
-		output.push_str("\x1b_Gm=0,q=2;\x1b\\");
-	}
-
-	// Once the data is sent, we render the previously transmitted image
-	// with the control `a=p`.
-	output.push_str(&format!(
+	format!(
 		"\x1b_G\
-		a=p,C=1,q=2,i={id},p={count},s={size},v={size},Y={off_y};\
-		\x1b\\"
-	));
-
-	output.push_str("\x1b[2C");
-
-	output
+		a=p,i={id},s={size},v={size},p={count},C=1,Y={off_y},q=2;\
+		\x1b\\\
+		\x1b[2C"
+	)
 }
 
 /// Strip the image data from the text.