fix: autofmt duplicate whitespace on attributes, flickering of vscode autoformat (#5159)

* wip: failing test

* edit harness a bit

* wip, not fixed entirely

* fix comment accumulation and whitespace

* fix: autofmt syncrhonize edits with vscode waitUntil edit pipeline

* fix optimization collapsing away whitespace

* comment out test case for now

packages/autofmt/src/writer.rs

@@ -312,8 +312,10 @@ impl<'a> Writer<'a> {
         let mut opt_level = ShortOptimization::NoOpt;
 
         // check if we have a lot of attributes
-        let attr_len = self.is_short_attrs(attributes, spreads);
-        let is_short_attr_list = (attr_len + self.out.indent_level * 4) < 80;
+        let attr_len = self.is_short_attrs(brace, attributes, spreads);
+        let has_postbrace_comments = self.brace_has_trailing_comments(brace);
+        let is_short_attr_list =
+            ((attr_len + self.out.indent_level * 4) < 80) && !has_postbrace_comments;
         let children_len = self
             .is_short_children(children)
             .map_err(|_| std::fmt::Error)?;
@@ -331,6 +333,7 @@ impl<'a> Writer<'a> {
             && attributes.len() <= 1
             && spreads.is_empty()
             && !has_trailing_comments
+            && !has_postbrace_comments
         {
             if children.is_empty() {
                 opt_level = ShortOptimization::Oneliner;
@@ -595,6 +598,15 @@ impl<'a> Writer<'a> {
         let attr_line = attr_span.start().line;
 
         if brace_line != attr_line {
+            // Get the raw line of the attribute
+            let line = self.src.get(attr_line - 1).unwrap_or(&"");
+
+            // Only write comments if the line is empty before the attribute start
+            let row_start = line.get(..attr_span.start().column - 1).unwrap_or("");
+            if !row_start.trim().is_empty() {
+                return Ok(());
+            }
+
             self.write_comments(attr_span.start())?;
         }
 
@@ -648,16 +660,35 @@ impl<'a> Writer<'a> {
             return comments;
         };
 
+        // We go backwards to collect comments and empty lines. We only want to keep one empty line,
+        // the rest should be `//` comments
+        let mut last_line_was_empty = false;
         for (id, line) in lines.iter().enumerate().rev() {
-            if line.trim().starts_with("//") || line.is_empty() && id != 0 {
-                if id != 0 {
+            let trimmed = line.trim();
+            if trimmed.starts_with("//") {
+                comments.push_front(id);
+                last_line_was_empty = false;
+            } else if trimmed.is_empty() {
+                if !last_line_was_empty {
                     comments.push_front(id);
+                    last_line_was_empty = true;
                 }
+
+                continue;
             } else {
                 break;
             }
         }
 
+        // If there is more than 1 comment, make sure the first comment is not an empty line
+        if comments.len() > 1 {
+            if let Some(&first) = comments.back() {
+                if self.src[first].trim().is_empty() {
+                    comments.pop_back();
+                }
+            }
+        }
+
         comments
     }
 
@@ -725,7 +756,12 @@ impl<'a> Writer<'a> {
         }
     }
 
-    fn is_short_attrs(&mut self, attributes: &[Attribute], spreads: &[Spread]) -> usize {
+    fn is_short_attrs(
+        &mut self,
+        _brace: &Brace,
+        attributes: &[Attribute],
+        spreads: &[Spread],
+    ) -> usize {
         let mut total = 0;
 
         // No more than 3 attributes before breaking the line
@@ -746,7 +782,7 @@ impl<'a> Writer<'a> {
                 };
             }
 
-            let name_len = match &attr.name {
+            total += match &attr.name {
                 AttributeName::BuiltIn(name) => {
                     let name = name.to_string();
                     name.len()
@@ -754,7 +790,6 @@ impl<'a> Writer<'a> {
                 AttributeName::Custom(name) => name.value().len() + 2,
                 AttributeName::Spread(_) => unreachable!(),
             };
-            total += name_len;
 
             if attr.can_be_shorthand() {
                 total += 2;
@@ -1090,6 +1125,13 @@ impl<'a> Writer<'a> {
         }
     }
 
+    fn brace_has_trailing_comments(&self, brace: &Brace) -> bool {
+        let span = brace.span.span();
+        let line = self.src.get(span.start().line - 1).unwrap_or(&"");
+        let after_brace = line.get(span.start().column + 1..).unwrap_or("").trim();
+        after_brace.starts_with("//")
+    }
+
     fn has_trailing_comments(&self, children: &[BodyNode], brace: &Brace) -> bool {
         let brace_span = brace.span.span();
 

packages/autofmt/tests/samples/comments.rsx

@@ -223,7 +223,6 @@ rsx! {
         rsx! { "one two three" }
     }
 
-
     // Please dont eat me 1
     //
     // Please dont eat me 2

packages/autofmt/tests/wrong.rs

@@ -39,3 +39,5 @@ twoway!("oneline-expand" => online_expand (IndentOptions::new(IndentType::Spaces
 twoway!("shortened" => shortened (IndentOptions::new(IndentType::Spaces, 4, false)));
 twoway!("syntax_error" => syntax_error (IndentOptions::new(IndentType::Spaces, 4, false)));
 twoway!("skipfail" => skipfail (IndentOptions::new(IndentType::Spaces, 4, false)));
+twoway!("comments-inline-4sp" => comments_inline_4sp (IndentOptions::new(IndentType::Spaces, 4, false)));
+twoway!("comments-attributes-4sp" => comments_attributes_4sp (IndentOptions::new(IndentType::Spaces, 4, false)));

packages/autofmt/tests/wrong/comments-attributes-4sp.rsx

@@ -0,0 +1,26 @@
+//! hello
+
+#[component]
+fn App() -> Element {
+    rsx! {
+        Stylesheet { href: MAIN_CSS }
+        div { id: "hero",
+            // ss
+            // s
+            // asd
+            img {
+                // blah
+                src: HEADER_SVG,
+                id: "header",
+            }
+            div { id: "links",
+                a { href: "https://discord.gg/XgGxMSkvUM", "👋 Community Discord" }
+                div {
+                    // blah
+                    src: HEADER_SVG,
+                    id: "header",
+                }
+            }
+        }
+    }
+}

packages/autofmt/tests/wrong/comments-attributes-4sp.wrong.rsx

@@ -0,0 +1,26 @@
+//! hello
+
+#[component]
+fn App() -> Element {
+    rsx! {
+        Stylesheet { href: MAIN_CSS }
+        div { id: "hero",
+            // ss
+            // s
+            // asd
+            img {
+                // blah
+                src: HEADER_SVG, id: "header" }
+            div { id: "links",
+                a { href: "https://discord.gg/XgGxMSkvUM", "👋 Community Discord" }
+                div {
+                    // blah
+
+
+
+                    src: HEADER_SVG, id: "header"
+                }
+            }
+        }
+    }
+}

packages/autofmt/tests/wrong/comments-inline-4sp.rsx

@@ -0,0 +1,9 @@
+rsx! {
+    div { // Comments
+        class: "asdasd",
+        "hello world"
+    }
+
+    // todo: fix this case
+    img { src: HEADER_SVG, id: "header" }
+}

packages/autofmt/tests/wrong/comments-inline-4sp.wrong.rsx

@@ -0,0 +1,7 @@
+rsx! {
+    div { // Comments
+        class: "asdasd", "hello world" }
+
+    // todo: fix this case
+    img { src: HEADER_SVG, id: "header" }
+}

packages/extension/src/main.ts

@@ -12,6 +12,13 @@ export async function activate(context: vscode.ExtensionContext) {
 	// VSCode doesn't have a `fetch` implementation, but we don't really care about polyfilling it
 	await init(wasmSourceCode);
 
+	// Register format-on-save handler using waitUntil() for proper synchronization
+	// This runs alongside (not instead of) other formatters like rust-analyzer,
+	// allowing both rustfmt and Dioxus RSX formatting to work together
+	context.subscriptions.push(
+		vscode.workspace.onWillSaveTextDocument(formatOnSave)
+	);
+
 	// Todo:
 	// I want a paste-handler that translates HTML to RSX whenever HTML is pasted into an Rsx block
 	// Or, a little tooltip that pops up and asks if you want to translate the HTML to RSX
@@ -19,12 +26,59 @@ export async function activate(context: vscode.ExtensionContext) {
 		vscode.commands.registerCommand('extension.htmlToDioxusRsx', () => translate(false)),
 		vscode.commands.registerCommand('extension.htmlToDioxusComponent', () => translate(true)),
 		vscode.commands.registerCommand('extension.formatRsx', fmtSelection),
-		vscode.commands.registerCommand('extension.formatRsxDocument', formatRsxDocument),
-		vscode.workspace.onWillSaveTextDocument(fmtDocumentOnSave)
+		vscode.commands.registerCommand('extension.formatRsxDocument', formatRsxDocument)
 	);
 
 	context.subscriptions.push(vscode.window.registerUriHandler(new UriLaunchServer()));
+}
 
+// Format RSX on save using waitUntil() for proper synchronization with VSCode's save pipeline
+function formatOnSave(e: vscode.TextDocumentWillSaveEvent) {
+	if (e.document.languageId !== 'rust') {
+		return;
+	}
+
+	// Check if Dioxus formatting is enabled
+	const dioxusConfig = vscode.workspace.getConfiguration('dioxus', e.document).get('formatOnSave');
+	if (dioxusConfig === 'disabled') {
+		return;
+	}
+
+	// Use waitUntil() to properly synchronize with VSCode's save pipeline
+	// This returns TextEdit[] which VSCode applies before completing the save
+	e.waitUntil(formatDocument(e.document));
+}
+
+// Returns a promise of TextEdit[] for use with waitUntil()
+function formatDocument(document: vscode.TextDocument): Thenable<vscode.TextEdit[]> {
+	return new Promise((resolve) => {
+		try {
+			const contents = document.getText();
+
+			// Get editor options for this document
+			const editor = vscode.window.visibleTextEditors.find(
+				e => e.document.uri.toString() === document.uri.toString()
+			);
+
+			const tabSize = (typeof editor?.options.tabSize === 'number') ? editor.options.tabSize : 4;
+			const useTabs = editor ? !editor.options.insertSpaces : false;
+
+			const formatted = dioxus.format_file(contents, useTabs, tabSize);
+
+			if (formatted.length() > 0) {
+				const fullRange = new vscode.Range(
+					document.positionAt(0),
+					document.positionAt(contents.length)
+				);
+				resolve([vscode.TextEdit.replace(fullRange, formatted.formatted())]);
+			} else {
+				resolve([]);
+			}
+		} catch (error) {
+			vscode.window.showWarningMessage(`Dioxus formatting error: ${error}`);
+			resolve([]);
+		}
+	});
 }
 
 function translate(component: boolean) {
@@ -49,11 +103,19 @@ function translate(component: boolean) {
 }
 
 
-function formatRsxDocument() {
+async function formatRsxDocument() {
 	const editor = vscode.window.activeTextEditor;
 	if (!editor) return;
 
-	fmtDocument(editor.document);
+	// Apply RSX formatting directly using a WorkspaceEdit
+	const edits = await formatDocument(editor.document);
+	if (edits.length > 0) {
+		const workspaceEdit = new vscode.WorkspaceEdit();
+		for (const edit of edits) {
+			workspaceEdit.replace(editor.document.uri, edit.range, edit.newText);
+		}
+		await vscode.workspace.applyEdit(workspaceEdit);
+	}
 }
 
 function fmtSelection() {
@@ -132,49 +194,6 @@ function fmtSelection() {
 
 }
 
-function fmtDocumentOnSave(e: vscode.TextDocumentWillSaveEvent) {
-	// check the settings to make sure format on save is configured
-	const dioxusConfig = vscode.workspace.getConfiguration('dioxus', e.document).get('formatOnSave');
-	const globalConfig = vscode.workspace.getConfiguration('editor', e.document).get('formatOnSave');
-	if (
-		(dioxusConfig === 'enabled') ||
-		(dioxusConfig !== 'disabled' && globalConfig)
-	) {
-		fmtDocument(e.document);
-	}
-}
-
-function fmtDocument(document: vscode.TextDocument) {
-	try {
-		if (document.languageId !== "rust") {
-			return;
-		}
-
-		const [editor,] = vscode.window.visibleTextEditors.filter(editor => editor.document.fileName === document.fileName);
-		if (!editor) return; // Need an editor to apply text edits.
-
-		const contents = editor.document.getText();
-		let tabSize: number;
-		if (typeof editor.options.tabSize === 'number') {
-			tabSize = editor.options.tabSize;
-		} else {
-			tabSize = 4;
-		}
-		const formatted = dioxus.format_file(contents, !editor.options.insertSpaces, tabSize);
-
-		// Replace the entire text document
-		// Yes, this is a bit heavy handed, but the dioxus side doesn't know the line/col scheme that vscode is using
-		if (formatted.length() > 0) {
-			editor.edit(editBuilder => {
-				const range = new vscode.Range(0, 0, document.lineCount, 0);
-				editBuilder.replace(range, formatted.formatted());
-			});
-		}
-	} catch (error) {
-		vscode.window.showWarningMessage(`Errors occurred while formatting. Make sure you have the most recent Dioxus-CLI installed! \n${error}`);
-	}
-}
-
 class UriLaunchServer implements vscode.UriHandler {
 	handleUri(uri: vscode.Uri): vscode.ProviderResult<void> {
 		if (uri.path === '/debugger') {