bitty_write_macro 0.1.2

// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Size-optimized versions of the standard library's `write!` and `writeln!` macros
//!
//! ## Why use these macros?
//!
//! In code size-constrained systems that perform string formatting,
//! using the `write!` or `writeln!` macro can [generate larger code][fmt_opt]
//! than writing the contents directly, even after linking and inlining.
//!
//! The [`write!`] and [`writeln!`] macros provided by this crate detect when
//! the write could be optimized as a direct write with `write_str`
//! (for `fmt::Write`) or `write` (for `io::Write`) and calls that instead of
//! `format_args!` and `write_fmt`.
//!
//! ## The `std` feature
//!
//! By default, this crate is `#![no_std]` and only supports [`fmt::Write`].
//! To write to [`io::Write`] sinks, enable the `std` feature.
//!
//! ## Known drawbacks
//!
//! These macros:
//!
//! - Target a `write_str` method on the trait, using an internal adapter to expose
//!   `write_str` on `io::Write`.
//! - Do not currently accept `concat!` in the format string.
//! - Do not optimize nested `format_args!` like `write!(w, "{}", format_args!("x"))`.
//!
//! [`fmt::Arguments`]: core::fmt::Arguments
//! [`fmt::Write`]: core::fmt::Write
//! [`io::Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
//! [fmt_opt]: https://godbolt.org/z/qaPa9WqdP

#![cfg_attr(not(test), no_std)]
#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]

pub use format_args_conditional::branch_on_format_capture;

/// An optimizing version of [`core::write!`] that calls `write_str`/`write` when possible.
///
/// This is mostly a drop-in replacement for the version in `core`.
/// Because this shares a name with an item in the prelude, it must be imported explicitly
/// and not with `use bitty_write_macro::fmt::*`.
///
/// This optimization is done on a best-effort basis, and may improve over time.
#[macro_export]
macro_rules! write {
    (@lit $w:expr, $lit:literal) => {{
        #[allow(unused_imports)]
        use $crate::__internal_unstable::IoWriteStr;
        $w.write_str($lit)
    }};
    (@fmt $w:expr, $($fmt:tt)+) => {
        $crate::__internal_unstable::core_write!($w, $($fmt)+)
    };
    ($w:expr $(,)?) => { Ok(()) };
    ($w:expr, $arg:tt $(,)?) => {
        $crate::branch_on_format_capture!(
            $crate::write {@fmt $w,},
            $crate::write {@lit $w,},
            $arg
        )
    };
    ($w:expr, $($arg:tt)+) => {
        $crate::__internal_unstable::core_write!($w, $($arg)+)
    };
}

/// An optimizing version of [`core::writeln!`] that calls `write_str`/`write` when possible.
///
/// This is mostly a drop-in replacement for the version in `core`.
/// Because this shares a name with an item in the prelude, it must be imported explicitly
/// and not with `use bitty_write_macro::*`.
///
/// This optimization is done on a best-effort basis, and may improve over time.
#[macro_export]
macro_rules! writeln {
    (@lit $w:expr, $lit:literal) => {{
        #[allow(unused_imports)]
        use $crate::__internal_unstable::IoWriteStr;
        $w.write_str(concat!($lit, "\n"))
    }};
    (@fmt $w:expr, $($fmt:tt)+) => {
        $crate::__internal_unstable::core_writeln!($w, $($fmt)+)
    };
    ($w:expr $(,)?) => {{
        #[allow(unused_imports)]
        use $crate::__internal_unstable::IoWriteStr;
        $w.write_str("\n")
    }};
    ($w:expr, $arg:tt $(,)?) => {
        $crate::branch_on_format_capture!(
            $crate::writeln {@fmt $w,},
            $crate::writeln {@lit $w,},
            $arg
        )
    };
    ($w:expr, $($arg:tt)+) => {
        $crate::__internal_unstable::core_writeln!($w, $($arg)+)
    }
}

// Items in this module may be broken across compatible versions.
#[doc(hidden)]
pub mod __internal_unstable {
    pub use core::write as core_write;
    pub use core::writeln as core_writeln;

    #[cfg(feature = "std")]
    mod std_only {
        extern crate std;
        use std::io;

        /// Used to have a consistent method name to target in the macro.
        pub trait IoWriteStr {
            // Same return value as `std::io::Write::write_fmt`
            fn write_str(&mut self, x: &str) -> io::Result<()>;
        }

        impl<T: io::Write + ?Sized> IoWriteStr for T {
            #[inline(always)]
            fn write_str(&mut self, x: &str) -> io::Result<()> {
                self.write(x.as_bytes()).map(|_| ())
            }
        }
    }

    #[cfg(feature = "std")]
    pub use std_only::IoWriteStr;

    #[cfg(not(feature = "std"))]
    pub trait IoWriteStr {}
}

#[cfg(test)]
mod tests {
    use core::fmt;
    use std::io;

    use fmt::Write as _;

    #[cfg(feature = "std")]
    use io::Write as _;

    use super::{write, writeln};

    #[derive(Default)]
    pub struct FormatTrack<T> {
        writer: T,
        used_write_fmt: bool,
    }

    impl<T: fmt::Write> fmt::Write for FormatTrack<T> {
        fn write_str(&mut self, s: &str) -> fmt::Result {
            self.writer.write_str(s)
        }

        fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result {
            self.used_write_fmt = true;
            fmt::write(&mut self.writer, args)
        }
    }

