Skip to main content

tess/
image_protocol.rs

1//! Pure terminal-graphics encoders: Kitty graphics protocol and Sixel.
2//! Decodes nothing and touches no terminal; callers pass a decoded `RgbaImage`.
3//! Mirrors `render`/`image_render` discipline: plain inputs, byte outputs,
4//! exhaustively unit-tested.
5
6use crate::render::{rgb_to_256, color_256_to_rgb};
7use image::RgbaImage;
8use std::collections::BTreeMap;
9
10const B64: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
11
12/// Standard base64 (RFC 4648) with `=` padding.
13fn base64_encode(bytes: &[u8]) -> String {
14    let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4);
15    for chunk in bytes.chunks(3) {
16        let b0 = chunk[0] as u32;
17        let b1 = *chunk.get(1).unwrap_or(&0) as u32;
18        let b2 = *chunk.get(2).unwrap_or(&0) as u32;
19        let n = (b0 << 16) | (b1 << 8) | b2;
20        out.push(B64[(n >> 18) as usize & 63] as char);
21        out.push(B64[(n >> 12) as usize & 63] as char);
22        out.push(if chunk.len() > 1 { B64[(n >> 6) as usize & 63] as char } else { '=' });
23        out.push(if chunk.len() > 2 { B64[n as usize & 63] as char } else { '=' });
24    }
25    out
26}
27
28/// Encode an image as a Kitty graphics-protocol "transmit and display" command.
29/// Format f=32 (RGBA), chunked at 4096 base64 chars with the m=1/m=0 convention.
30pub fn encode_kitty(img: &RgbaImage) -> Vec<u8> {
31    let (w, h) = img.dimensions();
32    let payload = base64_encode(img.as_raw());
33    let bytes = payload.as_bytes();
34    let mut out: Vec<u8> = Vec::new();
35    let chunks: Vec<&[u8]> = bytes.chunks(4096).collect();
36    let n = chunks.len().max(1);
37    for (i, chunk) in chunks.iter().enumerate() {
38        let more = if i + 1 < n { 1 } else { 0 };
39        out.extend_from_slice(b"\x1b_G");
40        if i == 0 {
41            out.extend_from_slice(format!("a=T,f=32,s={w},v={h},m={more}").as_bytes());
42        } else {
43            out.extend_from_slice(format!("m={more}").as_bytes());
44        }
45        out.push(b';');
46        out.extend_from_slice(chunk);
47        out.extend_from_slice(b"\x1b\\");
48    }
49    out
50}
51
52/// Encode an image as a Sixel data stream. Quantizes each pixel to the xterm
53/// 256-color palette via `rgb_to_256` (shared with `--truecolor` downsampling),
54/// then emits palette definitions and run-length-compressed sixel bands.
55pub fn encode_sixel(img: &RgbaImage) -> Vec<u8> {
56    let (w, h) = img.dimensions();
57    if w == 0 || h == 0 { return b"\x1bPq\x1b\\".to_vec(); }
58
59    let quant: Vec<u8> = img.pixels().map(|p| rgb_to_256(p.0[0], p.0[1], p.0[2])).collect();
60
61    let mut pal: BTreeMap<u8, usize> = BTreeMap::new();
62    for &idx in &quant {
63        let next = pal.len();
64        pal.entry(idx).or_insert(next);
65    }
66
67    let mut out: Vec<u8> = Vec::new();
68    out.extend_from_slice(b"\x1bPq"); // DCS sixel intro: ESC P <params> q
69    for (&idx, &p) in &pal {
70        let (r, g, b) = color_256_to_rgb(idx);
71        let to100 = |v: u8| v as u32 * 100 / 255;
72        out.extend_from_slice(
73            format!("#{};2;{};{};{}", p, to100(r), to100(g), to100(b)).as_bytes());
74    }
75
76    let bands = (h as usize).div_ceil(6);
77    for band in 0..bands {
78        let y0 = (band * 6) as u32;
79        let mut first_color = true;
80        for (&idx, &p) in &pal {
81            let mut values: Vec<u8> = Vec::with_capacity(w as usize);
82            let mut any = false;
83            for x in 0..w {
84                let mut v = 0u8;
85                for r in 0..6u32 {
86                    let y = y0 + r;
87                    if y < h && quant[(y * w + x) as usize] == idx {
88                        v |= 1 << r;
89                    }
90                }
91                if v != 0 { any = true; }
92                values.push(v);
93            }
94            if !any { continue; }
95            // carriage-return to band start so the next color overlays the same 6 rows
96            if !first_color { out.push(b'$'); }
97            first_color = false;
98            out.extend_from_slice(format!("#{p}").as_bytes());
99            let mut x = 0usize;
100            while x < values.len() {
101                let v = values[x];
102                let mut run = 1usize;
103                while x + run < values.len() && values[x + run] == v { run += 1; }
104                // sixel byte = 0x3F + 6-bit column mask, keeping it printable ASCII
105                let ch = (0x3F + v) as char;
106                if run >= 3 {
107                    out.extend_from_slice(format!("!{run}{ch}").as_bytes());
108                } else {
109                    for _ in 0..run { out.push(ch as u8); }
110                }
111                x += run;
112            }
113        }
114        // graphics newline: advance to the next 6-row band
115        if band + 1 < bands { out.push(b'-'); }
116    }
117    out.extend_from_slice(b"\x1b\\"); // ST: end the DCS string
118    out
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use image::{Rgba, RgbaImage};
125
126    #[test]
127    fn base64_known_vectors() {
128        assert_eq!(base64_encode(b""), "");
129        assert_eq!(base64_encode(b"f"), "Zg==");
130        assert_eq!(base64_encode(b"fo"), "Zm8=");
131        assert_eq!(base64_encode(b"foo"), "Zm9v");
132        assert_eq!(base64_encode(b"foob"), "Zm9vYg==");
133        assert_eq!(base64_encode(b"hello, world"), "aGVsbG8sIHdvcmxk");
134    }
135
136    #[test]
137    fn kitty_has_header_keys_and_terminator() {
138        let img = RgbaImage::from_pixel(2, 1, Rgba([1, 2, 3, 255]));
139        let out = encode_kitty(&img);
140        let s = String::from_utf8_lossy(&out);
141        assert!(s.starts_with("\x1b_G"));
142        assert!(s.contains("a=T,f=32,s=2,v=1,m=0"));
143        assert!(s.ends_with("\x1b\\"));
144        let expected_payload = base64_encode(img.as_raw());
145        assert!(s.contains(&expected_payload));
146    }
147
148    #[test]
149    fn kitty_chunks_large_payload_with_more_flags() {
150        let img = RgbaImage::from_pixel(1000, 1, Rgba([9, 9, 9, 9]));
151        let out = encode_kitty(&img);
152        let s = String::from_utf8_lossy(&out);
153        assert!(s.matches("\x1b_G").count() >= 2, "should split into multiple APCs");
154        assert!(s.contains("m=1"), "non-final chunks set m=1");
155        assert!(s.contains("m=0"), "final chunk sets m=0");
156    }
157
158    #[test]
159    fn sixel_has_intro_palette_and_terminator() {
160        let img = RgbaImage::from_pixel(2, 6, Rgba([255, 0, 0, 255]));
161        let out = encode_sixel(&img);
162        let s = String::from_utf8_lossy(&out);
163        assert!(s.starts_with("\x1bPq"), "DCS sixel intro");
164        assert!(s.ends_with("\x1b\\"), "ST terminator");
165        assert!(s.contains(";2;"), "RGB palette definition present");
166    }
167
168    #[test]
169    fn sixel_full_height_uses_all_six_bits() {
170        let img = RgbaImage::from_pixel(1, 6, Rgba([255, 255, 255, 255]));
171        let out = encode_sixel(&img);
172        let s = String::from_utf8_lossy(&out);
173        assert!(s.contains('~'), "all six rows set → '~' sixel char (0x3F+63)");
174    }
175
176    #[test]
177    fn sixel_emits_band_separator_for_tall_images() {
178        let img = RgbaImage::from_pixel(1, 12, Rgba([10, 20, 30, 255]));
179        let out = encode_sixel(&img);
180        let s = String::from_utf8_lossy(&out);
181        assert!(s.contains('-'), "two bands (12px = 2×6) separated by '-'");
182    }
183
184    #[test]
185    fn sixel_emits_rle_for_long_runs() {
186        // A solid 5px-wide, 6px-tall image → one color, a run of 5 identical sixel
187        // chars → RLE `!5<char>` (run >= 3 triggers compression).
188        let img = RgbaImage::from_pixel(5, 6, Rgba([255, 255, 255, 255]));
189        let out = encode_sixel(&img);
190        let s = String::from_utf8_lossy(&out);
191        assert!(s.contains("!5"), "a 5-wide solid run should emit RLE !5, got: {s:?}");
192    }
193
194    #[test]
195    fn sixel_overlays_multiple_colors_in_a_band_with_dollar() {
196        // Two distinct colors in the same 6px-tall band (side by side) must be
197        // emitted as two color passes separated by `$` (carriage-return to band start).
198        let mut img = RgbaImage::new(2, 6);
199        for y in 0..6 { img.put_pixel(0, y, Rgba([255, 0, 0, 255])); }   // col 0 red
200        for y in 0..6 { img.put_pixel(1, y, Rgba([0, 0, 255, 255])); }   // col 1 blue
201        let out = encode_sixel(&img);
202        let s = String::from_utf8_lossy(&out);
203        assert!(s.contains('$'), "two colors in one band must be overlaid with `$`, got: {s:?}");
204        // Two distinct palette entries defined.
205        assert!(s.matches(";2;").count() >= 2, "two palette colors defined");
206    }
207}