mz_ore/
fmt.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Formatting utilities.
17
18use std::{fmt, io};
19
20#[cfg(feature = "network")]
21use bytes::{BufMut, BytesMut};
22
23/// A trait for objects that can be infallibly written to.
24///
25/// Like [`std::fmt::Write`], except that the methods do not return errors.
26/// Implementations are provided for [`String`], `bytes::BytesMut`, and
27/// `Vec<u8>`, as writing to these types cannot fail.
28///
29/// Objects that implement `FormatBuffer` can be passed to the [`write!`] macro:
30///
31/// ```
32/// use mz_ore::fmt::FormatBuffer;
33///
34/// let mut buf = String::new();
35/// write!(buf, "{:.02}", 1.0 / 7.0);
36/// assert_eq!(buf, "0.14");
37/// ```
38///
39/// The trait is particularly useful for functions that need to generically
40/// write to either a [`String`] or a byte buffer:
41///
42/// ```
43/// use mz_ore::fmt::FormatBuffer;
44///
45/// fn write_timezone_offset<F>(buf: &mut F, mut offset_seconds: i32)
46/// where
47///     F: FormatBuffer,
48/// {
49///     if offset_seconds >= 0 {
50///         buf.write_char('+');
51///     } else {
52///         buf.write_char('-');
53///     }
54///     let offset_seconds = offset_seconds.abs();
55///     write!(buf, "{:02}:{:02}", offset_seconds / 60 / 60, offset_seconds / 60 % 60);
56/// }
57///
58/// let offset_seconds = -18000;
59///
60/// let mut s = String::new();
61/// write_timezone_offset(&mut s, offset_seconds);
62///
63/// let mut v = Vec::new();
64/// write_timezone_offset(&mut v, offset_seconds);
65///
66/// assert_eq!(s, "-05:00");
67/// assert_eq!(v, s.as_bytes());
68/// ```
69///
70/// The implementations of `FormatBuffer` for `Vec<u8>` and `BytesMut` are
71/// guaranteed to only write valid UTF-8 bytes into the underlying buffer.
72pub trait FormatBuffer: AsRef<[u8]> {
73    /// Glue for usage of the [`write!`] macro with implementors of this trait.
74    ///
75    /// This method should not be invoked manually, but rather through the
76    /// `write!` macro itself.
77    fn write_fmt(&mut self, fmt: fmt::Arguments);
78
79    /// Writes a [`char`] into this buffer.
80    fn write_char(&mut self, c: char);
81
82    /// Writes a string into this buffer.
83    fn write_str(&mut self, s: &str);
84
85    /// Returns the number of bytes in the buffer.
86    fn len(&self) -> usize;
87
88    /// Reports whether the buffer is empty.
89    fn is_empty(&self) -> bool {
90        self.len() == 0
91    }
92
93    /// Returns a mutable reference to the bytes in the buffer.
94    ///
95    /// # Safety
96    ///
97    /// If the byte slice was valid UTF-8, it must remain valid UTF-8 when the
98    // returned mutable reference is dropped. `FormatBuffer`s may be
99    /// [`String`]s or other data types whose memory safety depends upon only
100    /// containing valid UTF-8.
101    unsafe fn as_bytes_mut(&mut self) -> &mut [u8];
102}
103
104impl FormatBuffer for String {
105    fn write_fmt(&mut self, fmt: fmt::Arguments) {
106        fmt::Write::write_fmt(self, fmt).expect("fmt::Write::write_fmt cannot fail on a String");
107    }
108
109    fn write_char(&mut self, c: char) {
110        self.push(c)
111    }
112
113    fn write_str(&mut self, s: &str) {
114        self.push_str(s)
115    }
116
117    fn len(&self) -> usize {
118        self.len()
119    }
120
121    unsafe fn as_bytes_mut(&mut self) -> &mut [u8] {
122        unsafe { str::as_bytes_mut(self) }
123    }
124}
125
126impl FormatBuffer for Vec<u8> {
127    fn write_fmt(&mut self, fmt: fmt::Arguments) {
128        io::Write::write_fmt(self, fmt).expect("io::Write::write_fmt cannot fail on Vec<u8>")
129    }
130
131    fn write_char(&mut self, c: char) {
132        self.extend(c.encode_utf8(&mut [0; 4]).as_bytes())
133    }
134
135    fn write_str(&mut self, s: &str) {
136        self.extend(s.as_bytes())
137    }
138
139    fn len(&self) -> usize {
140        self.len()
141    }
142
143    unsafe fn as_bytes_mut(&mut self) -> &mut [u8] {
144        self
145    }
146}
147
148#[cfg(feature = "network")]
149impl FormatBuffer for BytesMut {
150    fn write_fmt(&mut self, fmt: fmt::Arguments) {
151        io::Write::write_fmt(&mut self.writer(), fmt)
152            .expect("io::Write::write_fmt cannot fail on BytesMut")
153    }
154
155    fn write_char(&mut self, c: char) {
156        self.put(c.encode_utf8(&mut [0; 4]).as_bytes())
157    }
158
159    fn write_str(&mut self, s: &str) {
160        self.put(s.as_bytes())
161    }
162
163    fn len(&self) -> usize {
164        self.len()
165    }
166
167    unsafe fn as_bytes_mut(&mut self) -> &mut [u8] {
168        self
169    }
170}