improve clean_interned

Author: Dunqing

crates/oxc_formatter/examples/formatter.rs

@@ -54,15 +54,15 @@ fn main() -> Result<(), String> {
 
     let formatter = Formatter::new(&allocator, options);
     let formatted = formatter.format(&ret.program);
-    // if show_ir {
-    //     println!("--- IR ---");
-    //     println!("{}", &formatted.document().to_string());
-    //     println!("--- End IR ---\n");
-    // }
+    if show_ir {
+        println!("--- IR ---");
+        println!("{}", &formatted.document().to_string());
+        println!("--- End IR ---\n");
+    }
 
-    // println!("--- Formatted Code ---");
+    println!("--- Formatted Code ---");
     let code = formatted.print().map_err(|e| e.to_string())?.into_code();
-    // println!("{code}");
-    // println!("--- End Formatted Code ---");
+    println!("{code}");
+    println!("--- End Formatted Code ---");
     Ok(())
 }

crates/oxc_formatter/src/formatter/buffer.rs

@@ -496,7 +496,7 @@ impl<'buf, 'ast> RemoveSoftLinesBuffer<'buf, 'ast> {
     }
 
     /// Removes the soft line breaks from an interned element.
-    fn clean_interned(&mut self, interned: &Interned<'ast>) -> Interned<'ast> {
+    fn clean_interned(&mut self, interned: Interned<'ast>) -> Interned<'ast> {
         clean_interned(
             interned,
             &mut self.interned_cache,
@@ -517,12 +517,12 @@ impl<'buf, 'ast> RemoveSoftLinesBuffer<'buf, 'ast> {
 
 // Extracted to function to avoid monomorphization
 fn clean_interned<'ast>(
-    interned: &Interned<'ast>,
+    interned: Interned<'ast>,
     interned_cache: &mut FxHashMap<Interned<'ast>, Interned<'ast>>,
     condition_content_stack: &mut Vec<Condition>,
     allocator: &'ast Allocator,
 ) -> Interned<'ast> {
-    if let Some(cleaned) = interned_cache.get(interned) {
+    if let Some(cleaned) = interned_cache.get(&interned) {
         cleaned.clone()
     } else {
         // Find the first soft line break element, interned element, or conditional expanded
@@ -531,21 +531,22 @@ fn clean_interned<'ast>(
             FormatElement::Line(LineMode::Soft | LineMode::SoftOrSpace)
             | FormatElement::Tag(Tag::StartConditionalContent(_) | Tag::EndConditionalContent)
             | FormatElement::BestFitting(_) => {
-                // TODO: avoid using clone here
-
                 let mut cleaned =
                     ArenaVec::from_iter_in(interned[..index].iter().cloned(), allocator);
                 Some((cleaned, &interned[index..]))
             }
             FormatElement::Interned(inner) => {
-                let cleaned_inner =
-                    clean_interned(inner, interned_cache, condition_content_stack, allocator);
+                let cleaned_inner = clean_interned(
+                    inner.clone(),
+                    interned_cache,
+                    condition_content_stack,
+                    allocator,
+                );
 
                 if &cleaned_inner == inner {
                     None
                 } else {
                     let mut cleaned = ArenaVec::with_capacity_in(interned.len(), allocator);
-                    // TODO: avoid using clone here
                     cleaned.extend(interned[..index].iter().cloned());
                     cleaned.push(FormatElement::Interned(cleaned_inner));
                     Some((cleaned, &interned[index + 1..]))
@@ -580,7 +581,7 @@ fn clean_interned<'ast>(
 
                         FormatElement::Interned(interned) => {
                             cleaned.push(FormatElement::Interned(clean_interned(
-                                interned,
+                                interned.clone(),
                                 interned_cache,
                                 condition_content_stack,
                                 allocator,
@@ -596,21 +597,20 @@ fn clean_interned<'ast>(
                     }
                 }
 
-                Interned::new(ArenaVec::from_iter_in(cleaned, allocator))
+                Interned::new(cleaned)
             }
             // No change necessary, return existing interned element
             None => interned.clone(),
         };
 
-        interned_cache.insert(interned.clone(), result.clone());
+        interned_cache.insert(interned, result.clone());
         result
     }
 }
 
 impl<'ast> Buffer<'ast> for RemoveSoftLinesBuffer<'_, 'ast> {
     fn write_element(&mut self, element: FormatElement<'ast>) -> FormatResult<()> {
-        let mut element_stack = Vec::new();
-        element_stack.push(element);
+        let mut element_stack = Vec::from_iter([element]);
         while let Some(element) = element_stack.pop() {
             match element {
                 FormatElement::Tag(Tag::StartConditionalContent(condition)) => {
@@ -628,14 +628,14 @@ impl<'ast> Buffer<'ast> for RemoveSoftLinesBuffer<'_, 'ast> {
                     self.inner.write_element(FormatElement::Space)?;
                 }
                 FormatElement::Interned(interned) => {
-                    let cleaned = self.clean_interned(&interned);
+                    let cleaned = self.clean_interned(interned);
                     self.inner.write_element(FormatElement::Interned(cleaned))?;
                 }
                 // Since this buffer aims to simulate infinite print width, we don't need to retain the best fitting.
                 // Just extract the flattest variant and then handle elements within it.
                 FormatElement::BestFitting(best_fitting) => {
                     let most_flat = best_fitting.most_flat();
-                    most_flat.iter().rev().for_each(|element| element_stack.push(element.clone()));
+                    element_stack.extend(most_flat.iter().rev().cloned());
                 }
                 element => self.inner.write_element(element)?,
             }

crates/oxc_formatter/src/formatter/builders.rs

@@ -2550,7 +2550,7 @@ impl<'ast> Format<'ast> for BestFitting<'_, 'ast> {
             buffer.write_fmt(Arguments::from(variant))?;
             buffer.write_element(FormatElement::Tag(EndEntry))?;
 
-            formatted_variants.push(buffer.take_vec().into_boxed_slice());
+            formatted_variants.push(buffer.take_vec().into_bump_slice());
         }
 
         let formatted_variants =

crates/oxc_formatter/src/formatter/format_element/document.rs

@@ -19,9 +19,9 @@ use crate::{
 use crate::{format, write};
 
 /// A formatted document.
-#[derive(Debug)]
+#[derive(Debug, Clone, Eq, PartialEq, Default)]
 pub struct Document<'a> {
-    elements: ArenaVec<'a, FormatElement<'a>>,
+    elements: &'a [FormatElement<'a>],
 }
 
 impl Document<'_> {
@@ -142,24 +142,24 @@ impl Document<'_> {
 
 impl<'a> From<ArenaVec<'a, FormatElement<'a>>> for Document<'a> {
     fn from(elements: ArenaVec<'a, FormatElement<'a>>) -> Self {
-        Self { elements }
+        Self { elements: elements.into_bump_slice() }
     }
 }
 
 impl<'a> Deref for Document<'a> {
     type Target = [FormatElement<'a>];
 
     fn deref(&self) -> &Self::Target {
-        self.elements.as_slice()
+        self.elements
     }
 }
 
 impl std::fmt::Display for Document<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let allocator = Allocator::default();
         let context = FormatContext::dummy(&allocator);
-        let formatted = format!(context, [self.elements.as_slice()])
-            .expect("Formatting not to throw any FormatErrors");
+        let formatted =
+            format!(context, [self.elements]).expect("Formatting not to throw any FormatErrors");
 
         f.write_str(formatted.print().expect("Expected a valid document").as_code())
     }

crates/oxc_formatter/src/formatter/format_element/mod.rs

@@ -1,12 +1,11 @@
 pub mod document;
 pub mod tag;
 
-use std::cell::Cell;
 // #[cfg(target_pointer_width = "64")]
 // use biome_rowan::static_assert;
 use std::hash::{Hash, Hasher};
-use std::mem::{self, ManuallyDrop};
 use std::num::NonZeroU32;
+use std::ptr;
 use std::{borrow::Cow, ops::Deref, rc::Rc};
 
 use unicode_width::{UnicodeWidthChar, UnicodeWidthStr};
@@ -53,7 +52,7 @@ const _: () = {
 /// Language agnostic IR for formatting source code.
 ///
 /// Use the helper functions like [crate::builders::space], [crate::builders::soft_line_break] etc. defined in this file to create elements.
-#[derive(Clone)]
+#[derive(Clone, Eq, PartialEq)]
 pub enum FormatElement<'a> {
     /// A space token, see [crate::builders::space] for documentation.
     Space,
@@ -150,29 +149,25 @@ impl PrintMode {
 }
 
 #[derive(Clone)]
-pub struct Interned<'a>(ManuallyDrop<Rc<ArenaVec<'a, FormatElement<'a>>>>);
+pub struct Interned<'a>(&'a [FormatElement<'a>]);
 
 impl<'a> Interned<'a> {
     pub(super) fn new(content: ArenaVec<'a, FormatElement<'a>>) -> Self {
-        Self(ManuallyDrop::new(Rc::new(content)))
+        Self(content.into_bump_slice())
     }
 }
 
 impl PartialEq for Interned<'_> {
     fn eq(&self, other: &Interned<'_>) -> bool {
-        std::ptr::eq(&raw const self.0, &raw const other.0)
+        ptr::eq(self.0, other.0)
     }
 }
 
 impl Eq for Interned<'_> {}
 
 impl Hash for Interned<'_> {
-    fn hash<H>(&self, hasher: &mut H)
-    where
-        H: Hasher,
-    {
-        Address::from_ptr(&raw const self.0).hash(hasher);
-        // Rc::as_ptr(&self.0).hash(hasher);
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.0.as_ptr().addr().hash(state);
     }
 }
 
@@ -186,7 +181,7 @@ impl<'a> Deref for Interned<'a> {
     type Target = [FormatElement<'a>];
 
     fn deref(&self) -> &Self::Target {
-        &self.0
+        self.0
     }
 }
 
@@ -305,19 +300,12 @@ impl FormatElements for FormatElement<'_> {
 /// can pick the best fitting variant.
 ///
 /// Best fitting is defined as the variant that takes the most horizontal space but fits on the line.
+#[derive(Clone, Eq, PartialEq)]
 pub struct BestFittingElement<'a> {
     /// The different variants for this element.
     /// The first element is the one that takes up the most space horizontally (the most flat),
     /// The last element takes up the least space horizontally (but most horizontal space).
-    variants: ArenaBox<'a, [ArenaBox<'a, [FormatElement<'a>]>]>,
-}
-
-impl Clone for BestFittingElement<'_> {
-    fn clone(&self) -> Self {
-        // SAFETY: NOT sure if it is safe, but needed to implement Clone.
-        // TODO: Once everything works, this should be replaced with a proper clone implementation.
-        Self { variants: unsafe { mem::transmute_copy(self) } }
-    }
+    variants: &'a [&'a [FormatElement<'a>]],
 }
 
 impl<'a> BestFittingElement<'a> {
@@ -330,15 +318,13 @@ impl<'a> BestFittingElement<'a> {
     /// ## Safety
     /// The slice must contain at least two variants.
     #[doc(hidden)]
-    pub unsafe fn from_vec_unchecked(
-        variants: ArenaVec<'a, ArenaBox<'a, [FormatElement<'a>]>>,
-    ) -> Self {
+    pub unsafe fn from_vec_unchecked(variants: ArenaVec<'a, &'a [FormatElement<'a>]>) -> Self {
         debug_assert!(
             variants.len() >= 2,
             "Requires at least the least expanded and most expanded variants"
         );
 
-        Self { variants: variants.into_boxed_slice() }
+        Self { variants: variants.into_bump_slice() }
     }
 
     /// Returns the most expanded variant
@@ -348,8 +334,8 @@ impl<'a> BestFittingElement<'a> {
         )
     }
 
-    pub fn variants(&self) -> &[ArenaBox<'a, [FormatElement<'a>]>] {
-        &self.variants
+    pub fn variants(&self) -> &[&'a [FormatElement<'a>]] {
+        self.variants
     }
 
     /// Returns the least expanded variant
@@ -362,7 +348,7 @@ impl<'a> BestFittingElement<'a> {
 
 impl std::fmt::Debug for BestFittingElement<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_list().entries(&*self.variants).finish()
+        f.debug_list().entries(self.variants).finish()
     }
 }
 

crates/oxc_formatter/src/formatter/formatter.rs

@@ -27,7 +27,7 @@ impl<'buf, 'ast> Formatter<'buf, 'ast> {
         Self { buffer }
     }
 
-    pub fn allocator(&self) -> &Allocator {
+    pub fn allocator(&self) -> &'ast Allocator {
         self.context().allocator()
     }
 

crates/oxc_formatter/src/formatter/mod.rs

@@ -73,7 +73,7 @@ pub use self::{
 };
 use self::{format_element::document::Document, group_id::UniqueGroupIdBuilder, prelude::TagKind};
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Formatted<'a> {
     document: Document<'a>,
     context: FormatContext<'a>,

crates/oxc_formatter/src/formatter/printer/mod.rs

@@ -357,8 +357,7 @@ impl<'a> Printer<'a> {
 
                 // Remove the content slice because printing needs the variant WITH the start entry
                 let popped_slice = queue.pop_slice();
-                // TODO: We should revisit here when everything works.
-                // debug_assert_eq!(popped_slice, Some(content));
+                debug_assert_eq!(popped_slice, Some(content));
 
                 if variant_fits {
                     queue.extend_back(variant);
@@ -512,7 +511,7 @@ impl<'a> Printer<'a> {
             }
         }
 
-        if queue.top().is_some_and(|e| matches!(e, FormatElement::Tag(Tag::EndFill))) {
+        if queue.top() == Some(&FormatElement::Tag(EndFill)) {
             Ok(())
         } else {
             invalid_end_tag(TagKind::Fill, stack.top_kind())

crates/oxc_formatter/src/ir_transform/sort_imports/mod.rs

@@ -2,9 +2,10 @@ mod import_unit;
 mod partitioned_chunk;
 mod source_line;
 
-use oxc_allocator::{Allocator, Vec as ArenaVec};
 use std::mem;
 
+use oxc_allocator::{Allocator, Vec as ArenaVec};
+
 use crate::{
     formatter::format_element::{FormatElement, LineMode, document::Document},
     options,
@@ -32,9 +33,7 @@ impl SortImportsTransform {
     pub fn transform<'a>(&self, document: &Document<'a>, allocator: &'a Allocator) -> Document<'a> {
         // Early return for empty files
         if document.len() == 1 && matches!(document[0], FormatElement::Line(LineMode::Hard)) {
-            // SAFETY: Not sure how to solve it.
-            // TODO: Refactor `Document` to avoid this.
-            return unsafe { mem::transmute_copy(document) };
+            return document.clone();
         }
 
         let prev_elements: &[FormatElement<'a>] = document;