1use std::io::Write;
8
9use stet_graphics::device::{PageSink, PageSinkFactory};
10
11pub struct PngSinkFactory;
13
14impl PageSinkFactory for PngSinkFactory {
15 fn create_sink(&self, output_path: &str) -> Result<Box<dyn PageSink>, String> {
16 Ok(Box::new(PngSink {
17 output_path: output_path.to_string(),
18 writer: None,
19 }))
20 }
21}
22
23struct PngSink {
25 output_path: String,
26 writer: Option<PngStreamWriter>,
27}
28
29struct PngStreamWriter {
31 stream: png::StreamWriter<'static, std::io::BufWriter<std::fs::File>>,
32 _writer_box: *mut png::Writer<std::io::BufWriter<std::fs::File>>,
36}
37
38unsafe impl Send for PngStreamWriter {}
40
41impl PngStreamWriter {
42 fn new(output_path: &str, width: u32, height: u32, band_h: u32) -> Result<Self, String> {
43 let file = std::fs::File::create(output_path)
44 .map_err(|e| format!("Failed to create PNG '{}': {}", output_path, e))?;
45 let buf_writer = std::io::BufWriter::with_capacity(256 * 1024, file);
46
47 let mut encoder = png::Encoder::new(buf_writer, width, height);
48 encoder.set_color(png::ColorType::Rgba);
49 encoder.set_depth(png::BitDepth::Eight);
50 encoder.set_compression(png::Compression::Default);
51 encoder.set_adaptive_filter(png::AdaptiveFilterType::Adaptive);
52
53 let writer = encoder
54 .write_header()
55 .map_err(|e| format!("Failed to write PNG header '{}': {}", output_path, e))?;
56
57 let writer_box = Box::into_raw(Box::new(writer));
59 let stream = unsafe { &mut *writer_box }
60 .stream_writer_with_size(band_h as usize * width as usize * 4)
61 .map_err(|e| format!("Failed to create PNG stream '{}': {}", output_path, e))?;
62
63 Ok(Self {
64 stream,
65 _writer_box: writer_box,
66 })
67 }
68
69 fn write_all(&mut self, data: &[u8]) -> Result<(), String> {
70 self.stream
71 .write_all(data)
72 .map_err(|e| format!("Failed to write PNG data: {}", e))
73 }
74
75 fn finish(self) -> Result<(), String> {
76 self.stream
77 .finish()
78 .map_err(|e| format!("Failed to finish PNG: {}", e))?;
79 unsafe {
81 let _ = Box::from_raw(self._writer_box);
82 }
83 Ok(())
84 }
85}
86
87impl PageSink for PngSink {
88 fn begin_page(&mut self, width: u32, height: u32) -> Result<(), String> {
89 let writer = PngStreamWriter::new(&self.output_path, width, height, height)?;
90 self.writer = Some(writer);
91 Ok(())
92 }
93
94 fn write_rows(&mut self, rgba_rows: &[u8], _num_rows: u32) -> Result<(), String> {
95 let writer = self
96 .writer
97 .as_mut()
98 .ok_or("PngSink: begin_page not called")?;
99 writer.write_all(rgba_rows)
100 }
101
102 fn end_page(&mut self) -> Result<(), String> {
103 let writer = self.writer.take().ok_or("PngSink: begin_page not called")?;
104 writer.finish()
105 }
106}