From 3e929eabe000fd8aa4dfa3fe1088553df3643ff8 Mon Sep 17 00:00:00 2001 From: Victor Giers Date: Thu, 7 May 2026 11:05:27 +0200 Subject: [PATCH] Add functionality to read prompt metadata from PNG files and include it in map listings --- src-tauri/src/main.rs | 89 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 87 insertions(+), 2 deletions(-) diff --git a/src-tauri/src/main.rs b/src-tauri/src/main.rs index 6eb44c4..1c7b3c6 100644 --- a/src-tauri/src/main.rs +++ b/src-tauri/src/main.rs @@ -2,8 +2,8 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; -use std::fs; -use std::io::{BufRead, BufReader, Read}; +use std::fs::{self, File}; +use std::io::{BufRead, BufReader, Read, Seek, SeekFrom}; use std::path::{Path, PathBuf}; use std::process::{Command, Stdio}; use std::sync::atomic::{AtomicBool, Ordering}; @@ -22,6 +22,7 @@ struct MapInfo { path: String, filename: String, modified: i64, + prompt: Option, } #[derive(Serialize)] @@ -78,6 +79,89 @@ fn python_binary(root: &Path) -> PathBuf { PathBuf::from("python3") } +fn decode_latin1(bytes: &[u8]) -> String { + bytes.iter().map(|&byte| byte as char).collect() +} + +fn prompt_from_png_text_chunk(chunk_type: &[u8; 4], data: &[u8]) -> Option { + match chunk_type { + b"tEXt" => { + let separator = data.iter().position(|&byte| byte == 0)?; + let key = decode_latin1(&data[..separator]); + if !key.eq_ignore_ascii_case("prompt") { + return None; + } + Some(decode_latin1(&data[separator + 1..])) + } + b"iTXt" => { + let keyword_end = data.iter().position(|&byte| byte == 0)?; + let key = decode_latin1(&data[..keyword_end]); + if !key.eq_ignore_ascii_case("prompt") { + return None; + } + + let compression_flag_index = keyword_end + 1; + let compression_method_index = compression_flag_index + 1; + let after_compression = compression_method_index + 1; + if data.len() < after_compression || data[compression_flag_index] != 0 { + return None; + } + + let rest = &data[after_compression..]; + let language_end = rest.iter().position(|&byte| byte == 0)?; + let translated_keyword_start = language_end + 1; + let translated_keyword_end = rest[translated_keyword_start..] + .iter() + .position(|&byte| byte == 0)? + + translated_keyword_start; + let text_start = translated_keyword_end + 1; + let text = &rest[text_start..]; + match String::from_utf8(text.to_vec()) { + Ok(value) => Some(value), + Err(_) => Some(String::from_utf8_lossy(text).to_string()), + } + } + _ => None, + } +} + +fn read_png_prompt(path: &Path) -> Option { + let mut file = File::open(path).ok()?; + let mut signature = [0_u8; 8]; + file.read_exact(&mut signature).ok()?; + if signature != [137, 80, 78, 71, 13, 10, 26, 10] { + return None; + } + + loop { + let mut length_buf = [0_u8; 4]; + if file.read_exact(&mut length_buf).is_err() { + break; + } + + let length = u32::from_be_bytes(length_buf) as usize; + let mut chunk_type = [0_u8; 4]; + file.read_exact(&mut chunk_type).ok()?; + + if &chunk_type == b"tEXt" || &chunk_type == b"iTXt" { + let mut data = vec![0_u8; length]; + file.read_exact(&mut data).ok()?; + if let Some(prompt) = prompt_from_png_text_chunk(&chunk_type, &data) { + return Some(prompt); + } + file.seek(SeekFrom::Current(4)).ok()?; + } else { + file.seek(SeekFrom::Current(length as i64 + 4)).ok()?; + } + + if &chunk_type == b"IEND" { + break; + } + } + + None +} + #[tauri::command] fn list_maps(state: State) -> Result, String> { let dir = output_dir(&state.project_root)?; @@ -116,6 +200,7 @@ fn list_maps(state: State) -> Result, String> { path: path.to_string_lossy().to_string(), filename, modified, + prompt: read_png_prompt(&path), }); } maps.sort_by(|a, b| b.modified.cmp(&a.modified));