Auto merge of #113106 - marcospb19:improve-path-with-extension-function, r=thomcc

std: remove an allocation in `Path::with_extension`

`Path::with_extension` used to reallocate (and copy) paths twice per call, now it does it once, by checking the size of the previous and new extensions it's possible to call `PathBuf::with_capacity` and pass the exact capacity required.

This also reduces the memory consumption of the path returned from `Path::with_extension` by using exact capacity instead of using amortized exponential growth.

2 files changed, 66 insertions, 8 deletions

diff --git a/library/std/src/path.rs b/library/std/src/path.rs
index 28cd3c4e4db..99f7a60f8ab 100644
--- a/library/std/src/path.rs
+++ b/library/std/src/path.rs
@@ -2608,9 +2608,27 @@ impl Path {
     }
 
     fn _with_extension(&self, extension: &OsStr) -> PathBuf {
-        let mut buf = self.to_path_buf();
-        buf.set_extension(extension);
-        buf
+        let self_len = self.as_os_str().len();
+        let self_bytes = self.as_os_str().as_os_str_bytes();
+
+        let (new_capacity, slice_to_copy) = match self.extension() {
+            None => {
+                // Enough capacity for the extension and the dot
+                let capacity = self_len + extension.len() + 1;
+                let whole_path = self_bytes.iter();
+                (capacity, whole_path)
+            }
+            Some(previous_extension) => {
+                let capacity = self_len + extension.len() - previous_extension.len();
+                let path_till_dot = self_bytes[..self_len - previous_extension.len()].iter();
+                (capacity, path_till_dot)
+            }
+        };
+
+        let mut new_path = PathBuf::with_capacity(new_capacity);
+        new_path.as_mut_vec().extend(slice_to_copy);
+        new_path.set_extension(extension);
+        new_path
     }
 
     /// Produces an iterator over the [`Component`]s of the path.
diff --git a/library/std/src/path/tests.rs b/library/std/src/path/tests.rs
index dd307022c6d..f12ffbf2e01 100644
--- a/library/std/src/path/tests.rs
+++ b/library/std/src/path/tests.rs
@@ -1183,7 +1183,7 @@ pub fn test_prefix_ext() {
 #[test]
 pub fn test_push() {
     macro_rules! tp (
-        ($path:expr, $push:expr, $expected:expr) => ( {
+        ($path:expr, $push:expr, $expected:expr) => ({
             let mut actual = PathBuf::from($path);
             actual.push($push);
             assert!(actual.to_str() == Some($expected),
@@ -1281,7 +1281,7 @@ pub fn test_push() {
 #[test]
 pub fn test_pop() {
     macro_rules! tp (
-        ($path:expr, $expected:expr, $output:expr) => ( {
+        ($path:expr, $expected:expr, $output:expr) => ({
             let mut actual = PathBuf::from($path);
             let output = actual.pop();
             assert!(actual.to_str() == Some($expected) && output == $output,
@@ -1335,7 +1335,7 @@ pub fn test_pop() {
 #[test]
 pub fn test_set_file_name() {
     macro_rules! tfn (
-            ($path:expr, $file:expr, $expected:expr) => ( {
+        ($path:expr, $file:expr, $expected:expr) => ({
             let mut p = PathBuf::from($path);
             p.set_file_name($file);
             assert!(p.to_str() == Some($expected),
@@ -1369,7 +1369,7 @@ pub fn test_set_file_name() {
 #[test]
 pub fn test_set_extension() {
     macro_rules! tfe (
-            ($path:expr, $ext:expr, $expected:expr, $output:expr) => ( {
+        ($path:expr, $ext:expr, $expected:expr, $output:expr) => ({
             let mut p = PathBuf::from($path);
             let output = p.set_extension($ext);
             assert!(p.to_str() == Some($expected) && output == $output,
@@ -1395,6 +1395,46 @@ pub fn test_set_extension() {
 }
 
 #[test]
+pub fn test_with_extension() {
+    macro_rules! twe (
+        ($input:expr, $extension:expr, $expected:expr) => ({
+            let input = Path::new($input);
+            let output = input.with_extension($extension);
+
+            assert!(
+                output.to_str() == Some($expected),
+                "calling Path::new({:?}).with_extension({:?}): Expected {:?}, got {:?}",
+                $input, $extension, $expected, output,
+            );
+        });
+    );
+
+    twe!("foo", "txt", "foo.txt");
+    twe!("foo.bar", "txt", "foo.txt");
+    twe!("foo.bar.baz", "txt", "foo.bar.txt");
+    twe!(".test", "txt", ".test.txt");
+    twe!("foo.txt", "", "foo");
+    twe!("foo", "", "foo");
+    twe!("", "foo", "");
+    twe!(".", "foo", ".");
+    twe!("foo/", "bar", "foo.bar");
+    twe!("foo/.", "bar", "foo.bar");
+    twe!("..", "foo", "..");
+    twe!("foo/..", "bar", "foo/..");
+    twe!("/", "foo", "/");
+
+    // New extension is smaller than file name
+    twe!("aaa_aaa_aaa", "bbb_bbb", "aaa_aaa_aaa.bbb_bbb");
+    // New extension is greater than file name
+    twe!("bbb_bbb", "aaa_aaa_aaa", "bbb_bbb.aaa_aaa_aaa");
+
+    // New extension is smaller than previous extension
+    twe!("ccc.aaa_aaa_aaa", "bbb_bbb", "ccc.bbb_bbb");
+    // New extension is greater than previous extension
+    twe!("ccc.bbb_bbb", "aaa_aaa_aaa", "ccc.aaa_aaa_aaa");
+}
+
+#[test]
 fn test_eq_receivers() {
     use crate::borrow::Cow;
 
@@ -1669,7 +1709,7 @@ fn into_rc() {
 #[test]
 fn test_ord() {
     macro_rules! ord(
-        ($ord:ident, $left:expr, $right:expr) => ( {
+        ($ord:ident, $left:expr, $right:expr) => ({
             use core::cmp::Ordering;
 
             let left = Path::new($left);