    impl<T: io::Write> io::Write for FormatTrack<T> {
        fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> io::Result<()> {
            self.used_write_fmt = true;
            io::Write::write_fmt(&mut self.writer, args)
        }

        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            self.writer.write(buf)
        }

        fn flush(&mut self) -> io::Result<()> {
            self.writer.flush()
        }
    }

    macro_rules! check {
        ($macro:ident, $use_write_fmt:ident, $expected_contents:literal, $writer:ident $($args:tt)*) => {
            assert!($macro!($writer $($args)*).is_ok());
            assert_eq!($writer.used_write_fmt, $use_write_fmt);
            $writer.used_write_fmt = false;
            assert_eq!(AsRef::<[u8]>::as_ref(&$writer.writer), $expected_contents.as_bytes());
            $writer.writer.clear();
        };
    }

    #[test]
    fn test_temporary_writer() {
        let x = 10;

        let _ = write!(String::new(), "Hello");
        let _ = write!(String::new(), "Hello {x}");
        let _ = write!(String::new(), "Hello {x}", x = "world");

        let _ = writeln!(String::new(), "Hello");
        let _ = writeln!(String::new(), "Hello {x}");
        let _ = writeln!(String::new(), "Hello {x}", x = "world");
    }

    #[test]
    fn test_generic_owned_fmt_writer() {
        fn generic_owned_fmt<W: fmt::Write>(mut w: W) -> W {
            let x = 10;
            let _ = write!(w, "A ");
            let _ = write!(w, "B {x} ");
            let _ = write!(w, "C {x} ", x = "D");

            let _ = writeln!(w, "A");
            let _ = writeln!(w, "B {x}");
            let _ = writeln!(w, "C {x}", x = "D");
            w
        }

        assert_eq!(
            generic_owned_fmt(String::new()),
            "A B 10 C D A\nB 10\nC D\n"
        );
    }

    #[test]
    fn test_generic_borrow_of_owned_fmt_writer() {
        fn generic_borrow_of_owned_fmt<W: fmt::Write>(mut w: W) -> W {
            let x = 10;
            let _ = write!(&mut w, "A ");
            let _ = write!(&mut w, "B {x} ");
            let _ = write!(&mut w, "C {x} ", x = "D");

            let _ = writeln!(&mut w, "A");
            let _ = writeln!(&mut w, "B {x}");
            let _ = writeln!(&mut w, "C {x}", x = "D");
            w
        }

        assert_eq!(
            generic_borrow_of_owned_fmt(String::new()),
            "A B 10 C D A\nB 10\nC D\n"
        );
    }

    #[test]
    fn test_generic_mut_fmt_writer() {
        // Note that the `w` binding itself is not `mut`.
        fn generic_mut_fmt<W: fmt::Write>(w: &mut W) -> &mut W {
            let x = 10;
            let _ = write!(w, "A ");
            let _ = write!(w, "B {x} ");
            let _ = write!(w, "C {x} ", x = "D");

            let _ = writeln!(w, "A");
            let _ = writeln!(w, "B {x}");
            let _ = writeln!(w, "C {x}", x = "D");
            w
        }

        assert_eq!(
            generic_mut_fmt(&mut String::new()),
            "A B 10 C D A\nB 10\nC D\n"
        );
    }

    #[test]
    fn test_fmt_write() {
        let mut w = FormatTrack::<String>::default();

        let x = 10;
        check!(write, true, "Hello world", w, "Hello {x}", x = "world");
        check!(
            writeln,
            true,
            "Hello newline\n",
            w,
            "Hello {x}",
            x = "newline"
        );
        check!(write, true, "Capture x as 10", w, "Capture x as {x}");
        check!(writeln, true, "Capture x as 10\n", w, "Capture x as {x}");
        check!(write, true, "Positional 10", w, "Positional {}", x);
        check!(writeln, true, "Positional 10\n", w, "Positional {}", x);

        check!(write, false, "Just a string", w, "Just a string");
        check!(write, false, "", w, "");
        check!(
            writeln,
            false,
            "Just a string\nwith a newline\n",
            w,
            "Just a string\nwith a newline"
        );
        check!(writeln, false, "\n", w, "",);
        check!(writeln, false, "\n", w, "");
        check!(writeln, false, "\n", w,);
        check!(writeln, false, "\n", w);
    }

    #[cfg(feature = "std")]
    #[test]
    fn test_io_write() {
        let mut w = FormatTrack::<Vec<u8>>::default();

        let x = 10;
        check!(write, true, "Hello world", w, "Hello {x}", x = "world");
        check!(
            writeln,
            true,
            "Hello newline\n",
            w,
            "Hello {x}",
            x = "newline"
        );
        check!(write, true, "Capture x as 10", w, "Capture x as {x}");
        check!(writeln, true, "Capture x as 10\n", w, "Capture x as {x}");
        check!(write, true, "Positional 10", w, "Positional {}", x);
        check!(writeln, true, "Positional 10\n", w, "Positional {}", x);

        check!(write, false, "Just a string", w, "Just a string");
        check!(write, false, "", w, "");
        check!(
            writeln,
            false,
            "Just a string\nwith a newline\n",
            w,
            "Just a string\nwith a newline"
        );
        check!(writeln, false, "\n", w, "",);
        check!(writeln, false, "\n", w, "");
        check!(writeln, false, "\n", w,);
        check!(writeln, false, "\n", w);
    }
}