From dacee201fa00ebeec122b0b4927eee61900cb4d3 Mon Sep 17 00:00:00 2001
From: Yaroslav Bolyukin <iam@lach.pw>
Date: Wed, 22 Jun 2022 19:08:09 +0000
Subject: [PATCH] feat(fmt): preserve comments in locals


---
--- a/cmds/jrsonnet-fmt/src/children.rs
+++ b/cmds/jrsonnet-fmt/src/children.rs
@@ -40,10 +40,9 @@
 	}
 	let mut iter = node.children_with_tokens().peekable();
 	while iter.peek() != start {
-		// println!("Skipped {}");
-		dbg!(&iter.next());
+		iter.next();
 	}
-	dbg!(&iter.next());
+	iter.next();
 	let mut out = Vec::new();
 	for item in iter {
 		if let Some(trivia) = item.as_token().cloned().and_then(Trivia::cast) {
@@ -59,6 +58,36 @@
 	out
 }
 
+pub fn trivia_between(
+	node: SyntaxNode,
+	start: Option<&SyntaxElement>,
+	end: Option<&SyntaxElement>,
+) -> ChildTrivia {
+	let mut iter = node.children_with_tokens().peekable();
+	while iter.peek() != start {
+		iter.next();
+	}
+	iter.next();
+
+	let loose = start.is_none() || end.is_none();
+
+	let mut out = Vec::new();
+	for item in iter.take_while(|i| Some(i) != end) {
+		if let Some(trivia) = item.as_token().cloned().and_then(Trivia::cast) {
+			out.push(trivia);
+		} else if loose {
+			break;
+		} else {
+			assert!(
+				TS![, ;].contains(item.kind()) || item.kind() == ERROR,
+				"silently eaten token: {:?}",
+				item.kind()
+			)
+		}
+	}
+	out
+}
+
 pub fn children_between<T: AstNode + Debug>(
 	node: SyntaxNode,
 	start: Option<&SyntaxElement>,
--- a/cmds/jrsonnet-fmt/src/comments.rs
+++ b/cmds/jrsonnet-fmt/src/comments.rs
@@ -12,6 +12,7 @@
 	EndOfItems,
 }
 
+#[must_use]
 pub fn format_comments(comments: &ChildTrivia, loc: CommentLocation) -> PrintItems {
 	let mut pi = p!(new:);
 
--- a/cmds/jrsonnet-fmt/src/main.rs
+++ b/cmds/jrsonnet-fmt/src/main.rs
@@ -13,7 +13,7 @@
 };
 
 use crate::{
-	children::{should_start_with_newline, trivia_after},
+	children::{should_start_with_newline, trivia_after, trivia_between},
 	comments::{format_comments, CommentLocation},
 };
 
@@ -436,11 +436,45 @@
 			}
 			Expr::ExprVar(n) => p!(new: {n.name()}),
 			Expr::ExprLocal(l) => {
-				let mut pi = p!(new: str("local") >i nl);
-				for bind in l.binds() {
-					p!(pi: {bind} str(",") nl);
+				let mut pi = p!(new:);
+				let (binds, end_comments) = children_between::<Bind>(
+					l.syntax().clone(),
+					l.local_kw_token().map(Into::into).as_ref(),
+					l.semi_token().map(Into::into).as_ref(),
+				);
+				if binds.len() == 1 {
+					let bind = &binds[0];
+					p!(pi: items(format_comments(&bind.before_trivia, CommentLocation::AboveItem)));
+					p!(pi: str("local ") {bind.value});
+				// TODO: keep end_comments, child.inline_trivia somehow, force multiple locals formatting in case of presence?
+				} else {
+					p!(pi: str("local") >i nl);
+					for bind in binds {
+						if bind.needs_newline_above() {
+							p!(pi: nl);
+						}
+						p!(pi: items(format_comments(&bind.before_trivia, CommentLocation::AboveItem)));
+						p!(pi: {bind.value} str(";"));
+						p!(pi: items(format_comments(&bind.inline_trivia, CommentLocation::ItemInline)) nl);
+					}
+					// TODO: needs_newline_above end_comments
+					p!(pi: items(format_comments(&end_comments, CommentLocation::EndOfItems)));
+					p!(pi: <i);
 				}
-				p!(pi: <i str(";") nl {l.expr()});
+				p!(pi: str(";") nl);
+
+				let expr_comments = trivia_between(
+					l.syntax().clone(),
+					l.semi_token().map(Into::into).as_ref(),
+					l.expr()
+						.map(|e| e.syntax().clone())
+						.map(Into::into)
+						.as_ref(),
+				);
+				p!(pi: items(format_comments(&expr_comments, CommentLocation::AboveItem)));
+
+				// TODO: needs_newline_above expr
+				p!(pi: {l.expr()});
 				pi
 			}
 			Expr::ExprIfThenElse(ite) => {
@@ -515,6 +549,25 @@
 
 		local unary = !a;
 
+		local
+			//   I am comment
+			singleLocalWithItemComment = 1,
+		;
+
+		// Comment between local and expression
+
+		local
+			a = 1, //   Inline
+			// Comment above b
+			b = 4,
+
+			// c needs some space
+			c = 5,
+
+			// Comment after everything
+		;
+
+
 		local Template = {z: "foo"};
 
 		{

-- 
gitstuff