Skip to main content

azathoth_utils/
formatter.rs

1use super::errors::{AzUtilResult, AzUtilErrorCode};
2use alloc::string::String;
3use alloc::vec::Vec;
4use core::str;
5
6/// Minimal writable buffer trait
7pub trait WriteBuffer {
8    /// Writes a `&str` into the buffer
9    fn write_str(&mut self, s: &str) -> AzUtilResult<()>;
10}
11
12/// Minimal string type
13///
14/// This type exists to make sure there are no surprises during the format operation
15pub struct AllocString {
16    buf: Vec<u8>,
17}
18
19impl AllocString {
20    /// Creates a new `AllocString`
21    pub fn new() -> AllocString {
22        Self { buf: Vec::new() }
23    }
24
25    /// Creates a new [`AllocString`] with a specified capacity
26    pub fn with_capacity(capacity: usize) -> AllocString {
27        Self {
28            buf: Vec::with_capacity(capacity),
29        }
30    }
31
32    /// Converts the [`AllocString`] into an [`String`] type
33    pub fn into_string(self) -> AzUtilResult<String> {
34        String::from_utf8(self.buf).map_err(|_| AzUtilErrorCode::FormatError)
35    }
36
37    /// Extend the current [`AllocString`] with an additional `&str`
38    pub fn push_str(&mut self, s: &str) -> AzUtilResult<()> {
39        self.buf.extend_from_slice(s.as_bytes());
40        Ok(())
41    }
42}
43
44impl WriteBuffer for AllocString {
45    fn write_str(&mut self, s: &str) -> AzUtilResult<()> {
46        self.push_str(s)
47    }
48}
49
50impl WriteBuffer for String {
51    fn write_str(&mut self, s: &str) -> AzUtilResult<()> {
52        self.extend(s.chars());
53        Ok(())
54    }
55}
56
57impl WriteBuffer for Vec<u8> {
58    fn write_str(&mut self, s: &str) -> AzUtilResult<()> {
59        self.extend_from_slice(s.as_bytes());
60        Ok(())
61    }
62}
63
64/// Format specifier struct
65///
66/// Used to track the format specifiers in a string
67#[derive(Default, Debug)]
68pub struct FormatSpec {
69    alternate: bool,
70    specifier: char,
71}
72
73impl FormatSpec {
74    /// Parses a string for extended format specifiers
75    ///
76    /// This function may be extended in the future, but for now it only searches for `:` and `#` chars
77    pub fn parse_spec(s: &str) -> Self {
78        let mut spec = FormatSpec::default();
79        if s.is_empty() {
80            return spec;
81        }
82        let mut chars = s.chars();
83        if s.starts_with(':') {
84            chars.next();
85        }
86        if chars.as_str().starts_with('#') {
87            spec.alternate = true;
88            chars.next();
89        }
90        if let Some(c) = chars.last() {
91            spec.specifier = c;
92        }
93        spec
94    }
95}
96/// Custom formatter replacement for the [`core::fmt::Display`] trait
97pub trait FDisplay {
98    /// The caller must implement this function to use the [`crate::format_str_inner!`] macro
99    fn fmt<W: WriteBuffer>(&self, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()>;
100}
101
102fn u64_to_str_decimal<W: WriteBuffer>(mut n: u64, buf: &mut W) -> AzUtilResult<()> {
103    if n == 0 {
104        return buf.write_str("0");
105    }
106    let mut temp_buf = [0u8; 20];
107    let mut i = 0;
108    while n > 0 {
109        temp_buf[i] = (n % 10) as u8 + b'0';
110        n /= 10;
111        i += 1;
112    }
113    let digits = &mut temp_buf[..i];
114    digits.reverse();
115    buf.write_str(unsafe { str::from_utf8_unchecked(digits) })
116}
117
118fn u64_to_str_radix<W: WriteBuffer>(
119    mut n: u64,
120    radix: u32,
121    uppercase: bool,
122    buf: &mut W,
123) -> AzUtilResult<()> {
124    if n == 0 {
125        return buf.write_str("0");
126    }
127    let mut temp_buf = [0u8; 64];
128    let mut i = 0;
129    let charset = if uppercase {
130        b"0123456789ABCDEF"
131    } else {
132        b"0123456789abcdef"
133    };
134    while n > 0 {
135        temp_buf[i] = charset[(n % (radix as u64)) as usize];
136        n /= radix as u64;
137        i += 1;
138    }
139    let digits = &mut temp_buf[..i];
140    digits.reverse();
141    buf.write_str(unsafe { str::from_utf8_unchecked(digits) })
142}
143
144impl<'a, T: ?Sized + FDisplay> FDisplay for &'a T {
145    fn fmt<W: WriteBuffer>(&self, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()> {
146        (*self).fmt(w, spec)
147    }
148}
149impl FDisplay for str {
150    fn fmt<W: WriteBuffer>(&self, w: &mut W, _spec: &FormatSpec) -> AzUtilResult<()> {
151        w.write_str(self)
152    }
153}
154impl FDisplay for String {
155    fn fmt<W: WriteBuffer>(&self, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()> {
156        self.as_str().fmt(w, spec)
157    }
158}
159impl FDisplay for char {
160    fn fmt<W: WriteBuffer>(&self, w: &mut W, _spec: &FormatSpec) -> AzUtilResult<()> {
161        let mut buffer = [0u8; 4];
162        w.write_str(self.encode_utf8(&mut buffer))
163    }
164}
165impl FDisplay for bool {
166    fn fmt<W: WriteBuffer>(&self, w: &mut W, _spec: &FormatSpec) -> AzUtilResult<()> {
167        w.write_str(if *self { "true" } else { "false" })
168    }
169}
170impl<T: FDisplay> FDisplay for Option<T> {
171    fn fmt<W: WriteBuffer>(&self, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()> {
172        match self {
173            Some(v) => v.fmt(w, spec),
174            None => w.write_str("None"),
175        }
176    }
177}
178impl<T: FDisplay, E: FDisplay> FDisplay for Result<T, E> {
179    fn fmt<W: WriteBuffer>(&self, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()> {
180        match self {
181            Ok(v) => v.fmt(w, spec),
182            Err(e) => {
183                w.write_str("Err(")?;
184                e.fmt(w, spec)?;
185                w.write_str(")")
186            }
187        }
188    }
189}
190
191impl<T: FDisplay> FDisplay for Vec<T> {
192    fn fmt<W: WriteBuffer>(&self, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()> {
193        w.write_str("[")?;
194        for (i, item) in self.iter().enumerate() {
195            if i > 0 {
196                w.write_str(", ")?;
197            }
198            item.fmt(w, spec)?;
199        }
200        w.write_str("]")
201    }
202}
203
204impl<T> FDisplay for *const T {
205    fn fmt<W: WriteBuffer>(&self, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()> {
206        if spec.specifier != 'p' && spec.specifier != '\0' {
207            return Err(AzUtilErrorCode::FormatError);
208        }
209        w.write_str("0x")?;
210        u64_to_str_radix(*self as usize as u64, 16, false, w)
211    }
212}
213impl<T> FDisplay for *mut T {
214    fn fmt<W: WriteBuffer>(&self, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()> {
215        (*self as *const T).fmt(w, spec)
216    }
217}
218
219macro_rules! impl_display_uint {
220    ($($t:ty),*) => {
221        $(impl FDisplay for $t { fn fmt<W: WriteBuffer>(&self, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()> {
222            let val = *self as u64;
223            fmt_spec(val, w, spec)
224        } })*
225    };
226}
227macro_rules! impl_display_int {
228    ($($t:ty),*) => {
229        $(impl FDisplay for $t { fn fmt<W: WriteBuffer>(&self, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()> {
230            let is_negative = *self < 0;
231            let val = self.unsigned_abs() as u64;
232            if is_negative { w.write_str("-")?; }
233            fmt_spec(val, w, spec)
234        } })*
235    };
236}
237impl_display_uint!(u8, u16, u32, u64, u128, usize);
238impl_display_int!(i8, i16, i32, i64, i128, isize);
239
240fn fmt_spec<W: WriteBuffer>(val: u64, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()> {
241    match spec.specifier {
242        'x' | 'X' => {
243            if spec.alternate {
244                w.write_str("0x")?;
245            }
246            u64_to_str_radix(val, 16, spec.specifier == 'X', w)
247        }
248        'b' => {
249            if spec.alternate {
250                w.write_str("0b")?;
251            }
252            u64_to_str_radix(val, 2, false, w)
253        }
254        _ => u64_to_str_decimal(val, w),
255    }
256}
257
258/// Argument formatting trait
259///
260/// The trait requires each implementor to implement the [`FormatArgs::format_at`] function
261pub trait FormatArgs {
262    /// Argument formatting function
263    ///
264    /// Expects a writable buffer that implements the [`WriteBuffer`] trait, an index for specifying where to place the formatted value,
265    /// and a [`FormatSpec`] object to allow the usage of the additional format specifiers
266    fn format_at<W: WriteBuffer>(
267        &self,
268        index: usize,
269        w: &mut W,
270        spec: &FormatSpec,
271    ) -> AzUtilResult<()>;
272}
273impl FormatArgs for () {
274    fn format_at<W: WriteBuffer>(
275        &self,
276        _index: usize,
277        _w: &mut W,
278        _spec: &FormatSpec,
279    ) -> AzUtilResult<()> {
280        Err(AzUtilErrorCode::FormatError)
281    }
282}
283macro_rules! impl_format_args {
284    ($($T:ident, $idx:tt),+) => {
285        impl<$($T: FDisplay),+> FormatArgs for ($($T),+,) {
286            #[allow(non_snake_case)]
287            fn format_at<W: WriteBuffer>(&self, index: usize, w: &mut W, spec: &FormatSpec) -> AzUtilResult<()> {
288                match index { $($idx => self.$idx.fmt(w, spec),)+ _ => Err(AzUtilErrorCode::ParseError) }
289            }
290        }
291    };
292}
293impl_format_args!(T0, 0);
294impl_format_args!(T0, 0, T1, 1);
295impl_format_args!(T0, 0, T1, 1, T2, 2);
296impl_format_args!(T0, 0, T1, 1, T2, 2, T3, 3);
297impl_format_args!(T0, 0, T1, 1, T2, 2, T3, 3, T4, 4);
298impl_format_args!(T0, 0, T1, 1, T2, 2, T3, 3, T4, 4, T5, 5);
299
300/// String formatting function
301///
302/// Accepts a mutable buffer that implements the [`WriteBuffer`] trait, a format string, and any type of argument that implements the [`FormatArgs`] trait
303/// Writes the formatted value to the buffer.
304pub fn format_rt<W, A>(buf: &mut W, fmt: &str, args: &A) -> AzUtilResult<()>
305where
306    W: WriteBuffer,
307    A: FormatArgs,
308{
309    let mut arg_idx = 0;
310    let mut parts = fmt.split('{');
311    if let Some(part) = parts.next() {
312        buf.write_str(part)?;
313    }
314
315    for part in parts {
316        if part.starts_with('{') {
317            buf.write_str("{")?;
318            buf.write_str(&part[1..])?;
319            continue;
320        }
321
322        if let Some(end_brace_idx) = part.find('}') {
323            let spec_str = &part[..end_brace_idx];
324            let spec = FormatSpec::parse_spec(spec_str);
325
326            args.format_at(arg_idx, buf, &spec)?;
327            arg_idx += 1;
328
329            buf.write_str(&part[end_brace_idx + 1..])?;
330        } else {
331            if !part.is_empty() {
332                return Err(AzUtilErrorCode::FormatError);
333            }
334        }
335    }
336
337    Ok(())
338}
339
340/// Wrapper around the [`format_rt`] function to simplify the [`crate::format_str_inner!`] macro definition
341pub fn format_str_inner<A: FormatArgs>(fmt: &str, args: &A) -> String {
342    let mut buffer = AllocString::new();
343    match format_rt(&mut buffer, fmt, args) {
344        Ok(()) => buffer.into_string().expect(""),
345        Err(e) => panic!("Failed to format value: {:?}", e),
346    }
347}