refactor(formatter,oxfmt): Remove EmbeddedIR and use FormatElement directly

apps/oxfmt/src/core/external_formatter.rs

@@ -276,7 +276,7 @@ impl ExternalFormatter {
         let embedded_doc_callback: Option<EmbeddedDocFormatterCallback> = if needs_embedded {
             let format_embedded_doc = Arc::clone(&self.format_embedded_doc);
             let options_for_doc = options.clone();
-            Some(Arc::new(move |language: &str, texts: &[&str]| {
+            Some(Arc::new(move |allocator, group_id_builder, language: &str, texts: &[&str]| {
                 let Some(parser_name) = language_to_prettier_parser(language) else {
                     return Err(format!("Unsupported language: {language}"));
                 };
@@ -302,7 +302,7 @@ impl ExternalFormatter {
                                     serde_json::from_str(&doc_json_str).map_err(|err| {
                                         format!("Failed to parse Doc JSON: {err}")
                                     })?;
-                                from_prettier_doc::doc_json_to_embedded_ir(&doc_json)
+                                from_prettier_doc::to_format_elements_for_template(&doc_json, allocator, group_id_builder)
                             })
                             .collect()
                     })

crates/oxc_formatter/src/external_formatter.rs

@@ -1,6 +1,8 @@
 use std::sync::Arc;
 
-use super::formatter::format_element::{LineMode, PrintMode};
+use oxc_allocator::Allocator;
+
+use super::formatter::{FormatElement, group_id::UniqueGroupIdBuilder};
 
 /// Callback function type for formatting embedded code.
 /// Takes (tag_name, code) and returns formatted code or an error.
@@ -9,13 +11,24 @@ pub type EmbeddedFormatterCallback =
 
 /// Callback function type for formatting embedded code via Doc in batch.
 ///
-/// Takes (tag_name, texts) and returns one `Vec<EmbeddedIR>` per input.
-/// Used for the Doc→IR path (e.g., `printToDoc` → Doc JSON → `EmbeddedIR`).
+/// Takes (allocator, group_id_builder, tag_name, texts) and returns one `Vec<FormatElement<'a>>` per input.
+/// Used for the Doc→IR path (e.g., `JS:printToDoc()` → Doc JSON → `Rust:FormatElement`).
+///
+/// The `&Allocator` allows the callback to allocate arena strings for `FormatElement::Text`.
+/// The `&UniqueGroupIdBuilder` allows the callback to create `GroupId`s for group/conditional constructs.
 ///
 /// For GraphQL, each quasi is a separate text (`texts.len() == quasis.len()`).
 /// For CSS/HTML, quasis are joined with placeholders into a single text (`texts.len() == 1`).
-pub type EmbeddedDocFormatterCallback =
-    Arc<dyn Fn(&str, &[&str]) -> Result<Vec<Vec<EmbeddedIR>>, String> + Send + Sync>;
+pub type EmbeddedDocFormatterCallback = Arc<
+    dyn for<'a> Fn(
+            &'a Allocator,
+            &UniqueGroupIdBuilder,
+            &str,
+            &[&str],
+        ) -> Result<Vec<Vec<FormatElement<'a>>>, String>
+        + Send
+        + Sync,
+>;
 
 /// Callback function type for sorting Tailwind CSS classes.
 /// Takes classes and returns the sorted versions.
@@ -79,23 +92,29 @@ impl ExternalCallbacks {
 
     /// Format embedded code as Doc in batch.
     ///
-    /// Takes multiple texts and returns one `Vec<EmbeddedIR>` per input text.
+    /// Takes multiple texts and returns one `Vec<FormatElement<'a>>` per input text.
     /// The caller is responsible for interleaving the results with JS expressions.
     ///
     /// # Arguments
+    /// * `allocator` - The arena allocator for allocating strings in `FormatElement::Text`
+    /// * `group_id_builder` - Builder for creating unique `GroupId`s
     /// * `tag_name` - The template tag (e.g., "css", "gql", "html")
     /// * `texts` - The code texts to format (multiple quasis for GraphQL, single joined text for CSS/HTML)
     ///
     /// # Returns
-    /// * `Some(Ok(Vec<Vec<EmbeddedIR>>))` - The formatted code as a vector of `EmbeddedIR` for each input text
+    /// * `Some(Ok(Vec<Vec<FormatElement<'a>>>))` - The formatted code as FormatElements for each input text
     /// * `Some(Err(String))` - An error message if formatting failed
     /// * `None` - No embedded formatter callback is set
-    pub fn format_embedded_doc(
+    pub fn format_embedded_doc<'a>(
         &self,
+        allocator: &'a Allocator,
+        group_id_builder: &UniqueGroupIdBuilder,
         tag_name: &str,
         texts: &[&str],
-    ) -> Option<Result<Vec<Vec<EmbeddedIR>>, String>> {
-        self.embedded_doc_formatter.as_ref().map(|cb| cb(tag_name, texts))
+    ) -> Option<Result<Vec<Vec<FormatElement<'a>>>, String>> {
+        self.embedded_doc_formatter
+            .as_ref()
+            .map(|cb| cb(allocator, group_id_builder, tag_name, texts))
     }
 
     /// Sort Tailwind CSS classes.
@@ -116,58 +135,3 @@ impl ExternalCallbacks {
         }
     }
 }
-
-// ---
-
-/// Owned intermediate IR for embedded language formatting.
-///
-/// This type bridges the callback boundary between `apps/oxfmt` (or other callers) and `oxc_formatter`.
-/// Unlike `FormatElement<'a>`, it has no lifetime parameter and owns all its data,
-/// so it can be returned from `Arc<dyn Fn>` callbacks.
-///
-/// The `oxc_formatter` side converts `EmbeddedIR` → `FormatElement<'a>` using the allocator.
-#[derive(Debug, Clone)]
-pub enum EmbeddedIR {
-    Space,
-    HardSpace,
-    Line(LineMode),
-    ExpandParent,
-    /// Owned string (unlike `FormatElement::Text` which borrows from the arena).
-    Text(String),
-    LineSuffixBoundary,
-    // --- Tag equivalents (all fields pub, no lifetime) ---
-    StartIndent,
-    EndIndent,
-    /// Positive integer only. Converted to `Tag::StartAlign(Align(NonZeroU8))`.
-    StartAlign(u8),
-    EndAlign,
-    /// - `to_root: false` → `DedentMode::Level`
-    /// - `to_root: true` → `DedentMode::Root`
-    StartDedent {
-        to_root: bool,
-    },
-    EndDedent {
-        to_root: bool,
-    },
-    /// `id` is a numeric group ID (mapped to `GroupId` via `HashMap<u32, GroupId>`).
-    StartGroup {
-        id: Option<u32>,
-        should_break: bool,
-    },
-    EndGroup,
-    /// `mode` = Break or Flat, `group_id` references a group by numeric ID.
-    StartConditionalContent {
-        mode: PrintMode,
-        group_id: Option<u32>,
-    },
-    EndConditionalContent,
-    /// GroupId is mandatory (matches `Tag::StartIndentIfGroupBreaks(GroupId)`).
-    StartIndentIfGroupBreaks(u32),
-    EndIndentIfGroupBreaks(u32),
-    StartFill,
-    EndFill,
-    StartEntry,
-    EndEntry,
-    StartLineSuffix,
-    EndLineSuffix,
-}

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

@@ -158,11 +158,13 @@ impl Group {
         Self { id: None, mode: Cell::new(GroupMode::Flat) }
     }
 
+    #[must_use]
     pub fn with_id(mut self, id: Option<GroupId>) -> Self {
         self.id = id;
         self
     }
 
+    #[must_use]
     pub fn with_mode(mut self, mode: GroupMode) -> Self {
         self.mode = Cell::new(mode);
         self
@@ -208,6 +210,7 @@ impl Condition {
         Self { mode, group_id: None }
     }
 
+    #[must_use]
     pub fn with_group_id(mut self, id: Option<GroupId>) -> Self {
         self.group_id = id;
         self
@@ -226,6 +229,10 @@ impl Condition {
 pub struct Align(pub(crate) NonZeroU8);
 
 impl Align {
+    pub fn new(count: NonZeroU8) -> Self {
+        Self(count)
+    }
+
     pub fn count(&self) -> NonZeroU8 {
         self.0
     }

crates/oxc_formatter/src/formatter/formatter.rs

@@ -67,6 +67,12 @@ impl<'buf, 'ast> Formatter<'buf, 'ast> {
         self.state().group_id(debug_name)
     }
 
+    /// Returns a reference to the unique group id builder for this document.
+    #[inline]
+    pub fn group_id_builder(&self) -> &super::UniqueGroupIdBuilder {
+        self.state().group_id_builder()
+    }
+
     /// Joins multiple [Format] together without any separator
     ///
     /// ## Examples

crates/oxc_formatter/src/formatter/group_id.rs

@@ -66,7 +66,7 @@ impl From<GroupId> for u32 {
 }
 
 /// Builder to construct unique group ids that are unique if created with the same builder.
-pub(super) struct UniqueGroupIdBuilder {
+pub struct UniqueGroupIdBuilder {
     next_id: AtomicU32,
 }
 

crates/oxc_formatter/src/formatter/mod.rs

@@ -41,7 +41,7 @@ use std::fmt::Debug;
 
 pub use buffer::{Buffer, BufferExtensions, VecBuffer};
 pub use format_element::FormatElement;
-pub use group_id::GroupId;
+pub use group_id::{GroupId, UniqueGroupIdBuilder};
 
 pub use self::comments::Comments;
 use self::printer::Printer;
@@ -54,7 +54,7 @@ pub use self::{
     state::FormatState,
     text_range::TextRange,
 };
-use self::{format_element::document::Document, group_id::UniqueGroupIdBuilder, prelude::TagKind};
+use self::{format_element::document::Document, prelude::TagKind};
 
 #[derive(Debug)]
 pub struct Formatted<'a> {

crates/oxc_formatter/src/formatter/state.rs

@@ -52,6 +52,11 @@ impl<'ast> FormatState<'ast> {
         self.group_id_builder.group_id(debug_name)
     }
 
+    /// Returns a reference to the unique group id builder.
+    pub fn group_id_builder(&self) -> &UniqueGroupIdBuilder {
+        &self.group_id_builder
+    }
+
     #[expect(clippy::mutable_key_type)]
     pub fn printed_interned_elements(&mut self) -> &mut FxHashMap<Interned<'ast>, usize> {
         &mut self.printed_interned_elements

crates/oxc_formatter/src/lib.rs

@@ -19,15 +19,16 @@ use oxc_span::SourceType;
 
 pub use crate::ast_nodes::{AstNode, AstNodes};
 pub use crate::external_formatter::{
-    EmbeddedDocFormatterCallback, EmbeddedFormatterCallback, EmbeddedIR, ExternalCallbacks,
-    TailwindCallback,
+    EmbeddedDocFormatterCallback, EmbeddedFormatterCallback, ExternalCallbacks, TailwindCallback,
+};
+pub use crate::formatter::format_element::tag::{
+    Align, Condition, DedentMode, Group, GroupMode, Tag,
 };
-pub use crate::formatter::GroupId;
-pub use crate::formatter::format_element::tag::{DedentMode, Tag};
 pub use crate::formatter::format_element::{
-    BestFittingElement, FormatElement, LineMode, PrintMode,
+    BestFittingElement, FormatElement, LineMode, PrintMode, TextWidth,
 };
 pub use crate::formatter::{Format, Formatted};
+pub use crate::formatter::{GroupId, UniqueGroupIdBuilder};
 pub use crate::ir_transform::options::*;
 pub use crate::options::*;
 pub use crate::print::{FormatVueBindingParams, FormatVueScriptGeneric};

crates/oxc_formatter/src/print/template/embed/graphql.rs

@@ -1,19 +1,19 @@
-use rustc_hash::FxHashMap;
-
+use oxc_allocator::Allocator;
 use oxc_ast::ast::*;
 
 use crate::{
     ast_nodes::AstNode,
-    external_formatter::EmbeddedIR,
-    formatter::{Formatter, format_element::LineMode, prelude::*},
+    formatter::{
+        FormatElement, Formatter,
+        format_element::{LineMode, TextWidth},
+        prelude::*,
+    },
     print::template::{
         FormatTemplateExpression, FormatTemplateExpressionOptions, TemplateExpression,
     },
     write,
 };
 
-use super::write_embedded_ir;
-
 /// Format a GraphQL template literal via the Doc→IR path.
 ///
 /// Handles both no-substitution and `${}` templates uniformly.
@@ -87,26 +87,30 @@ pub(super) fn format_graphql_doc<'a>(
     let all_irs = if texts_to_format.is_empty() {
         vec![]
     } else {
-        let Some(Ok(irs)) = f
-            .context()
-            .external_callbacks()
-            .format_embedded_doc("tagged-graphql", &texts_to_format)
-        else {
+        let allocator = f.allocator();
+        let group_id_builder = f.group_id_builder();
+        let Some(Ok(irs)) = f.context().external_callbacks().format_embedded_doc(
+            allocator,
+            group_id_builder,
+            "tagged-graphql",
+            &texts_to_format,
+        ) else {
             return false;
         };
         irs
     };
 
     // Phase 3: Build `ir_parts` by mapping formatted results back to original indices.
-    // Use `into_iter` to take ownership and avoid cloning `Vec<EmbeddedIR>`.
+    // Use `into_iter` to take ownership and avoid cloning.
     let mut irs_iter = all_irs.into_iter();
-    let mut ir_parts: Vec<Option<Vec<EmbeddedIR>>> = Vec::with_capacity(num_quasis);
+    let mut ir_parts: Vec<Option<Vec<FormatElement<'a>>>> = Vec::with_capacity(num_quasis);
     for (idx, info) in infos.iter().enumerate() {
         if format_index_map[idx].is_some() {
             ir_parts.push(irs_iter.next());
         } else if info.comments_only {
             // Build IR for comment-only quasis manually
-            let comment_ir = build_graphql_comment_ir(info.text);
+            let comment_ir =
+                build_graphql_comment_ir(info.text, f.allocator(), f.options().indent_width);
             ir_parts.push(comment_ir);
         } else {
             ir_parts.push(None);
@@ -127,23 +131,22 @@ pub(super) fn format_graphql_doc<'a>(
     // Phase 4: Write the template structure
     // `["`", indent([hardline, join(hardline, parts)]), hardline, "`"]`
     // https://github.com/prettier/prettier/blob/90983f40dce5e20beea4e5618b5e0426a6a7f4f0/src/language-js/embed/graphql.js#L68C10-L68C73
-    let format_content = format_with(|f: &mut Formatter<'_, 'a>| {
-        let mut group_id_map = FxHashMap::default();
+    let format_content = format_once(|f: &mut Formatter<'_, 'a>| {
         let mut has_prev_part = false;
 
-        for (idx, maybe_ir) in ir_parts.iter().enumerate() {
+        for (idx, mut maybe_ir) in ir_parts.into_iter().enumerate() {
             let is_first = idx == 0;
             let is_last = idx == num_quasis - 1;
 
-            if let Some(ir) = maybe_ir {
+            if let Some(ir) = maybe_ir.take() {
                 if !is_first && infos[idx].starts_with_blank_line {
                     if has_prev_part {
                         write!(f, [empty_line()]);
                     }
                 } else if has_prev_part {
                     write!(f, [hard_line_break()]);
                 }
-                write_embedded_ir(ir, f, &mut group_id_map);
+                f.write_elements(ir);
                 has_prev_part = true;
             } else if !is_first && !is_last && infos[idx].starts_with_blank_line && has_prev_part {
                 write!(f, [empty_line()]);
@@ -188,9 +191,14 @@ struct QuasiInfo<'a> {
 /// <https://github.com/prettier/prettier/blob/90983f40dce5e20beea4e5618b5e0426a6a7f4f0/src/language-js/embed/graphql.js#L71>
 ///
 /// Extracts comment lines, joins with hardline, and preserves blank lines between comment groups.
-fn build_graphql_comment_ir(text: &str) -> Option<Vec<EmbeddedIR>> {
+fn build_graphql_comment_ir<'a>(
+    text: &str,
+    allocator: &'a Allocator,
+    indent_width: crate::IndentWidth,
+) -> Option<Vec<FormatElement<'a>>> {
+    // This comes from `.cooked`, which has normalized line terminators
     let lines: Vec<&str> = text.split('\n').map(str::trim).collect();
-    let mut parts: Vec<EmbeddedIR> = vec![];
+    let mut parts: Vec<FormatElement<'a>> = vec![];
     let mut seen_comment = false;
 
     for (i, line) in lines.iter().enumerate() {
@@ -199,16 +207,20 @@ fn build_graphql_comment_ir(text: &str) -> Option<Vec<EmbeddedIR>> {
         }
 
         if i > 0 && lines[i - 1].is_empty() && seen_comment {
-            // Blank line before this comment group → emit empty line + text
-            parts.push(EmbeddedIR::Line(LineMode::Empty));
-            parts.push(EmbeddedIR::ExpandParent);
-            parts.push(EmbeddedIR::Text((*line).to_string()));
+            // Blank line before this comment group -> emit empty line + text
+            parts.push(FormatElement::Line(LineMode::Empty));
+            parts.push(FormatElement::ExpandParent);
+            let arena_text = allocator.alloc_str(line);
+            let width = TextWidth::from_text(arena_text, indent_width);
+            parts.push(FormatElement::Text { text: arena_text, width });
         } else {
             if seen_comment {
-                parts.push(EmbeddedIR::Line(LineMode::Hard));
-                parts.push(EmbeddedIR::ExpandParent);
+                parts.push(FormatElement::Line(LineMode::Hard));
+                parts.push(FormatElement::ExpandParent);
             }
-            parts.push(EmbeddedIR::Text((*line).to_string()));
+            let arena_text = allocator.alloc_str(line);
+            let width = TextWidth::from_text(arena_text, indent_width);
+            parts.push(FormatElement::Text { text: arena_text, width });
         }
 
         seen_comment = true;