From 928dee830e5078a18a17c76db81fb4b2f8debec1 Mon Sep 17 00:00:00 2001
From: Stepan Koltsov <stepan.koltsov@gmail.com>
Date: Thu, 31 Oct 2024 16:39:32 +0000
Subject: [PATCH] Indices::extend (#16023)

# Objective

There's integer overflow in `Mesh::merge` in branches like this:


https://github.com/bevyengine/bevy/blob/405fa3e8ea1d8ff23f0590a18242667e2b21e2d9/crates/bevy_mesh/src/mesh.rs#L857-L859

we truncate `u32` to `u16` and ignore integer overflow on `u16`. This
may lead to unexpected results when the number of vertices exceeds
`u16::MAX`.

## Solution

Convert indices storage to `u32` when necessary.

## Testing

- Unit test added for `extend` function
- For changes in `Mesh`, I presume it is already tested elsewhere
---
 crates/bevy_mesh/src/index.rs | 62 ++++++++++++++++++++++++++++-------
 crates/bevy_mesh/src/mesh.rs  | 15 +--------
 2 files changed, 52 insertions(+), 25 deletions(-)

diff --git a/crates/bevy_mesh/src/index.rs b/crates/bevy_mesh/src/index.rs
index af2b3635ba..096cd031d6 100644
--- a/crates/bevy_mesh/src/index.rs
+++ b/crates/bevy_mesh/src/index.rs
@@ -103,19 +103,36 @@ impl Indices {
     /// Add an index. If the index is greater than `u16::MAX`,
     /// the storage will be converted to `u32`.
     pub fn push(&mut self, index: u32) {
+        self.extend([index]);
+    }
+}
+
+/// Extend the indices with indices from an iterator.
+/// Semantically equivalent to calling [`push`](Indices::push) for each element in the iterator,
+/// but more efficient.
+impl Extend<u32> for Indices {
+    fn extend<T: IntoIterator<Item = u32>>(&mut self, iter: T) {
+        let mut iter = iter.into_iter();
         match self {
-            Indices::U32(vec) => vec.push(index),
-            Indices::U16(vec) => match u16::try_from(index) {
-                Ok(index) => vec.push(index),
-                Err(_) => {
-                    let new_vec = vec
-                        .iter()
-                        .map(|&index| u32::from(index))
-                        .chain(iter::once(index))
-                        .collect::<Vec<u32>>();
-                    *self = Indices::U32(new_vec);
+            Indices::U32(indices) => indices.extend(iter),
+            Indices::U16(indices) => {
+                indices.reserve(iter.size_hint().0);
+                while let Some(index) = iter.next() {
+                    match u16::try_from(index) {
+                        Ok(index) => indices.push(index),
+                        Err(_) => {
+                            let new_vec = indices
+                                .iter()
+                                .map(|&index| u32::from(index))
+                                .chain(iter::once(index))
+                                .chain(iter)
+                                .collect::<Vec<u32>>();
+                            *self = Indices::U32(new_vec);
+                            break;
+                        }
+                    }
                 }
-            },
+            }
         }
     }
 }
@@ -181,4 +198,27 @@ mod tests {
             indices.iter().collect::<Vec<_>>()
         );
     }
+
+    #[test]
+    fn test_indices_extend() {
+        let mut indices = Indices::U16(Vec::new());
+        indices.extend([10, 11]);
+        assert_eq!(IndexFormat::Uint16, IndexFormat::from(&indices));
+        assert_eq!(vec![10, 11], indices.iter().collect::<Vec<_>>());
+
+        // Add a value that is too large for `u16` so the storage should be converted to `U32`.
+        indices.extend([12, 0x10013, 0x10014]);
+        assert_eq!(IndexFormat::Uint32, IndexFormat::from(&indices));
+        assert_eq!(
+            vec![10, 11, 12, 0x10013, 0x10014],
+            indices.iter().collect::<Vec<_>>()
+        );
+
+        indices.extend([15, 0x10016]);
+        assert_eq!(IndexFormat::Uint32, IndexFormat::from(&indices));
+        assert_eq!(
+            vec![10, 11, 12, 0x10013, 0x10014, 15, 0x10016],
+            indices.iter().collect::<Vec<_>>()
+        );
+    }
 }
diff --git a/crates/bevy_mesh/src/mesh.rs b/crates/bevy_mesh/src/mesh.rs
index f2a9646d04..22698c4e68 100644
--- a/crates/bevy_mesh/src/mesh.rs
+++ b/crates/bevy_mesh/src/mesh.rs
@@ -842,20 +842,7 @@ impl Mesh {
 
         // Extend indices of `self` with indices of `other`.
         if let (Some(indices), Some(other_indices)) = (self.indices_mut(), other.indices()) {
-            match (indices, other_indices) {
-                (Indices::U16(i1), Indices::U16(i2)) => {
-                    i1.extend(i2.iter().map(|i| *i + index_offset as u16));
-                }
-                (Indices::U32(i1), Indices::U32(i2)) => {
-                    i1.extend(i2.iter().map(|i| *i + index_offset as u32));
-                }
-                (Indices::U16(i1), Indices::U32(i2)) => {
-                    i1.extend(i2.iter().map(|i| *i as u16 + index_offset as u16));
-                }
-                (Indices::U32(i1), Indices::U16(i2)) => {
-                    i1.extend(i2.iter().map(|i| *i as u32 + index_offset as u32));
-                }
-            }
+            indices.extend(other_indices.iter().map(|i| (i + index_offset) as u32));
         }
     }