Skip to main content

pfx_editor_shell/
offscreen.rs

1use egui::{
2    Context, Event, FullOutput, Id, Key, Modifiers, MouseWheelUnit, PointerButton, Pos2, RawInput,
3    Rect, TouchPhase, Ui, Vec2, ViewportId, ViewportInfo,
4};
5use egui_wgpu::{RendererOptions, ScreenDescriptor};
6use pfx_gpu::{Gpu, wgpu};
7
8use crate::paint::Painter;
9use crate::{App, Error, Mods, Script, Step, Theme, block_on, context};
10
11pub const MAX_SIDE: u32 = 8192;
12const DT: f32 = 1.0 / 60.0;
13const DRAG_STEPS: u32 = 4;
14const SETTLE: usize = 2;
15const CLOCK: &str = "pfx-editor-shell/drive-passes";
16
17struct Driver<'a> {
18    ctx: &'a Context,
19    screen: Rect,
20    passes: u64,
21    modifiers: Modifiers,
22    textures: egui::TexturesDelta,
23}
24
25impl Driver<'_> {
26    fn pass(&mut self, events: Vec<Event>, ui: &mut impl FnMut(&mut Ui)) -> FullOutput {
27        let events: Vec<Event> = std::iter::once(Event::ModifiersChanged(self.modifiers))
28            .chain(events)
29            .collect();
30        let mut input = RawInput {
31            screen_rect: Some(self.screen),
32            max_texture_side: Some(MAX_SIDE as usize),
33            time: Some(self.passes as f64 * f64::from(DT)),
34            predicted_dt: DT,
35            focused: true,
36            events,
37            ..RawInput::default()
38        };
39        input.viewports.insert(
40            ViewportId::ROOT,
41            ViewportInfo {
42                native_pixels_per_point: Some(1.0),
43                focused: Some(true),
44                ..ViewportInfo::default()
45            },
46        );
47        self.passes += 1;
48        let mut output = self.ctx.run_ui(input, |root| ui(root));
49        self.textures
50            .append(std::mem::take(&mut output.textures_delta));
51        output
52    }
53}
54
55fn modifiers_of(mods: Mods) -> Modifiers {
56    Modifiers {
57        alt: mods.alt,
58        ctrl: mods.ctrl,
59        shift: mods.shift,
60        mac_cmd: false,
61        command: mods.ctrl,
62    }
63}
64
65fn key(name: &str, pressed: bool, modifiers: Modifiers) -> Option<Event> {
66    Some(Event::Key {
67        key: Key::from_name(name)?,
68        physical_key: None,
69        pressed,
70        repeat: false,
71        modifiers,
72    })
73}
74
75fn button(at: Pos2, pressed: bool) -> Event {
76    press(at, PointerButton::Primary, pressed, Modifiers::NONE)
77}
78
79fn press(at: Pos2, button: PointerButton, pressed: bool, modifiers: Modifiers) -> Event {
80    Event::PointerButton {
81        pos: at,
82        button,
83        pressed,
84        modifiers,
85    }
86}
87
88pub fn drive(
89    ctx: &Context,
90    size: [u32; 2],
91    script: &Script,
92    ui: impl FnMut(&mut Ui),
93) -> FullOutput {
94    let mut output = drive_painted(ctx, size, script, ui);
95    output.textures_delta.clear();
96    output
97}
98
99fn drive_painted(
100    ctx: &Context,
101    size: [u32; 2],
102    script: &Script,
103    mut ui: impl FnMut(&mut Ui),
104) -> FullOutput {
105    let mut driver = Driver {
106        ctx,
107        screen: Rect::from_min_size(Pos2::ZERO, Vec2::new(size[0] as f32, size[1] as f32)),
108        passes: ctx.data(|data| data.get_temp(Id::new(CLOCK))).unwrap_or(0),
109        modifiers: Modifiers::NONE,
110        textures: egui::TexturesDelta::default(),
111    };
112    let mut last = driver.pass(Vec::new(), &mut ui);
113    let mut events = Vec::new();
114    let mut fresh = true;
115    for step in &script.steps {
116        match step {
117            Step::Key(name, mods) if *mods == Mods::default() => {
118                let modifiers = driver.modifiers;
119                events.extend(
120                    [key(name, true, modifiers), key(name, false, modifiers)]
121                        .into_iter()
122                        .flatten(),
123                )
124            }
125            Step::Key(name, mods) => {
126                if !events.is_empty() {
127                    driver.pass(std::mem::take(&mut events), &mut ui);
128                }
129                let held = driver.modifiers;
130                let modifiers = modifiers_of(*mods);
131                driver.modifiers = held | modifiers;
132                let chord = driver.modifiers;
133                events.extend(
134                    [key(name, true, chord), key(name, false, chord)]
135                        .into_iter()
136                        .flatten(),
137                );
138                last = driver.pass(std::mem::take(&mut events), &mut ui);
139                driver.modifiers = held;
140                fresh = false;
141            }
142            Step::Text(text) => events.push(Event::Text(text.clone())),
143            Step::Move([x, y]) => events.push(Event::PointerMoved(Pos2::new(*x, *y))),
144            Step::Click([x, y]) => {
145                let at = Pos2::new(*x, *y);
146                events.push(Event::PointerMoved(at));
147                events.push(button(at, true));
148                last = driver.pass(std::mem::take(&mut events), &mut ui);
149                events.push(button(at, false));
150                fresh = false;
151            }
152            Step::Press { at, mods } => {
153                let at = Pos2::new(at[0], at[1]);
154                let modifiers = modifiers_of(*mods);
155                driver.modifiers = modifiers;
156                events.push(Event::PointerMoved(at));
157                events.push(press(at, PointerButton::Primary, true, modifiers));
158                last = driver.pass(std::mem::take(&mut events), &mut ui);
159                fresh = false;
160            }
161            Step::Release([x, y]) => {
162                let at = Pos2::new(*x, *y);
163                events.push(Event::PointerMoved(at));
164                events.push(press(at, PointerButton::Primary, false, driver.modifiers));
165                last = driver.pass(std::mem::take(&mut events), &mut ui);
166                driver.modifiers = Modifiers::NONE;
167                fresh = false;
168            }
169            Step::Drag {
170                from,
171                to,
172                middle,
173                mods,
174            } => {
175                let (from, to) = (Pos2::new(from[0], from[1]), Pos2::new(to[0], to[1]));
176                let held = if *middle {
177                    PointerButton::Middle
178                } else {
179                    PointerButton::Primary
180                };
181                let modifiers = modifiers_of(*mods);
182                driver.modifiers = modifiers;
183                events.push(Event::PointerMoved(from));
184                events.push(press(from, held, true, modifiers));
185                driver.pass(std::mem::take(&mut events), &mut ui);
186                for step in 1..=DRAG_STEPS {
187                    let at = from.lerp(to, step as f32 / DRAG_STEPS as f32);
188                    events.push(Event::PointerMoved(at));
189                    driver.pass(std::mem::take(&mut events), &mut ui);
190                }
191                events.push(press(to, held, false, modifiers));
192                last = driver.pass(std::mem::take(&mut events), &mut ui);
193                driver.modifiers = Modifiers::NONE;
194                fresh = false;
195            }
196            Step::Scroll { at, lines } => {
197                events.push(Event::PointerMoved(Pos2::new(at[0], at[1])));
198                events.push(Event::MouseWheel {
199                    unit: MouseWheelUnit::Line,
200                    delta: Vec2::new(0.0, *lines),
201                    phase: TouchPhase::Move,
202                    modifiers: Modifiers::NONE,
203                });
204                last = driver.pass(std::mem::take(&mut events), &mut ui);
205                fresh = false;
206            }
207            Step::Frame => {
208                last = driver.pass(std::mem::take(&mut events), &mut ui);
209                fresh = false;
210            }
211        }
212    }
213    if fresh || !events.is_empty() {
214        last = driver.pass(events, &mut ui);
215    }
216    ctx.data_mut(|data| data.insert_temp(Id::new(CLOCK), driver.passes));
217    last.textures_delta = driver.textures;
218    last
219}
220
221pub fn check(size: [u32; 2]) -> Result<(), Error> {
222    if size.iter().any(|side| *side == 0 || *side > MAX_SIDE) {
223        return Err(Error(format!(
224            "a shot is 1 to {MAX_SIDE} pixels a side, not {}x{}",
225            size[0], size[1]
226        )));
227    }
228    Ok(())
229}
230
231pub fn pixels(app: &mut impl App, size: [u32; 2], script: &Script) -> Result<Vec<u8>, Error> {
232    check(size)?;
233    let theme = app.theme();
234    let gpu = block_on(Gpu::headless()).map_err(Error)?;
235    let format = wgpu::TextureFormat::Rgba8UnormSrgb;
236    let mut painter = Painter::new(&gpu, format, RendererOptions::PREDICTABLE);
237    if !app.viewport() {
238        return finish(&gpu, &mut painter, size, script, &theme, |ui, _| app.ui(ui));
239    }
240    app.forget();
241    let mut settled = script.clone();
242    settled
243        .steps
244        .extend(std::iter::repeat_n(Step::Frame, SETTLE));
245    finish(&gpu, &mut painter, size, &settled, &theme, |ui, painter| {
246        app.ui(ui);
247        app.frame(&gpu, painter.renderer(), ui.ctx().pixels_per_point());
248    })
249}
250
251fn finish(
252    gpu: &Gpu,
253    painter: &mut Painter,
254    size: [u32; 2],
255    script: &Script,
256    app_theme: &Theme,
257    mut ui: impl FnMut(&mut Ui, &mut Painter),
258) -> Result<Vec<u8>, Error> {
259    let ctx = context(app_theme);
260    let mut output = {
261        let painter = &mut *painter;
262        drive_painted(&ctx, size, script, |root| ui(root, painter))
263    };
264    let primitives = ctx.tessellate(output.shapes, output.pixels_per_point);
265    let target = gpu
266        .offscreen(size[0], size[1], wgpu::TextureFormat::Rgba8UnormSrgb)
267        .map_err(Error)?;
268    let screen = ScreenDescriptor {
269        size_in_pixels: size,
270        pixels_per_point: output.pixels_per_point,
271    };
272    painter.paint(
273        gpu,
274        &target.view,
275        &screen,
276        &output.textures_delta,
277        &primitives,
278        app_theme.bg,
279    );
280    output.textures_delta.clear();
281    gpu.readback_rgba8(&target).map_err(Error)
282}
283
284pub fn encode_png(size: [u32; 2], rgba: &[u8]) -> Result<Vec<u8>, Error> {
285    let mut bytes = Vec::new();
286    let mut encoder = png::Encoder::new(&mut bytes, size[0], size[1]);
287    encoder.set_color(png::ColorType::Rgba);
288    encoder.set_depth(png::BitDepth::Eight);
289    let mut writer = encoder
290        .write_header()
291        .map_err(|error| Error(error.to_string()))?;
292    writer
293        .write_image_data(rgba)
294        .map_err(|error| Error(error.to_string()))?;
295    writer.finish().map_err(|error| Error(error.to_string()))?;
296    Ok(bytes)
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302    use pfx_editor_style::fixture;
303
304    #[test]
305    fn sizes_outside_the_limit_are_refused() {
306        assert!(check([0, 10]).is_err());
307        assert!(check([10, MAX_SIDE + 1]).is_err());
308        assert!(check([1, MAX_SIDE]).is_ok());
309    }
310
311    #[test]
312    fn a_click_spends_a_pass_and_releases_on_the_next() {
313        let ctx = context(&fixture::theme());
314        let script = Script::parse("click 5 5\nframe").unwrap();
315        let mut seen = Vec::new();
316        drive(&ctx, [64, 32], &script, |ui| {
317            seen.push(ui.input(|input| input.pointer.primary_down()));
318        });
319        assert_eq!(seen, vec![false, true, false]);
320    }
321
322    #[test]
323    fn two_drives_on_one_context_never_move_the_clock_back() {
324        let ctx = context(&fixture::theme());
325        let script = Script::parse("click 5 5\nframe").unwrap();
326        let mut times = Vec::new();
327        for _ in 0..2 {
328            drive(&ctx, [64, 32], &script, |ui| {
329                times.push(ui.input(|input| input.time));
330            });
331        }
332        assert_eq!(times.len(), 6);
333        assert!(times.windows(2).all(|pair| pair[0] < pair[1]), "{times:?}");
334    }
335
336    #[test]
337    fn a_script_without_frames_still_draws_once() {
338        let ctx = context(&fixture::theme());
339        let mut passes = 0;
340        drive(&ctx, [64, 32], &Script::default(), |_| passes += 1);
341        assert_eq!(passes, 2);
342    }
343
344    #[test]
345    fn the_png_round_trips() {
346        let rgba: Vec<u8> = (0..4 * 3 * 2).map(|i| i as u8).collect();
347        let bytes = encode_png([3, 2], &rgba).unwrap();
348        let decoder = png::Decoder::new(std::io::Cursor::new(bytes));
349        let mut reader = decoder.read_info().unwrap();
350        let mut back = vec![0; reader.output_buffer_size().unwrap()];
351        reader.next_frame(&mut back).unwrap();
352        assert_eq!(back, rgba);
353    }
354
355    struct Blank;
356
357    impl App for Blank {
358        fn ui(&mut self, _: &mut Ui) {}
359
360        fn theme(&self) -> Theme {
361            fixture::theme()
362        }
363    }
364
365    #[test]
366    #[ignore = "GPU: pgpu run --class interactive --as pfx -- cargo test -- --ignored"]
367    fn an_empty_context_shoots_the_themes_solid_ground() {
368        let since = std::time::Instant::now();
369        let size = [96, 60];
370        let rgba = pixels(&mut Blank, size, &Script::default()).unwrap();
371        assert_eq!(rgba.len(), 96 * 60 * 4);
372        let ground = fixture::theme().bg;
373        let base = [ground.r(), ground.g(), ground.b(), 255];
374        assert!(rgba.chunks_exact(4).all(|pixel| pixel == base));
375        crate::gpu_turn(since);
376    }
377}