Fix too many points crash (#1822)

* Simplify point cloud builder and fix crash on too many points
Fixes #1779
* faster point cloud population by not chaining iterators with default values

crates/re_renderer/examples/2d.rs

@@ -150,28 +150,25 @@ impl framework::Example for Render2D {
         // Also, it looks different under perspective projection.
         // The third point is automatic thickness which is determined by the point renderer implementation.
         let mut point_cloud_builder = PointCloudBuilder::new(re_ctx);
-        point_cloud_builder
-            .batch("points")
-            .add_points_2d(
-                4,
-                [
-                    glam::vec2(500.0, 120.0),
-                    glam::vec2(520.0, 120.0),
-                    glam::vec2(540.0, 120.0),
-                    glam::vec2(560.0, 120.0),
-                ]
-                .into_iter(),
-            )
-            .radii(
-                [
-                    Size::new_scene(4.0),
-                    Size::new_points(4.0),
-                    Size::AUTO,
-                    Size::AUTO_LARGE,
-                ]
-                .into_iter(),
-            )
-            .colors(std::iter::repeat(Color32::from_rgb(55, 180, 1)).take(4));
+        point_cloud_builder.batch("points").add_points_2d(
+            4,
+            [
+                glam::vec2(500.0, 120.0),
+                glam::vec2(520.0, 120.0),
+                glam::vec2(540.0, 120.0),
+                glam::vec2(560.0, 120.0),
+            ]
+            .into_iter(),
+            [
+                Size::new_scene(4.0),
+                Size::new_points(4.0),
+                Size::AUTO,
+                Size::AUTO_LARGE,
+            ]
+            .into_iter(),
+            std::iter::repeat(Color32::from_rgb(55, 180, 1)),
+            std::iter::repeat(re_renderer::PickingLayerInstanceId::default()),
+        );
 
         // Pile stuff to test for overlap handling
         {

crates/re_renderer/examples/depth_cloud.rs

@@ -99,11 +99,13 @@ impl RenderDepthClouds {
                 .multiunzip();
 
             let mut builder = PointCloudBuilder::new(re_ctx);
-            builder
-                .batch("backprojected point cloud")
-                .add_points(num_points as _, points.into_iter())
-                .colors(colors.into_iter())
-                .radii(radii.into_iter());
+            builder.batch("backprojected point cloud").add_points(
+                num_points as _,
+                points.into_iter(),
+                radii.into_iter(),
+                colors.into_iter(),
+                std::iter::empty::<re_renderer::PickingLayerInstanceId>(),
+            );
 
             builder.to_draw_data(re_ctx).unwrap()
         };

crates/re_renderer/examples/multiview.rs

@@ -323,9 +323,10 @@ impl Example for Multiview {
             .add_points(
                 self.random_points_positions.len(),
                 self.random_points_positions.iter().cloned(),
-            )
-            .radii(self.random_points_radii.iter().cloned())
-            .colors(self.random_points_colors.iter().cloned());
+                self.random_points_radii.iter().cloned(),
+                self.random_points_colors.iter().cloned(),
+                std::iter::empty::<re_renderer::PickingLayerInstanceId>(),
+            );
 
         let point_cloud = builder.to_draw_data(re_ctx).unwrap();
         let meshes = build_mesh_instances(

crates/re_renderer/examples/picking.rs

@@ -165,10 +165,10 @@ impl framework::Example for Picking {
                 .add_points(
                     point_set.positions.len(),
                     point_set.positions.iter().cloned(),
-                )
-                .radii(point_set.radii.iter().cloned())
-                .colors(point_set.colors.iter().cloned())
-                .picking_instance_ids(point_set.picking_ids.iter().cloned());
+                    point_set.radii.iter().cloned(),
+                    point_set.colors.iter().cloned(),
+                    point_set.picking_ids.iter().cloned(),
+                );
         }
         view_builder.queue_draw(&point_builder.to_draw_data(re_ctx).unwrap());
 

crates/re_renderer/src/allocator/cpu_write_gpu_read_belt.rs

@@ -65,11 +65,15 @@ where
     /// Pushes several elements into the buffer.
     ///
     /// Panics if there are more elements than there is space in the buffer.
+    ///
+    /// Returns number of elements pushed.
     #[inline]
-    pub fn extend(&mut self, elements: impl Iterator<Item = T>) {
+    pub fn extend(&mut self, elements: impl Iterator<Item = T>) -> usize {
+        let num_written_before = self.num_written();
         for element in elements {
             self.push(element);
         }
+        self.num_written() - num_written_before
     }
 
     /// Pushes a single element into the buffer and advances the write pointer.

crates/re_renderer/src/point_cloud_builder.rs

@@ -114,7 +114,6 @@ impl<'a> Drop for PointCloudBatchBuilder<'a> {
         if self.0.batches.last().unwrap().point_count == 0 {
             self.0.batches.pop();
         }
-        self.extend_defaults();
     }
 }
 
@@ -141,92 +140,110 @@ impl<'a> PointCloudBatchBuilder<'a> {
         self
     }
 
-    /// Each time we `add_points`, or upon builder drop, make sure that we
-    /// fill in any additional colors and user-data to have matched vectors.
-    fn extend_defaults(&mut self) {
-        if self.0.color_buffer.num_written() < self.0.vertices.len() {
-            self.0.color_buffer.extend(
-                std::iter::repeat(Color32::WHITE)
-                    .take(self.0.vertices.len() - self.0.color_buffer.num_written()),
-            );
-        }
-
-        if self.0.picking_instance_ids_buffer.num_written() < self.0.vertices.len() {
-            self.0
-                .picking_instance_ids_buffer
-                .extend(std::iter::repeat(Default::default()).take(
-                    self.0.vertices.len() - self.0.picking_instance_ids_buffer.num_written(),
-                ));
-        }
-    }
-
-    #[inline]
     /// Add several 3D points
     ///
     /// Returns a `PointBuilder` which can be used to set the colors, radii, and user-data for the points.
     ///
-    /// Params:
-    ///  - `size_hint`: The `PointBuilder` will pre-allocate buffers to accommodate up to this number of points.
-    ///                 The resulting point batch, will still be determined by the length of the iterator.
-    ///  - `positions`: An iterable of the positions of the collection of points
+    /// Will *always* add `num_points`, no matter how many elements are in the iterators.
+    /// Missing elements will be filled up with defaults (in case of positions that's the origin)
+    ///
+    /// TODO(#957): Clamps number of points to the allowed per-builder maximum.
+    #[inline]
     pub fn add_points(
-        &mut self,
-        size_hint: usize,
+        mut self,
+        mut num_points: usize,
         positions: impl Iterator<Item = glam::Vec3>,
-    ) -> PointsBuilder<'_> {
+        radii: impl Iterator<Item = Size>,
+        colors: impl Iterator<Item = Color32>,
+        picking_instance_ids: impl Iterator<Item = PickingLayerInstanceId>,
+    ) -> Self {
         // TODO(jleibs): Figure out if we can plumb-through proper support for `Iterator::size_hints()`
         // or potentially make `FixedSizedIterator` work correctly. This should be possible size the
         // underlying arrow structures are of known-size, but carries some complexity with the amount of
         // chaining, joining, filtering, etc. that happens along the way.
         crate::profile_function!();
 
-        self.extend_defaults();
-
         debug_assert_eq!(self.0.vertices.len(), self.0.color_buffer.num_written());
+        debug_assert_eq!(
+            self.0.vertices.len(),
+            self.0.picking_instance_ids_buffer.num_written()
+        );
 
-        let old_size = self.0.vertices.len();
-
-        self.0.vertices.reserve(size_hint);
-        self.0.vertices.extend(positions.map(|p| PointCloudVertex {
-            position: p,
-            radius: Size::AUTO,
-        }));
-
-        let num_points = self.0.vertices.len() - old_size;
+        if num_points + self.0.vertices.len() > PointCloudDrawData::MAX_NUM_POINTS {
+            re_log::error_once!(
+                "Reached maximum number of supported points of {}.
+     See also https://github.com/rerun-io/rerun/issues/957",
+                PointCloudDrawData::MAX_NUM_POINTS
+            );
+            num_points = PointCloudDrawData::MAX_NUM_POINTS - self.0.vertices.len();
+        }
+        if num_points == 0 {
+            return self;
+        }
         self.batch_mut().point_count += num_points as u32;
 
-        let new_range = old_size..self.0.vertices.len();
-
-        let max_points = self.0.vertices.len();
-
-        PointsBuilder {
-            vertices: &mut self.0.vertices[new_range],
-            max_points,
-            colors: &mut self.0.color_buffer,
-            picking_instance_ids: &mut self.0.picking_instance_ids_buffer,
-            additional_outline_mask_ids: &mut self
+        {
+            crate::profile_scope!("positions");
+            let num_before = self.0.vertices.len();
+            self.0.vertices.extend(
+                positions
+                    .take(num_points)
+                    .zip(radii.take(num_points))
+                    .map(|(position, radius)| PointCloudVertex { position, radius }),
+            );
+            // Fill up with defaults. Doing this in a separate step is faster than chaining the iterator.
+            let num_default = num_points - (self.0.vertices.len() - num_before);
+            self.0.vertices.extend(
+                std::iter::repeat(PointCloudVertex {
+                    position: glam::Vec3::ZERO,
+                    radius: Size::AUTO,
+                })
+                .take(num_default),
+            );
+        }
+        {
+            crate::profile_scope!("colors");
+            let num_written = self.0.color_buffer.extend(colors.take(num_points));
+            // Fill up with defaults. Doing this in a separate step is faster than chaining the iterator.
+            self.0
+                .color_buffer
+                .extend(std::iter::repeat(Color32::TRANSPARENT).take(num_points - num_written));
+        }
+        {
+            crate::profile_scope!("picking_instance_ids");
+            let num_written = self
                 .0
-                .batches
-                .last_mut()
-                .unwrap()
-                .additional_outline_mask_ids_vertex_ranges,
-            start_vertex_index: old_size as _,
+                .picking_instance_ids_buffer
+                .extend(picking_instance_ids.take(num_points));
+            // Fill up with defaults. Doing this in a separate step is faster than chaining the iterator.
+            self.0.picking_instance_ids_buffer.extend(
+                std::iter::repeat(PickingLayerInstanceId::default()).take(num_points - num_written),
+            );
         }
+
+        self
     }
 
     /// Adds several 2D points. Uses an autogenerated depth value, the same for all points passed.
     ///
-    /// Params:
-    ///  - `size_hint`: The `PointBuilder` will pre-allocate buffers to accommodate up to this number of points.
-    ///                 The resulting point batch, will be the size of the length of the `positions` iterator.
-    ///  - `positions`: An iterable of the positions of the collection of points
+    /// Will *always* add `num_points`, no matter how many elements are in the iterators.
+    /// Missing elements will be filled up with defaults (in case of positions that's the origin)
     #[inline]
     pub fn add_points_2d(
-        &mut self,
-        size_hint: usize,
+        self,
+        num_points: usize,
         positions: impl Iterator<Item = glam::Vec2>,
-    ) -> PointsBuilder<'_> {
-        self.add_points(size_hint, positions.map(|p| p.extend(0.0)))
+        radii: impl Iterator<Item = Size>,
+        colors: impl Iterator<Item = Color32>,
+        picking_instance_ids: impl Iterator<Item = PickingLayerInstanceId>,
+    ) -> Self {
+        self.add_points(
+            num_points,
+            positions.map(|p| p.extend(0.0)),
+            radii,
+            colors,
+            picking_instance_ids,
+        )
     }
 
     /// Set flags for this batch.
@@ -235,80 +252,25 @@ impl<'a> PointCloudBatchBuilder<'a> {
         self
     }
 
+    /// Sets the picking object id for the current batch.
     pub fn picking_object_id(mut self, picking_object_id: PickingLayerObjectId) -> Self {
         self.batch_mut().picking_object_id = picking_object_id;
         self
     }
-}
-
-pub struct PointsBuilder<'a> {
-    // Vertices is a slice, which radii will update
-    vertices: &'a mut [PointCloudVertex],
-    max_points: usize,
-    colors: &'a mut CpuWriteGpuReadBuffer<Color32>,
-    picking_instance_ids: &'a mut CpuWriteGpuReadBuffer<PickingLayerInstanceId>,
-    additional_outline_mask_ids: &'a mut Vec<(std::ops::Range<u32>, OutlineMaskPreference)>,
-    start_vertex_index: u32,
-}
-
-impl<'a> PointsBuilder<'a> {
-    /// Assigns radii to all points.
-    ///
-    /// This mustn't call this more than once.
-    ///
-    /// If the iterator doesn't cover all points, some will not be assigned.
-    /// If the iterator provides more values than there are points, the extra values will be ignored.
-    #[inline]
-    pub fn radii(self, radii: impl Iterator<Item = Size>) -> Self {
-        // TODO(andreas): This seems like an argument for moving radius
-        // to a separate storage
-        crate::profile_function!();
-        for (point, radius) in self.vertices.iter_mut().zip(radii) {
-            point.radius = radius;
-        }
-        self
-    }
-
-    /// Assigns colors to all points.
-    ///
-    /// This mustn't call this more than once.
-    ///
-    /// If the iterator doesn't cover all points, some will not be assigned.
-    /// If the iterator provides more values than there are points, the extra values will be ignored.
-    #[inline]
-    pub fn colors(self, colors: impl Iterator<Item = Color32>) -> Self {
-        crate::profile_function!();
-        self.colors
-            .extend(colors.take(self.max_points - self.colors.num_written()));
-        self
-    }
-
-    #[inline]
-    pub fn picking_instance_ids(
-        self,
-        picking_instance_ids: impl Iterator<Item = PickingLayerInstanceId>,
-    ) -> Self {
-        crate::profile_function!();
-        self.picking_instance_ids.extend(
-            picking_instance_ids.take(self.max_points - self.picking_instance_ids.num_written()),
-        );
-        self
-    }
 
     /// Pushes additional outline mask ids for a specific range of points.
-    /// The range is relative to this builder's range, not the entire batch.
+    /// The range is relative to this batch.
     ///
     /// Prefer the `overall_outline_mask_ids` setting to set the outline mask ids for the entire batch whenever possible!
     #[inline]
     pub fn push_additional_outline_mask_ids_for_range(
-        self,
+        mut self,
         range: std::ops::Range<u32>,
         ids: OutlineMaskPreference,
     ) -> Self {
-        self.additional_outline_mask_ids.push((
-            (range.start + self.start_vertex_index)..(range.end + self.start_vertex_index),
-            ids,
-        ));
+        self.batch_mut()
+            .additional_outline_mask_ids_vertex_ranges
+            .push((range, ids));
         self
     }
 }

crates/re_renderer/src/renderer/point_cloud.rs

@@ -142,6 +142,7 @@ pub struct PointCloudBatchInfo {
 }
 
 /// Description of a point cloud.
+#[derive(Clone)]
 pub struct PointCloudVertex {
     /// Connected points. Must be at least 2.
     pub position: glam::Vec3,
@@ -225,7 +226,7 @@ impl PointCloudDrawData {
             0
         );
 
-        let vertices = if vertices.len() >= Self::MAX_NUM_POINTS {
+        let vertices = if vertices.len() > Self::MAX_NUM_POINTS {
             re_log::error_once!(
                 "Reached maximum number of supported points. Clamping down to {}, passed were {}.
  See also https://github.com/rerun-io/rerun/issues/957",