From 56e35fc3f99bc68fb160c3ee975290995613f4b8 Mon Sep 17 00:00:00 2001
From: Ian Manske <ian.manske@pm.me>
Date: Thu, 2 Nov 2023 19:01:46 +0000
Subject: [PATCH] Reduce element shifting in `Record::retain_mut` (#10915)

# Description
Replaces the `Vec::remove` in `Record::retain_mut` with some swaps which
should eliminate the `O(n^2)` complexity due to repeated shifting of
elements.
---
 crates/nu-protocol/src/value/record.rs | 27 +++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/crates/nu-protocol/src/value/record.rs b/crates/nu-protocol/src/value/record.rs
index 5ba2b70b86..674c4670e7 100644
--- a/crates/nu-protocol/src/value/record.rs
+++ b/crates/nu-protocol/src/value/record.rs
@@ -172,24 +172,33 @@ impl Record {
     where
         F: FnMut(&str, &mut Value) -> bool,
     {
+        // `Vec::retain` is able to optimize memcopies internally.
+        // For maximum benefit, `retain` is used on `vals`,
+        // as `Value` is a larger struct than `String`.
+        //
+        // To do a simultaneous retain on the `cols`, three portions of it are tracked:
+        //     [..retained, ..dropped, ..unvisited]
+
+        // number of elements keep so far, start of ..dropped and length of ..retained
+        let mut retained = 0;
+        // current index of element being checked, start of ..unvisited
         let mut idx = 0;
 
-        // `Vec::retain` is able to optimize memcopies internally. For maximum benefit as `Value`
-        // is a larger struct than `String` use `retain` on `vals`
-        //
-        // The calls to `Vec::remove` are suboptimal as they need memcopies to shift each time.
-        //
-        // As the operations should remain inplace, we don't allocate a separate index `Vec` which
-        // could be used to avoid the repeated shifting of `Vec::remove` in cols.
         self.vals.retain_mut(|val| {
-            if keep(self.cols[idx].as_str(), val) {
+            if keep(&self.cols[idx], val) {
+                // skip swaps for first consecutive run of kept elements
+                if idx != retained {
+                    self.cols.swap(idx, retained);
+                }
+                retained += 1;
                 idx += 1;
                 true
             } else {
-                self.cols.remove(idx);
+                idx += 1;
                 false
             }
         });
+        self.cols.truncate(retained);
     }
 
     pub fn columns(&self) -> Columns {