Fix performance on large canvases

src/handwriting/canvas_rasterizer.rs

@@ -2,80 +2,175 @@ use std::sync::Arc;
 
 use egui::{
     Color32, ColorImage, CornerRadius, Painter, Pos2, Rect, Stroke, StrokeKind, TextureHandle,
+    Vec2,
+    ahash::HashMap,
     emath::TSTransform,
     epaint::{Brush, RectShape, Vertex},
     load::SizedTexture,
 };
 
-use crate::rasterizer::rasterize_onto;
+use crate::rasterizer::{PxBoundingBox, rasterize_triangle_onto, triangle_bounding_box};
 
 use super::StrokeBlendMode;
 
+const CHUNK_SIZE: usize = 64;
+
 /// Rasterize onto a resizeable canvas.
 #[derive(Default)]
 pub struct CanvasRasterizer {
+    image_size: [usize; 2],
+    tiles: HashMap<[usize; 2], Tile>,
+}
+
+struct Tile {
+    bounding_box: PxBoundingBox,
     image: ColorImage,
     texture: Option<TextureHandle>,
     texture_is_dirty: bool,
 }
 
+impl Tile {
+    fn new(xi: usize, yi: usize) -> Self {
+        let x_from = xi * CHUNK_SIZE;
+        let y_from = yi * CHUNK_SIZE;
+        let bounding_box = PxBoundingBox {
+            x_from,
+            y_from,
+            x_to: x_from + CHUNK_SIZE,
+            y_to: y_from + CHUNK_SIZE,
+        };
+
+        Self {
+            bounding_box,
+            image: ColorImage::new([CHUNK_SIZE, CHUNK_SIZE], Color32::TRANSPARENT),
+            texture: None,
+            texture_is_dirty: false,
+        }
+    }
+}
+
 impl CanvasRasterizer {
-    //pub fn clear(&mut self) {
-    //    self.image = ColorImage::new(self.image.size, Color32::TRANSPARENT);
-    //    self.texture_is_dirty = true;
-    //}
+    pub fn set_size(&mut self, width: usize, height: usize) {
+        self.image_size = [width, height];
+        self.populate_tiles();
+    }
 
-    //pub fn set_size(&mut self, width: usize, height: usize) {
-    //    if self.image.size != [width, height] {
-    //        self.image = ColorImage::new([width, height], Color32::TRANSPARENT);
-    //        self.texture_is_dirty = true;
-    //    }
-    //}
+    pub fn clear(&mut self) {
+        log::error!("clearing all tiles");
+        self.tiles.clear();
+        self.populate_tiles();
+    }
 
-    pub fn clear_and_set_size(&mut self, width: usize, height: usize) {
-        self.image = ColorImage::new([width, height], Color32::TRANSPARENT);
-        self.texture_is_dirty = true;
+    fn populate_tiles(&mut self) {
+        let [width, height] = self.image_size;
+
+        // discard tiles that are out of bounds
+        self.tiles.retain(|_, tile| {
+            tile.bounding_box.x_from <= width && tile.bounding_box.y_from <= height
+        });
+
+        let chunk = |max: usize| {
+            (0..)
+                .step_by(CHUNK_SIZE)
+                .take_while(move |n| n <= &max)
+                .enumerate()
+        };
+
+        // create new tiles where we need them
+        for (xi, _x) in chunk(width) {
+            for (yi, _y) in chunk(height) {
+                self.tiles
+                    .entry([xi, yi])
+                    .or_insert_with(|| Tile::new(xi, yi));
+            }
+        }
     }
 
     pub fn rasterize<'a>(
         &mut self,
         point_to_pixel: TSTransform,
-        triangles: impl Iterator<Item = [&'a Vertex; 3]>,
+        triangles: impl Iterator<Item = [&'a Vertex; 3]> + Clone,
     ) {
-        rasterize_onto::<StrokeBlendMode>(&mut self.image, point_to_pixel, triangles);
-        self.texture_is_dirty = true;
+        for triangle in triangles {
+            let triangle_bounding_box = triangle_bounding_box(&triangle, point_to_pixel);
+            for chunk in chunks_from_bounding_box(triangle_bounding_box) {
+                let Some(tile) = self.tiles.get_mut(&chunk) else {
+                    continue;
+                };
+
+                let mut point_to_tile_pixel = point_to_pixel;
+                point_to_tile_pixel.translation -= Vec2::new(
+                    tile.bounding_box.x_from as f32,
+                    tile.bounding_box.y_from as f32,
+                );
+
+                tile.texture_is_dirty = true;
+                rasterize_triangle_onto::<StrokeBlendMode>(
+                    &mut tile.image,
+                    point_to_tile_pixel,
+                    triangle,
+                );
+            }
+        }
     }
 
+    /// `at` defines the location in screen-coordinates where the canvas should be drawn.
     pub fn show(&mut self, ctx: &egui::Context, painter: &Painter, at: Rect) {
-        if self.texture_is_dirty {
-            self.texture_is_dirty = false;
-            let new_image = || {
-                let image = ColorImage::new(self.image.size, Color32::TRANSPARENT);
-                ctx.load_texture("handwriting", image, Default::default())
-            };
-            let texture = self.texture.get_or_insert_with(new_image);
-            texture.set(self.image.clone(), Default::default());
-        }
+        let pixels_per_point = ctx.pixels_per_point();
+        let chunk_vec = Vec2::splat(CHUNK_SIZE as f32) / pixels_per_point;
 
-        if let Some(texture) = &mut self.texture {
-            let texture = SizedTexture::new(texture.id(), texture.size_vec2());
-            let shape = RectShape {
-                rect: at,
-                corner_radius: CornerRadius::ZERO,
-                fill: Color32::WHITE,
-                stroke: Stroke::NONE,
-                stroke_kind: StrokeKind::Inside,
-                round_to_pixels: None,
-                blur_width: 0.0,
-                brush: Some(Arc::new(Brush {
-                    fill_texture_id: texture.id,
-                    uv: Rect {
-                        min: Pos2::ZERO,
-                        max: Pos2::new(1.0, 1.0),
-                    },
-                })),
-            };
-            painter.add(shape);
+        for ([xi, yi], tile) in &mut self.tiles {
+            if tile.texture_is_dirty {
+                tile.texture_is_dirty = false;
+                if let Some(texture) = &mut tile.texture {
+                    texture.set(tile.image.clone(), Default::default());
+                } else {
+                    tile.texture = Some(ctx.load_texture(
+                        "handwriting",
+                        tile.image.clone(),
+                        Default::default(),
+                    ));
+                }
+            }
+
+            if let Some(texture) = &mut tile.texture {
+                let texture = SizedTexture::new(texture.id(), texture.size_vec2());
+                let shape = RectShape {
+                    rect: Rect::from_min_size(
+                        at.min + Vec2::new(*xi as f32, *yi as f32) * chunk_vec,
+                        chunk_vec,
+                    ),
+                    corner_radius: CornerRadius::ZERO,
+                    fill: Color32::WHITE,
+                    stroke: Stroke::NONE,
+                    stroke_kind: StrokeKind::Inside,
+                    round_to_pixels: None,
+                    blur_width: 0.0,
+                    brush: Some(Arc::new(Brush {
+                        fill_texture_id: texture.id,
+                        uv: Rect {
+                            min: Pos2::ZERO,
+                            max: Pos2::new(1.0, 1.0),
+                        },
+                    })),
+                };
+                painter.add(shape);
+            }
         }
     }
 }
+
+/// Get all chunk indices that overlaps with a PxBoundingBox.
+fn chunks_from_bounding_box(
+    triangle_bounding_box: PxBoundingBox,
+) -> impl Iterator<Item = [usize; 2]> {
+    let x_from_chunk = triangle_bounding_box.x_from / CHUNK_SIZE;
+    let y_from_chunk = triangle_bounding_box.y_from / CHUNK_SIZE;
+    let x_to_chunk = triangle_bounding_box.x_to.saturating_sub(1) / CHUNK_SIZE;
+    let y_to_chunk = triangle_bounding_box.y_to.saturating_sub(1) / CHUNK_SIZE;
+
+    let xs = x_from_chunk..=x_to_chunk;
+    let ys = y_from_chunk..=y_to_chunk;
+
+    ys.flat_map(move |yi| xs.clone().map(move |xi| [xi, yi]))
+}

src/handwriting/mod.rs

@@ -105,20 +105,9 @@ struct MeshContext {
 
     pub pixels_per_point: f32,
 
-    /// Canvas size in points
-    pub size: Vec2,
-
     pub stroke: Stroke,
 }
 
-impl MeshContext {
-    /// Calculate canvas size in pixels
-    pub fn pixel_size(&self) -> [usize; 2] {
-        let Vec2 { x, y } = self.size * self.pixels_per_point;
-        [x, y].map(|f| f.ceil() as usize)
-    }
-}
-
 impl Default for Handwriting {
     fn default() -> Self {
         Self {
@@ -379,7 +368,6 @@ impl Handwriting {
         let new_context = MeshContext {
             ui_theme: ui.ctx().theme(),
             pixels_per_point: ui.pixels_per_point(),
-            size: mesh_rect.size(),
             stroke: style.stroke,
         };
 
@@ -388,6 +376,13 @@ impl Handwriting {
             self.e.refresh_texture = true;
         }
 
+        let [px_width, px_height] = {
+            let Vec2 { x, y } = mesh_rect.size() * new_context.pixels_per_point;
+            [x, y].map(|f| f.ceil() as usize)
+        };
+
+        self.e.canvas_rasterizer.set_size(px_width, px_height);
+
         if self.e.refresh_texture {
             // ...if we do, rasterize the entire texture from scratch
             self.refresh_texture(style, new_context);
@@ -457,11 +452,10 @@ impl Handwriting {
         //    debug_assert!(vertex.pos.y.is_finite(), "{} must be finite", vertex.pos.y);
         //}
 
-        let [px_x, px_y] = mesh_context.pixel_size();
         let point_to_pixel = TSTransform::from_scaling(mesh_context.pixels_per_point);
         let triangles = mesh_triangles(&self.e.mesh);
 
-        self.e.canvas_rasterizer.clear_and_set_size(px_x, px_y);
+        self.e.canvas_rasterizer.clear();
         self.e
             .canvas_rasterizer
             .rasterize(point_to_pixel, triangles);
@@ -721,8 +715,10 @@ impl HandwritingStyle {
     }
 }
 
-fn mesh_triangles(mesh: &Mesh) -> impl Iterator<Item = [&Vertex; 3]> {
-    mesh.triangles()
+fn mesh_triangles(mesh: &Mesh) -> impl Iterator<Item = [&Vertex; 3]> + Clone {
+    mesh.indices
+        .chunks_exact(3)
+        .map(|chunk| [chunk[0], chunk[1], chunk[2]])
         .map(|indices| indices.map(|i| &mesh.vertices[i as usize]))
 }
 

src/rasterizer.rs

@@ -1,5 +1,6 @@
 use core::f32;
 use egui::{Color32, ColorImage, Pos2, Rect, Vec2, emath::TSTransform, epaint::Vertex};
+use std::ops::Range;
 
 pub trait BlendFn {
     fn blend(a: Color32, b: Color32) -> Color32;
@@ -86,9 +87,20 @@ pub fn rasterize_onto<'a, Blend: BlendFn>(
     }
 }
 
+/// Rasterize a single triangles onto an image,
+///
+/// Triangle positions must be in image-local point-coords.
+pub fn rasterize_triangle_onto<'a, Blend: BlendFn>(
+    image: &mut ColorImage,
+    point_to_pixel: TSTransform,
+    triangle: [&'a Vertex; 3],
+) {
+    rasterize_onto::<Blend>(image, point_to_pixel, [triangle].into_iter());
+}
+
 /// A bounding box, measured in pixels.
 #[derive(Debug, PartialEq, Eq)]
-struct PxBoundingBox {
+pub struct PxBoundingBox {
     pub x_from: usize,
     pub y_from: usize,
     pub x_to: usize,
@@ -104,9 +116,42 @@ impl PxBoundingBox {
             y_to: self.y_to.min(other.y_to),
         }
     }
+
+    pub fn union(&self, other: &PxBoundingBox) -> PxBoundingBox {
+        PxBoundingBox {
+            x_from: self.x_from.min(other.x_from),
+            y_from: self.y_from.min(other.y_from),
+            x_to: self.x_to.max(other.x_to),
+            y_to: self.y_to.max(other.y_to),
+        }
+    }
+
+    pub fn x_range(&self) -> Range<usize> {
+        self.x_from..self.x_to
+    }
+
+    pub fn y_range(&self) -> Range<usize> {
+        self.y_from..self.y_to
+    }
+
+    /// Test whether two boxes do NOT overlap
+    pub fn overlaps_with(&self, other: &PxBoundingBox) -> bool {
+        !self.is_disjoint_from(other)
+    }
+
+    pub fn is_disjoint_from(&self, other: &PxBoundingBox) -> bool {
+        false
+            || self.x_from > other.x_to
+            || self.y_from > other.y_to
+            || other.x_from > self.x_to
+            || other.y_from > self.y_to
+    }
 }
 
-fn triangle_bounding_box(triangle: &[&Vertex; 3], point_to_pixel: TSTransform) -> PxBoundingBox {
+pub fn triangle_bounding_box(
+    triangle: &[&Vertex; 3],
+    point_to_pixel: TSTransform,
+) -> PxBoundingBox {
     // calculate bounding box in point coords
     let mut rect = Rect::NOTHING;
     for vertex in triangle {