fs: avoid some copies in tokio::fs::write (#7199)

tokio/src/fs/write.rs

@@ -25,7 +25,7 @@ use std::{io, path::Path};
 /// ```
 pub async fn write(path: impl AsRef<Path>, contents: impl AsRef<[u8]>) -> io::Result<()> {
     let path = path.as_ref().to_owned();
-    let contents = contents.as_ref().to_owned();
+    let contents = crate::util::as_ref::upgrade(contents);
 
     asyncify(move || std::fs::write(path, contents)).await
 }

tokio/src/util/as_ref.rs

@@ -0,0 +1,38 @@
+use super::typeid;
+
+#[derive(Debug)]
+pub(crate) enum OwnedBuf {
+    Vec(Vec<u8>),
+    #[cfg(feature = "io-util")]
+    Bytes(bytes::Bytes),
+}
+
+impl AsRef<[u8]> for OwnedBuf {
+    fn as_ref(&self) -> &[u8] {
+        match self {
+            Self::Vec(vec) => vec,
+            #[cfg(feature = "io-util")]
+            Self::Bytes(bytes) => bytes,
+        }
+    }
+}
+
+pub(crate) fn upgrade<B: AsRef<[u8]>>(buf: B) -> OwnedBuf {
+    let buf = match unsafe { typeid::try_transmute::<B, Vec<u8>>(buf) } {
+        Ok(vec) => return OwnedBuf::Vec(vec),
+        Err(original_buf) => original_buf,
+    };
+
+    let buf = match unsafe { typeid::try_transmute::<B, String>(buf) } {
+        Ok(string) => return OwnedBuf::Vec(string.into_bytes()),
+        Err(original_buf) => original_buf,
+    };
+
+    #[cfg(feature = "io-util")]
+    let buf = match unsafe { typeid::try_transmute::<B, bytes::Bytes>(buf) } {
+        Ok(bytes) => return OwnedBuf::Bytes(bytes),
+        Err(original_buf) => original_buf,
+    };
+
+    OwnedBuf::Vec(buf.as_ref().to_owned())
+}

tokio/src/util/mod.rs

@@ -2,6 +2,9 @@ cfg_io_driver! {
     pub(crate) mod bit;
 }
 
+#[cfg(feature = "fs")]
+pub(crate) mod as_ref;
+
 #[cfg(feature = "rt")]
 pub(crate) mod atomic_cell;
 
@@ -81,6 +84,9 @@ cfg_rt_multi_thread! {
 
 pub(crate) mod trace;
 
+#[cfg(feature = "fs")]
+pub(crate) mod typeid;
+
 pub(crate) mod error;
 
 #[cfg(feature = "io-util")]

tokio/src/util/typeid.rs

@@ -0,0 +1,44 @@
+use std::{
+    any::TypeId,
+    marker::PhantomData,
+    mem::{self, ManuallyDrop},
+};
+
+// SAFETY: this function does not compare lifetimes. Values returned as `Ok`
+// may have their lifetimes extended.
+pub(super) unsafe fn try_transmute<Src, Target: 'static>(x: Src) -> Result<Target, Src> {
+    if nonstatic_typeid::<Src>() == TypeId::of::<Target>() {
+        let x = ManuallyDrop::new(x);
+        Ok(mem::transmute_copy::<Src, Target>(&x))
+    } else {
+        Err(x)
+    }
+}
+
+// https://github.com/dtolnay/typeid/blob/b06a3c08a0eaccc7df6091ade1ae4e3fb53609d5/src/lib.rs#L197-L222
+#[inline(always)]
+fn nonstatic_typeid<T>() -> TypeId
+where
+    T: ?Sized,
+{
+    trait NonStaticAny {
+        fn get_type_id(&self) -> TypeId
+        where
+            Self: 'static;
+    }
+
+    impl<T: ?Sized> NonStaticAny for PhantomData<T> {
+        #[inline(always)]
+        fn get_type_id(&self) -> TypeId
+        where
+            Self: 'static,
+        {
+            TypeId::of::<T>()
+        }
+    }
+
+    let phantom_data = PhantomData::<T>;
+    NonStaticAny::get_type_id(unsafe {
+        mem::transmute::<&dyn NonStaticAny, &(dyn NonStaticAny + 'static)>(&phantom_data)
+    })
+}