From 2b72b9f9f3c796ec32018b004d787a700b994399 Mon Sep 17 00:00:00 2001
From: Armin Becher <Armin Becher>
Date: Sun, 30 Apr 2023 23:10:09 +0200
Subject: [PATCH] perf: improve remove_col speed

---
 benches/benches.rs |  2 +-
 src/lib.rs         | 17 +++++++++++++----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/benches/benches.rs b/benches/benches.rs
index 3e55238..cfa77f5 100644
--- a/benches/benches.rs
+++ b/benches/benches.rs
@@ -157,7 +157,7 @@ fn criterion_benchmark(c: &mut Criterion) {
             criterion::BatchSize::SmallInput,
         )
     });
-    
+
     // Remove
     c.bench_function("grid_remove_row", |b| {
         let grid = init_grid();
diff --git a/src/lib.rs b/src/lib.rs
index e1cb174..aa5501a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -40,6 +40,7 @@ extern crate alloc;
 #[cfg(all(not(feature = "std")))]
 use alloc::{vec::Vec, vec, format};
 
+use core::cmp;
 use core::cmp::Eq;
 use core::fmt;
 use core::iter::StepBy;
@@ -575,6 +576,9 @@ impl<T> Grid<T> {
             .drain((row_index * self.cols)..((row_index + 1) * self.cols));
 
         self.rows -= 1;
+        if self.rows == 0 {
+            self.cols = 0;
+        }
         Some(residue.collect())
     }
 
@@ -624,12 +628,17 @@ impl<T> Grid<T> {
         if self.cols == 0 || self.rows == 0 || col_index >= self.cols {
             return None;
         }
-        let mut residue: Vec<T> = vec![];
         for i in 0..self.rows {
-			residue.push(self.data.remove(i * self.cols + col_index - i));
+            let row_idx = col_index + i * (self.cols - 1);
+            let end = cmp::min(row_idx+self.cols+i, self.data.len());
+            self.data[row_idx..end].rotate_left(i + 1);
         }
-       	self.cols -= 1;
-        Some(residue)
+        let col = self.data.split_off(self.data.len() - self.rows);
+        self.cols -= 1;
+        if self.cols == 0 {
+            self.rows = 0;
+        }
+        Some(col)
     }
 
     /// Insert a new row at the index and shifts all rows after down.