feature(rome_formatter): Add inspect to Memoized

This PR adds the `inspect` method to `Memoized` that allows inspecting the formatted content.

This can be useful in scenarios where some formatting must determine if some content will break to decide the best formatting.

This is different from #2771 because the buffer should be preferred in cases where the fact whatever some content breaks is only used to determine the formatting of another element because it avoids interning/memoizing the inner content.

The API shown here should be preferred when the formatting of the element itself depends on whatever its content breaks where using the `WillBreakBuffer` would need to rewind to the previous state.

crates/rome_formatter/src/format_element.rs

@@ -553,12 +553,11 @@ impl FormatElement {
 
     /// Interns a format element.
     ///
-    /// Returns `self` for an empty list AND an already interned elements.
-    pub fn intern(self) -> FormatElement {
+    /// Returns `self` for an already interned element.
+    pub fn intern(self) -> Interned {
         match self {
-            FormatElement::List(list) if list.is_empty() => list.into(),
-            element @ FormatElement::Interned(_) => element,
-            element => FormatElement::Interned(Interned(Rc::new(element))),
+            FormatElement::Interned(interned) => interned,
+            element => Interned(Rc::new(element)),
         }
     }
 }

crates/rome_formatter/src/format_extensions.rs

@@ -1,6 +1,7 @@
 use crate::prelude::*;
 use std::cell::RefCell;
 use std::marker::PhantomData;
+use std::ops::Deref;
 
 use crate::{write, Buffer, VecBuffer};
 
@@ -144,7 +145,7 @@ impl<T, Context> MemoizeFormat<Context> for T where T: Format<Context> {}
 #[derive(Debug)]
 pub struct Memoized<F, Context> {
     inner: F,
-    memory: RefCell<Option<FormatResult<FormatElement>>>,
+    memory: RefCell<Option<FormatResult<Interned>>>,
     options: PhantomData<Context>,
 }
 
@@ -159,43 +160,117 @@ where
             options: PhantomData,
         }
     }
-}
 
-impl<F, Context> Format<Context> for Memoized<F, Context>
-where
-    F: Format<Context>,
-{
-    fn fmt(&self, f: &mut Formatter<Context>) -> FormatResult<()> {
-        // Cached
-        if let Some(memory) = self.memory.borrow().as_ref() {
-            return match memory {
-                Ok(elements) => {
-                    f.write_element(elements.clone())?;
-
-                    Ok(())
-                }
-                Err(err) => Err(*err),
-            };
+    /// Gives access to the memoized content.
+    ///
+    /// Performs the formatting if the content hasn't been formatted at this point.
+    ///
+    /// # Example
+    ///
+    /// Inspect if some memoized content breaks.
+    ///
+    /// ```rust
+    /// use std::cell::Cell;
+    /// use rome_formatter::{format, write};
+    /// use rome_formatter::prelude::*;
+    /// use rome_rowan::TextSize;
+    ///
+    /// #[derive(Default)]
+    /// struct Counter {
+    ///   value: Cell<u64>
+    /// }
+    ///
+    /// impl Format<SimpleFormatContext> for Counter {fn fmt(&self, f: &mut Formatter<SimpleFormatContext>) -> FormatResult<()> {
+    ///         let current = self.value.get();
+    ///
+    ///         write!(f, [
+    ///             token("Count:"),
+    ///             space_token(),
+    ///             dynamic_token(&std::format!("{current}"), TextSize::default()),
+    ///             hard_line_break()
+    ///         ])?;
+    ///         self.value.set(current + 1);
+    ///         Ok(())
+    ///     }
+    /// }
+    ///
+    /// let content = format_with(|f| {
+    ///     let mut counter = Counter::default().memoized();
+    ///     let counter_content = counter.inspect(f)?;
+    ///
+    ///     if counter_content.will_break() {
+    ///         write!(f, [token("Counter:"), block_indent(&counter)])
+    ///     } else {
+    ///         write!(f, [token("Counter:"), counter])
+    ///     }?;
+    ///
+    ///     write!(f, [counter])
+    /// });
+    ///
+    ///
+    /// let formatted = format!(SimpleFormatContext::default(), [content]).unwrap();
+    /// assert_eq!("Counter:\n\tCount: 0\nCount: 0\n", formatted.print().as_code())
+    ///
+    /// ```
+    pub fn inspect(&mut self, f: &mut Formatter<Context>) -> FormatResult<&FormatElement> {
+        let cache = self.memory.get_mut();
+        if cache.is_none() {
+            Self::format_and_store(cache, &self.inner, f)?;
+        }
+
+        // SAFETY: Safe because `format_and_store` populates the cache
+        match cache.as_ref().unwrap() {
+            Ok(content) => Ok(content.deref()),
+            Err(error) => Err(*error),
         }
+    }
+
+    /// Formats the inner content and stores it in the provided store.
+    ///
+    /// Guarantees that the cache is populated after this call.
+    fn format_and_store(
+        cache: &mut Option<FormatResult<Interned>>,
+        inner: &F,
+        f: &mut Formatter<Context>,
+    ) -> FormatResult<()> {
         let mut buffer = VecBuffer::new(f.state_mut());
 
-        let result = write!(buffer, [self.inner]);
+        let result = write!(buffer, [inner]);
 
         match result {
             Ok(_) => {
                 let elements = buffer.into_element();
                 let interned = elements.intern();
 
-                f.write_element(interned.clone())?;
-
-                *self.memory.borrow_mut() = Some(Ok(interned));
+                *cache = Some(Ok(interned));
 
                 Ok(())
             }
             Err(err) => {
-                *self.memory.borrow_mut() = Some(Err(err));
+                *cache = Some(Err(err));
                 Err(err)
             }
         }
     }
 }
+
+impl<F, Context> Format<Context> for Memoized<F, Context>
+where
+    F: Format<Context>,
+{
+    fn fmt(&self, f: &mut Formatter<Context>) -> FormatResult<()> {
+        if self.memory.borrow().is_none() {
+            Self::format_and_store(&mut self.memory.borrow_mut(), &self.inner, f)?;
+        }
+
+        // SAFETY: format_and_store guarantees that content is stored in memory
+        return match self.memory.borrow().as_ref().unwrap() {
+            Ok(elements) => {
+                f.write_element(FormatElement::Interned(elements.clone()))?;
+
+                Ok(())
+            }
+            Err(err) => Err(*err),
+        };
+    }
+}