[PLATFORM-860]: Veil memory optimisations and refactoring (#60)

* Internal code refactor + optimisations

* Format

* Compile time error if a specialisation isn't handled

Co-authored-by: MaeIsBad <26093674+MaeIsBad@users.noreply.github.com>

* `str` -> `to_redact` more descriptive name

* Use an enum for much cleaner handling of redaction length modifier

Co-authored-by: William Venner <14863743+WilliamVenner@users.noreply.github.com>
Co-authored-by: MaeIsBad <26093674+MaeIsBad@users.noreply.github.com>

src/private.rs

@@ -1,47 +1,44 @@
-use std::fmt::{Debug, Display};
+use std::{
+    fmt::{Debug, Display, Write},
+    num::NonZeroU8,
+};
 
-#[repr(transparent)]
-pub struct DisplayDebug(String);
-impl std::fmt::Debug for DisplayDebug {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.write_str(self.0.as_str())
-    }
-}
-impl AsRef<str> for DisplayDebug {
-    #[inline(always)]
-    fn as_ref(&self) -> &str {
-        self.0.as_str()
-    }
-}
-
-pub struct RedactFlags {
+pub enum RedactSpecialization {
     /// Whether the type we're redacting is an Option<T> or not. Poor man's specialization! This is detected
     /// by the proc macro reading the path to the type, so it's not perfect.
     ///
     /// This could be improved & rid of in a number of different ways in the future:
     ///
-    /// * Once specialization is stabilized, we can use a trait to override redacting behaviour for some types,
+    /// * Once specialization is stabilized, we can use a trait to override redacting behavior for some types,
     /// one of which would be Option<T>.
     ///
     /// * Once std::ptr::metadata and friends are stabilized, we could use it to unsafely cast the dyn Debug pointer
     /// to a concrete Option<T> and redact it directly. Probably not the best idea.
     ///
     /// * Once trait upcasting is stabilized, we could use it to upcast the dyn Debug pointer to a dyn Any and then
     /// downcast it to a concrete Option<T> and redact it directly.
-    pub is_option: bool,
+    Option,
+}
 
-    /// Whether to only partially redact the data.
-    ///
-    /// Incompatible with `fixed`.
-    pub partial: bool,
+#[derive(Clone, Copy)]
+pub enum RedactionLength {
+    /// Redact the entire data.
+    Full,
 
-    /// What character to use for redacting.
-    pub redact_char: char,
+    /// Redact a portion of the data.
+    Partial,
 
     /// Whether to redact with a fixed width, ignoring the length of the data.
-    ///
-    /// Incompatible with `partial`.
-    pub fixed: u8,
+    Fixed(NonZeroU8),
+}
+
+#[derive(Clone, Copy)]
+pub struct RedactFlags {
+    /// How much of the data to redact.
+    pub redact_length: RedactionLength,
+
+    /// What character to use for redacting.
+    pub redact_char: char,
 }
 impl RedactFlags {
     /// How many characters must a word be for it to be partially redacted?
@@ -52,14 +49,14 @@ impl RedactFlags {
     /// Maximum number of characters to expose at the beginning and end of a partial redact.
     const MAX_PARTIAL_EXPOSE: usize = 3;
 
-    fn redact_partial(&self, str: &str, redacted: &mut String) {
-        let count = str.chars().filter(|char| char.is_alphanumeric()).count();
+    pub(crate) fn redact_partial(&self, fmt: &mut std::fmt::Formatter, to_redact: &str) -> std::fmt::Result {
+        let count = to_redact.chars().filter(|char| char.is_alphanumeric()).count();
         if count < Self::MIN_PARTIAL_CHARS {
-            for char in str.chars() {
+            for char in to_redact.chars() {
                 if char.is_alphanumeric() {
-                    redacted.push(self.redact_char);
+                    fmt.write_char(self.redact_char)?;
                 } else {
-                    redacted.push(char);
+                    fmt.write_char(char)?;
                 }
             }
         } else {
@@ -68,104 +65,122 @@ impl RedactFlags {
 
             let mut prefix_gas = redact_count;
             let mut middle_gas = count - redact_count - redact_count;
-            for char in str.chars() {
+            for char in to_redact.chars() {
                 if char.is_alphanumeric() {
                     if prefix_gas > 0 {
                         prefix_gas -= 1;
-                        redacted.push(char);
+                        fmt.write_char(char)?;
                     } else if middle_gas > 0 {
                         middle_gas -= 1;
-                        redacted.push(self.redact_char);
+                        fmt.write_char(self.redact_char)?;
                     } else {
-                        redacted.push(char);
+                        fmt.write_char(char)?;
                     }
                 } else {
-                    redacted.push(char);
+                    fmt.write_char(char)?;
                 }
             }
         }
+        Ok(())
     }
 
-    fn redact_full(&self, str: &str, redacted: &mut String) {
-        for char in str.chars() {
+    pub(crate) fn redact_full(&self, fmt: &mut std::fmt::Formatter, to_redact: &str) -> std::fmt::Result {
+        for char in to_redact.chars() {
             if char.is_whitespace() || !char.is_alphanumeric() {
-                redacted.push(char);
+                fmt.write_char(char)?;
             } else {
-                redacted.push(self.redact_char);
+                fmt.write_char(self.redact_char)?;
             }
         }
+        Ok(())
     }
 
-    fn redact_fixed(&self, width: usize, redacted: &mut String) {
-        redacted.reserve_exact(width);
+    pub(crate) fn redact_fixed(fmt: &mut std::fmt::Formatter, width: usize, char: char) -> std::fmt::Result {
+        let mut buf = String::with_capacity(width);
         for _ in 0..width {
-            redacted.push(self.redact_char);
+            buf.push(char);
         }
+        fmt.write_str(&buf)
     }
 }
 
 pub enum RedactionTarget<'a> {
-    /// Redact the output of the type's [`std::fmt::Debug`] implementation.
+    /// Redact the output of the type's [`Debug`] implementation.
     Debug {
         this: &'a dyn Debug,
 
         /// Sourced from [`std::fmt::Formatter::alternate`]
         alternate: bool,
     },
 
-    /// Redact the output of the type's [`std::fmt::Display`] implementation.
+    /// Redact the output of the type's [`Display`] implementation.
     Display(&'a dyn Display),
 }
-
-pub fn redact(this: RedactionTarget, flags: RedactFlags) -> DisplayDebug {
-    let mut redacted = String::new();
-
-    let to_redactable_string = || match this {
-        RedactionTarget::Debug { this, alternate: false } => format!("{:?}", this),
-        RedactionTarget::Debug { this, alternate: true } => format!("{:#?}", this),
-        RedactionTarget::Display(this) => this.to_string(),
-    };
-
+impl RedactionTarget<'_> {
+    /// Pass through directly to the formatter.
     #[cfg(feature = "toggle")]
-    if crate::toggle::get_redaction_behavior().is_plaintext() {
-        return DisplayDebug(to_redactable_string());
+    pub(crate) fn passthrough(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            RedactionTarget::Debug { this, .. } => std::fmt::Debug::fmt(this, fmt),
+            RedactionTarget::Display(this) => std::fmt::Display::fmt(this, fmt),
+        }
     }
+}
+impl ToString for RedactionTarget<'_> {
+    fn to_string(&self) -> String {
+        match self {
+            RedactionTarget::Debug { this, alternate: false } => format!("{:?}", this),
+            RedactionTarget::Debug { this, alternate: true } => format!("{:#?}", this),
+            RedactionTarget::Display(this) => this.to_string(),
+        }
+    }
+}
 
-    (|| {
-        if flags.fixed > 0 {
-            flags.redact_fixed(flags.fixed as usize, &mut redacted);
-            return;
+pub struct RedactionFormatter<'a> {
+    pub this: RedactionTarget<'a>,
+    pub flags: RedactFlags,
+    pub specialization: Option<RedactSpecialization>,
+}
+impl std::fmt::Debug for RedactionFormatter<'_> {
+    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        #[cfg(feature = "toggle")]
+        if crate::toggle::get_redaction_behavior().is_plaintext() {
+            return self.this.passthrough(fmt);
         }
 
-        let redactable_string = to_redactable_string();
-
-        redacted.reserve(redactable_string.len());
-
-        // Specialize for Option<T>
-        if flags.is_option {
-            if redactable_string == "None" {
-                // We don't need to do any redacting
-                // https://prima.slack.com/archives/C03URH9N43U/p1661423554871499
-            } else if let Some(inner) = redactable_string
-                .strip_prefix("Some(")
-                .and_then(|inner| inner.strip_suffix(')'))
-            {
-                redacted.push_str("Some(");
-                flags.redact_partial(inner, &mut redacted);
-                redacted.push(')');
-            } else {
-                // This should never happen, but just in case...
-                flags.redact_full(&redactable_string, &mut redacted);
+        if let RedactionLength::Fixed(n) = &self.flags.redact_length {
+            return RedactFlags::redact_fixed(fmt, n.get() as usize, self.flags.redact_char);
+        }
+
+        let redactable_string = self.this.to_string();
+
+        #[allow(clippy::single_match)]
+        match self.specialization {
+            Some(RedactSpecialization::Option) => {
+                if redactable_string == "None" {
+                    // We don't need to do any redacting
+                    // https://prima.slack.com/archives/C03URH9N43U/p1661423554871499
+                    return fmt.write_str("None");
+                } else if let Some(inner) = redactable_string
+                    .strip_prefix("Some(")
+                    .and_then(|inner| inner.strip_suffix(')'))
+                {
+                    fmt.write_str("Some(")?;
+                    self.flags.redact_partial(fmt, inner)?;
+                    return fmt.write_char(')');
+                } else {
+                    // This should never happen, but just in case...
+                    return self.flags.redact_full(fmt, &redactable_string);
+                }
             }
-            return;
+
+            None => {}
         }
 
-        if flags.partial {
-            flags.redact_partial(&redactable_string, &mut redacted);
+        if let RedactionLength::Partial = &self.flags.redact_length {
+            self.flags.redact_partial(fmt, &redactable_string)
         } else {
-            flags.redact_full(&redactable_string, &mut redacted);
+            self.flags.redact_full(fmt, &redactable_string)
         }
-    })();
-
-    DisplayDebug(redacted)
+    }
 }

veil-macros/src/enums.rs

@@ -1,7 +1,7 @@
 use crate::{
-    flags::FieldFlags,
+    flags::{ExtractFlags, FieldFlags, FieldFlagsParse},
     fmt::{self, FormatData},
-    UnusedDiagnostic,
+    redact::UnusedDiagnostic,
 };
 use proc_macro::TokenStream;
 use quote::ToTokens;
@@ -21,7 +21,7 @@ pub(super) fn derive_redact(
     unused: &mut UnusedDiagnostic,
 ) -> Result<TokenStream, syn::Error> {
     // Parse #[redact(all, variant, ...)] from the enum attributes, if present.
-    let top_level_flags = match FieldFlags::extract::<1>(&attrs, false)? {
+    let top_level_flags = match FieldFlags::extract::<1>("Redact", &attrs, FieldFlagsParse { skip_allowed: false })? {
         [Some(flags)] => {
             if !flags.all || !flags.variant {
                 return Err(syn::Error::new(
@@ -41,7 +41,13 @@ pub(super) fn derive_redact(
     // Collect each variant's flags
     let mut variant_flags = Vec::with_capacity(e.variants.len());
     for variant in &e.variants {
-        let mut flags = match FieldFlags::extract::<2>(&variant.attrs, top_level_flags.is_some())? {
+        let mut flags = match FieldFlags::extract::<2>(
+            "Redact",
+            &variant.attrs,
+            FieldFlagsParse {
+                skip_allowed: top_level_flags.is_some(),
+            },
+        )? {
             [None, None] => EnumVariantFieldFlags::default(),
 
             [Some(flags), None] => {
@@ -163,7 +169,17 @@ pub(super) fn derive_redact(
         // Variant name redacting
         let variant_name = variant.ident.to_string();
         let variant_name = if let Some(flags) = &flags.variant_flags {
-            fmt::generate_redact_call(quote! { &#variant_name }, false, flags, unused)
+            // The variant name must always be formatted with the Display impl.
+            let flags = FieldFlags {
+                display: true,
+                ..*flags
+            };
+
+            // Generate the RedactionFormatter expression for the variant name
+            let redact = fmt::generate_redact_call(quote! { &#variant_name }, false, &flags, unused);
+
+            // Because the other side is expecting a &str, we need to convert the RedactionFormatter to a String (and then to a &str)
+            quote! { format!("{:?}", #redact).as_str() }
         } else {
             variant_name.into_token_stream()
         };
@@ -182,7 +198,7 @@ pub(super) fn derive_redact(
                         "unit structs do not need redacting as they contain no data",
                     ));
                 } else {
-                    quote! { write!(f, "{:?}", #variant_name)? }
+                    quote! { fmt.write_str(#variant_name)? }
                 }
             }
         });
@@ -191,9 +207,9 @@ pub(super) fn derive_redact(
     let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
     Ok(quote! {
         impl #impl_generics ::std::fmt::Debug for #name_ident #ty_generics #where_clause {
-            fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
+            fn fmt(&self, fmt: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
                 #[allow(unused)] // Suppresses unused warning with `#[redact(display)]`
-                let alternate = f.alternate();
+                let alternate = fmt.alternate();
 
                 match self {
                     #(Self::#variant_idents #variant_destructures => { #variant_bodies; },)*