1use super::lzw::LzwEncoder;
2use super::quant::{Dither, Histogram, Lut, Mapper, Palette};
3use super::{Error, Result};
4use crate::media::Rgba;
5use rayon::prelude::*;
6use std::cell::RefCell;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum PaletteMode {
10 Global,
11 PerFrame,
12}
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub struct EncodeOptions {
16 pub colors: usize,
17 pub dither: Dither,
18 pub palette: PaletteMode,
19 pub loop_count: Option<u16>,
20 pub diff: bool,
21}
22
23impl Default for EncodeOptions {
24 fn default() -> Self {
25 EncodeOptions {
26 colors: 256,
27 dither: Dither::FloydSteinberg,
28 palette: PaletteMode::Global,
29 loop_count: Some(0),
30 diff: true,
31 }
32 }
33}
34
35struct Mapped {
36 idx: Vec<u8>,
37 palette: Option<Palette>,
38 alpha: bool,
39 delay: u16,
40}
41
42struct Record {
43 rect: (usize, usize, usize, usize),
44 idx: Vec<u8>,
45 palette: Option<Palette>,
46 transparent: u8,
47 disposal: u8,
48 delay: u16,
49 body: Vec<u8>,
50}
51
52pub struct Encoder {
53 width: usize,
54 height: usize,
55 opts: EncodeOptions,
56 global: Option<(Palette, Lut)>,
57 out: Vec<u8>,
58 started: bool,
59 batch: Vec<(Rgba, u16)>,
60 pending: Option<Mapped>,
61 carry: Option<Rgba>,
62 prev: Option<Rgba>,
63 canvas: Vec<u32>,
64 held: Option<Record>,
65 spare: Vec<Vec<u8>>,
66}
67
68thread_local! {
69 static LZW: RefCell<LzwEncoder> = RefCell::new(LzwEncoder::default());
70 static ERR: RefCell<Vec<i16>> = const { RefCell::new(Vec::new()) };
71 static HIST: RefCell<Option<Histogram>> = const { RefCell::new(None) };
72}
73
74fn keep(slot: &mut Option<Rgba>, from: &Rgba) {
75 match slot {
76 Some(r) => {
77 r.width = from.width;
78 r.height = from.height;
79 r.data.clear();
80 r.data.extend_from_slice(&from.data);
81 }
82 None => *slot = Some(from.clone()),
83 }
84}
85
86fn batch_size() -> usize {
87 rayon::current_num_threads().clamp(8, 32)
88}
89
90fn bits_for(entries: usize) -> u8 {
91 let mut b = 1;
92 while (1usize << b) < entries {
93 b += 1;
94 }
95 b
96}
97
98fn check_size(width: u32, height: u32) -> Result<()> {
99 if !(1..=65535).contains(&width) || !(1..=65535).contains(&height) {
100 return Err(Error::Invalid("GIF sides are 1 to 65535"));
101 }
102 Ok(())
103}
104
105fn check_frame(frame: &Rgba, width: usize, height: usize) -> Result<()> {
106 if frame.width as usize != width || frame.height as usize != height {
107 return Err(Error::Invalid("every frame has the canvas size"));
108 }
109 if frame.data.len() < width * height * 4 {
110 return Err(Error::Invalid("frame data shorter than its size"));
111 }
112 Ok(())
113}
114
115fn packed_colors(p: &Palette) -> [u32; 256] {
116 let mut t = [0u32; 256];
117 for (d, c) in t.iter_mut().zip(p.colors()) {
118 *d = 0xFF00_0000 | c[0] as u32 | (c[1] as u32) << 8 | (c[2] as u32) << 16;
119 }
120 t
121}
122
123impl Encoder {
124 pub fn new(width: u32, height: u32, opts: EncodeOptions) -> Result<Encoder> {
125 check_size(width, height)?;
126 let opts = EncodeOptions {
127 colors: opts.colors.clamp(2, 256),
128 ..opts
129 };
130 Ok(Encoder {
131 width: width as usize,
132 height: height as usize,
133 opts,
134 global: None,
135 out: Vec::new(),
136 started: false,
137 batch: Vec::new(),
138 pending: None,
139 carry: None,
140 prev: None,
141 canvas: vec![0; width as usize * height as usize],
142 held: None,
143 spare: Vec::new(),
144 })
145 }
146
147 pub fn with_palette(
148 width: u32,
149 height: u32,
150 opts: EncodeOptions,
151 palette: Palette,
152 ) -> Result<Encoder> {
153 let mut e = Encoder::new(width, height, opts)?;
154 let palette = Palette::new(palette.colors()[..palette.len().min(255)].to_vec())?;
155 let lut = Lut::new(&palette);
156 e.global = Some((palette, lut));
157 Ok(e)
158 }
159
160 fn max_colors(&self) -> usize {
161 self.opts.colors.min(255)
162 }
163
164 pub fn push(&mut self, frame: Rgba, delay_cs: u16) -> Result<()> {
165 check_frame(&frame, self.width, self.height)?;
166 self.batch.push((frame, delay_cs));
167 if self.batch.len() >= batch_size() {
168 self.flush_batch(false);
169 }
170 Ok(())
171 }
172
173 fn flush_batch(&mut self, last: bool) {
174 let mut batch = std::mem::take(&mut self.batch);
175 let refs: Vec<(&Rgba, u16)> = batch.iter().map(|(f, d)| (f, *d)).collect();
176 self.process(&refs, last);
177 drop(refs);
178 batch.clear();
179 self.batch = batch;
180 }
181
182 pub fn finish(mut self) -> Vec<u8> {
183 self.flush_batch(true);
184 self.close()
185 }
186
187 fn close(mut self) -> Vec<u8> {
188 if let Some(r) = self.held.take() {
189 self.write(&r);
190 }
191 if !self.started {
192 self.header();
193 }
194 self.out.push(0x3B);
195 self.out
196 }
197
198 fn process(&mut self, frames: &[(&Rgba, u16)], last: bool) {
199 if self.opts.palette == PaletteMode::Global && self.global.is_none() && !frames.is_empty() {
200 let refs: Vec<&Rgba> = frames.iter().map(|f| f.0).collect();
201 let palette = Palette::from_frames(&refs, self.max_colors());
202 let lut = Lut::new(&palette);
203 self.global = Some((palette, lut));
204 }
205 let n = self.width * self.height;
206 let mut bufs: Vec<Vec<u8>> = (0..frames.len())
207 .map(|_| self.spare.pop().unwrap_or_default())
208 .collect();
209 for b in bufs.iter_mut() {
210 b.resize(n, 0);
211 }
212 let colors = self.max_colors();
213 let dither = self.opts.dither;
214 let global = self.global.as_ref();
215 let mapped: Vec<Mapped> = frames
216 .par_iter()
217 .zip(bufs.into_par_iter())
218 .map(|(&(src, delay), mut idx)| {
219 let alpha = src.data.chunks_exact(4).any(|p| p[3] < 128);
220 let own;
221 let (palette, lut, local) = match global {
222 Some((p, l)) => (p, l, None),
223 None => {
224 let mut h = HIST.take().unwrap_or_default();
225 h.clear();
226 h.add(src);
227 let p = h.palette(colors);
228 HIST.set(Some(h));
229 let l = Lut::new(&p);
230 own = (p, l);
231 (&own.0, &own.1, Some(own.0.clone()))
232 }
233 };
234 let mapper = Mapper {
235 lut,
236 palette: palette.colors(),
237 dither,
238 transparent: palette.len() as u8,
239 };
240 let mut err = ERR.take();
241 mapper.map(src, &mut idx, &mut err);
242 ERR.set(err);
243 Mapped {
244 idx,
245 palette: local,
246 alpha,
247 delay,
248 }
249 })
250 .collect();
251 let carry = self.carry.take();
252 let had_pending = self.pending.is_some();
253 let mut items: Vec<Mapped> = self.pending.take().into_iter().chain(mapped).collect();
254 let srcs: Vec<&Rgba> = carry
255 .iter()
256 .filter(|_| had_pending)
257 .chain(frames.iter().map(|f| f.0))
258 .collect();
259 let count = if last {
260 items.len()
261 } else {
262 items.len().saturating_sub(1)
263 };
264 let alphas: Vec<bool> = items.iter().map(|m| m.alpha).collect();
265 let rest = items.split_off(count);
266 let mut prev = self.prev.take();
267 let mut records = Vec::with_capacity(count);
268 for (i, f) in items.into_iter().enumerate() {
269 let before = if i == 0 {
270 prev.as_ref()
271 } else {
272 Some(srcs[i - 1])
273 };
274 let next_alpha = alphas.get(i + 1).copied().unwrap_or(false);
275 records.push(self.diff(f, srcs[i], before, next_alpha));
276 }
277 if count > 0 {
278 keep(&mut prev, srcs[count - 1]);
279 }
280 drop(srcs);
281 self.prev = prev;
282 if let Some(m) = rest.into_iter().next() {
283 if count == 0 && had_pending {
284 self.carry = carry;
285 } else {
286 let mut buf = carry;
287 keep(&mut buf, frames[count - had_pending as usize].0);
288 self.carry = buf;
289 }
290 self.pending = Some(m);
291 }
292 let w = self.width;
293 records.par_iter_mut().for_each(|r| encode_body(r, w));
294 for r in records {
295 self.emit(r);
296 }
297 }
298
299 fn diff(&mut self, f: Mapped, src: &Rgba, prev: Option<&Rgba>, next_alpha: bool) -> Record {
300 let palette = f
301 .palette
302 .as_ref()
303 .unwrap_or_else(|| &self.global.as_ref().unwrap().0);
304 let t = palette.len() as u8;
305 let colors = packed_colors(palette);
306 let full = prev.is_none() || f.alpha || next_alpha || !self.opts.diff;
307 let mut idx = f.idx;
308 let w = self.width;
309 let rect = if full {
310 for (c, &i) in self.canvas.iter_mut().zip(&idx) {
311 *c = colors[i as usize];
312 }
313 (0, 0, w, self.height)
314 } else {
315 diff_rows(
316 &mut idx,
317 &mut self.canvas,
318 &src.data,
319 &prev.unwrap().data,
320 &colors,
321 t,
322 w,
323 )
324 };
325 if next_alpha {
326 self.canvas.fill(0);
327 }
328 Record {
329 rect,
330 idx,
331 palette: f.palette,
332 transparent: t,
333 disposal: if next_alpha { 2 } else { 1 },
334 delay: f.delay,
335 body: Vec::new(),
336 }
337 }
338
339 fn emit(&mut self, mut r: Record) {
340 if r.rect.2 == 0 {
341 if let Some(h) = self.held.as_mut()
342 && let Some(d) = h.delay.checked_add(r.delay)
343 {
344 h.delay = d;
345 self.spare.push(r.idx);
346 return;
347 }
348 r.rect = (0, 0, 1, 1);
349 r.idx[0] = r.transparent;
350 encode_body(&mut r, self.width);
351 }
352 if let Some(h) = self.held.take() {
353 self.write(&h);
354 self.spare.push(h.idx);
355 }
356 self.held = Some(r);
357 }
358
359 fn header(&mut self) {
360 self.started = true;
361 let o = &mut self.out;
362 o.extend_from_slice(b"GIF89a");
363 o.extend_from_slice(&(self.width as u16).to_le_bytes());
364 o.extend_from_slice(&(self.height as u16).to_le_bytes());
365 match &self.global {
366 Some((p, _)) => {
367 let bits = bits_for(p.len() + 1);
368 o.push(0x80 | 0x70 | (bits - 1));
369 o.push(p.len() as u8);
370 o.push(0);
371 table(o, p, bits);
372 }
373 None => {
374 o.push(0x70);
375 o.push(0);
376 o.push(0);
377 }
378 }
379 if let Some(n) = self.opts.loop_count {
380 o.extend_from_slice(&[0x21, 0xFF, 11]);
381 o.extend_from_slice(b"NETSCAPE2.0");
382 o.extend_from_slice(&[3, 1]);
383 o.extend_from_slice(&n.to_le_bytes());
384 o.push(0);
385 }
386 }
387
388 fn write(&mut self, r: &Record) {
389 if !self.started {
390 self.header();
391 }
392 let o = &mut self.out;
393 o.extend_from_slice(&[0x21, 0xF9, 4, (r.disposal << 2) | 1]);
394 o.extend_from_slice(&r.delay.to_le_bytes());
395 o.extend_from_slice(&[r.transparent, 0, 0x2C]);
396 let (x, y, w, h) = r.rect;
397 for v in [x, y, w, h] {
398 o.extend_from_slice(&(v as u16).to_le_bytes());
399 }
400 match &r.palette {
401 Some(p) => {
402 let bits = bits_for(p.len() + 1);
403 o.push(0x80 | (bits - 1));
404 table(o, p, bits);
405 }
406 None => o.push(0),
407 }
408 o.extend_from_slice(&r.body);
409 }
410}
411
412fn table(o: &mut Vec<u8>, p: &Palette, bits: u8) {
413 for c in p.colors() {
414 o.extend_from_slice(c);
415 }
416 o.resize(o.len() + ((1usize << bits) - p.len()) * 3, 0);
417}
418
419fn encode_body(r: &mut Record, stride: usize) {
420 let (x, y, w, h) = r.rect;
421 if w == 0 {
422 return;
423 }
424 let entries = r
425 .palette
426 .as_ref()
427 .map_or(0, |p| p.len())
428 .max(r.transparent as usize)
429 + 1;
430 let min = bits_for(entries).max(2);
431 let idx = &r.idx;
432 let rows = (y..y + h).map(|row| &idx[row * stride + x..row * stride + x + w]);
433 let mut body = std::mem::take(&mut r.body);
434 body.clear();
435 LZW.with_borrow_mut(|l| l.encode(min, rows, &mut body));
436 r.body = body;
437}
438
439fn diff_rows(
440 idx: &mut [u8],
441 canvas: &mut [u32],
442 src: &[u8],
443 prev: &[u8],
444 colors: &[u32; 256],
445 t: u8,
446 w: usize,
447) -> (usize, usize, usize, usize) {
448 let band = 16;
449 let boxes: Vec<(usize, usize, usize, usize)> = idx
450 .par_chunks_mut(w * band)
451 .zip(canvas.par_chunks_mut(w * band))
452 .zip(
453 src.par_chunks(w * 4 * band)
454 .zip(prev.par_chunks(w * 4 * band)),
455 )
456 .enumerate()
457 .map(|(i, ((ib, cb), (sb, pb)))| {
458 let (x0, y0, x1, y1) = diff_band(ib, cb, sb, pb, colors, t, w);
459 (x0, y0 + i * band, x1, y1 + i * band)
460 })
461 .collect();
462 let mut r = (usize::MAX, usize::MAX, 0, 0);
463 for b in boxes {
464 if b.2 > b.0 {
465 r = (r.0.min(b.0), r.1.min(b.1), r.2.max(b.2), r.3.max(b.3));
466 }
467 }
468 if r.0 == usize::MAX {
469 (0, 0, 0, 0)
470 } else {
471 (r.0, r.1, r.2 - r.0, r.3 - r.1)
472 }
473}
474
475crate::media::isa::dispatch! {
476 fn diff_band(
477 idx: &mut [u8],
478 canvas: &mut [u32],
479 src: &[u8],
480 prev: &[u8],
481 colors: &[u32; 256],
482 t: u8,
483 w: usize,
484 ) -> (usize, usize, usize, usize) {
485 let mut x0 = usize::MAX;
486 let mut x1 = 0;
487 let mut y0 = 0;
488 let mut y1 = 0;
489 for (y, ((irow, crow), (srow, prow))) in idx
490 .chunks_exact_mut(w)
491 .zip(canvas.chunks_exact_mut(w))
492 .zip(src.chunks_exact(w * 4).zip(prev.chunks_exact(w * 4)))
493 .enumerate()
494 {
495 let mut first = usize::MAX;
496 let mut lastx = 0;
497 for (x, ((i, c), (s, p))) in irow
498 .iter_mut()
499 .zip(crow.iter_mut())
500 .zip(srow.chunks_exact(4).zip(prow.chunks_exact(4)))
501 .enumerate()
502 {
503 let col = colors[*i as usize];
504 if s == p || col == *c {
505 *i = t;
506 } else {
507 *c = col;
508 first = first.min(x);
509 lastx = x;
510 }
511 }
512 if first != usize::MAX {
513 if x0 == usize::MAX {
514 y0 = y;
515 }
516 x0 = x0.min(first);
517 x1 = x1.max(lastx + 1);
518 y1 = y + 1;
519 }
520 }
521 if x0 == usize::MAX {
522 (0, 0, 0, 0)
523 } else {
524 (x0, y0, x1, y1)
525 }
526 }
527}
528
529pub fn encode(frames: &[Rgba], delays_cs: &[u16], opts: EncodeOptions) -> Result<Vec<u8>> {
530 let Some(first) = frames.first() else {
531 return Err(Error::Invalid("no frames to encode"));
532 };
533 check_size(first.width, first.height)?;
534 for f in frames {
535 check_frame(f, first.width as usize, first.height as usize)?;
536 }
537 let mut e = match opts.palette {
538 PaletteMode::Global => {
539 let refs: Vec<&Rgba> = frames.iter().collect();
540 let palette = Palette::from_frames(&refs, opts.colors.clamp(2, 256).min(255));
541 Encoder::with_palette(first.width, first.height, opts, palette)?
542 }
543 PaletteMode::PerFrame => Encoder::new(first.width, first.height, opts)?,
544 };
545 let delays: Vec<u16> = (0..frames.len())
546 .map(|i| delays_cs.get(i).copied().unwrap_or(10))
547 .collect();
548 let size = batch_size();
549 let refs: Vec<(&Rgba, u16)> = frames.iter().zip(delays).collect();
550 let chunks = refs.chunks(size).count();
551 for (i, chunk) in refs.chunks(size).enumerate() {
552 e.process(chunk, i + 1 == chunks);
553 }
554 Ok(e.close())
555}