From 5a11b814d4cd84d43f7b96a81f3e1595187c26a1 Mon Sep 17 00:00:00 2001
From: Markus Reiter <me@reitermark.us>
Date: Tue, 16 Aug 2022 14:42:00 +0200
Subject: [PATCH] Add `IpDisplayBuffer` helper struct.

---
 library/std/src/net/ip.rs | 77 ++++++++++++++++++++++++---------------
 1 file changed, 47 insertions(+), 30 deletions(-)

diff --git a/library/std/src/net/ip.rs b/library/std/src/net/ip.rs
index 41ca9ba8425..f4fba7f7386 100644
--- a/library/std/src/net/ip.rs
+++ b/library/std/src/net/ip.rs
@@ -3,12 +3,44 @@
 mod tests;
 
 use crate::cmp::Ordering;
-use crate::fmt::{self, Write as FmtWrite};
-use crate::io::Write as IoWrite;
+use crate::fmt::{self, Write};
 use crate::mem::transmute;
+use crate::str;
 use crate::sys::net::netc as c;
 use crate::sys_common::{FromInner, IntoInner};
 
+/// Used for slow path in `Display` implementations when alignment is required.
+struct IpDisplayBuffer<const SIZE: usize> {
+    buf: [u8; SIZE],
+    len: usize,
+}
+
+impl<const SIZE: usize> IpDisplayBuffer<SIZE> {
+    #[inline(always)]
+    pub const fn new(_ip: &[u8; SIZE]) -> Self {
+        Self { buf: [0; SIZE], len: 0 }
+    }
+
+    #[inline(always)]
+    pub fn as_str(&self) -> &str {
+        // SAFETY: `buf` is only written to by the `fmt::Write::write_str` implementation
+        // which writes a valid UTF-8 string to `buf` and correctly sets `len`.
+        unsafe { str::from_utf8_unchecked(&self.buf[..self.len]) }
+    }
+}
+
+impl<const SIZE: usize> fmt::Write for IpDisplayBuffer<SIZE> {
+    fn write_str(&mut self, s: &str) -> fmt::Result {
+        if let Some(buf) = self.buf.get_mut(self.len..(self.len + s.len())) {
+            buf.copy_from_slice(s.as_bytes());
+            self.len += s.len();
+            Ok(())
+        } else {
+            Err(fmt::Error)
+        }
+    }
+}
+
 /// An IP address, either IPv4 or IPv6.
 ///
 /// This enum can contain either an [`Ipv4Addr`] or an [`Ipv6Addr`], see their
@@ -991,21 +1023,17 @@ fn from(ipv6: Ipv6Addr) -> IpAddr {
 impl fmt::Display for Ipv4Addr {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
         let octets = self.octets();
-        // Fast Path: if there's no alignment stuff, write directly to the buffer
+
+        // If there are no alignment requirements, write the IP address directly to `f`.
+        // Otherwise, write it to a local buffer and then use `f.pad`.
         if fmt.precision().is_none() && fmt.width().is_none() {
             write!(fmt, "{}.{}.{}.{}", octets[0], octets[1], octets[2], octets[3])
         } else {
-            const IPV4_BUF_LEN: usize = 15; // Long enough for the longest possible IPv4 address
-            let mut buf = [0u8; IPV4_BUF_LEN];
-            let mut buf_slice = &mut buf[..];
+            let mut buf = IpDisplayBuffer::new(b"255.255.255.255");
+            // Buffer is long enough for the longest possible IPv4 address, so this should never fail.
+            write!(buf, "{}.{}.{}.{}", octets[0], octets[1], octets[2], octets[3]).unwrap();
 
-            // Note: The call to write should never fail, hence the unwrap
-            write!(buf_slice, "{}.{}.{}.{}", octets[0], octets[1], octets[2], octets[3]).unwrap();
-            let len = IPV4_BUF_LEN - buf_slice.len();
-
-            // This unsafe is OK because we know what is being written to the buffer
-            let buf = unsafe { crate::str::from_utf8_unchecked(&buf[..len]) };
-            fmt.pad(buf)
+            fmt.pad(buf.as_str())
         }
     }
 }
@@ -1708,8 +1736,8 @@ pub const fn to_canonical(&self) -> IpAddr {
 #[stable(feature = "rust1", since = "1.0.0")]
 impl fmt::Display for Ipv6Addr {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        // If there are no alignment requirements, write out the IP address to
-        // f. Otherwise, write it to a local buffer, then use f.pad.
+        // If there are no alignment requirements, write the IP address directly to `f`.
+        // Otherwise, write it to a local buffer and then use `f.pad`.
         if f.precision().is_none() && f.width().is_none() {
             let segments = self.segments();
 
@@ -1780,22 +1808,11 @@ fn fmt_subslice(f: &mut fmt::Formatter<'_>, chunk: &[u16]) -> fmt::Result {
                 }
             }
         } else {
-            // Slow path: write the address to a local buffer, then use f.pad.
-            // Defined recursively by using the fast path to write to the
-            // buffer.
+            let mut buf = IpDisplayBuffer::new(b"ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff");
+            // Buffer is long enough for the longest possible IPv6 address, so this should never fail.
+            write!(buf, "{}", self).unwrap();
 
-            // This is the largest possible size of an IPv6 address
-            const IPV6_BUF_LEN: usize = (4 * 8) + 7;
-            let mut buf = [0u8; IPV6_BUF_LEN];
-            let mut buf_slice = &mut buf[..];
-
-            // Note: This call to write should never fail, so unwrap is okay.
-            write!(buf_slice, "{}", self).unwrap();
-            let len = IPV6_BUF_LEN - buf_slice.len();
-
-            // This is safe because we know exactly what can be in this buffer
-            let buf = unsafe { crate::str::from_utf8_unchecked(&buf[..len]) };
-            f.pad(buf)
+            f.pad(buf.as_str())
         }
     }
 }