From 016538aa1e6dca920b524ec8d8549ba66de17ee8 Mon Sep 17 00:00:00 2001
From: Yaroslav Bolyukin <iam@lach.pw>
Date: Wed, 22 Jun 2022 18:30:04 +0000
Subject: [PATCH] feat(fmt): preserve comments at root of source code


---
--- a/cmds/jrsonnet-fmt/src/children.rs
+++ b/cmds/jrsonnet-fmt/src/children.rs
@@ -1,3 +1,5 @@
+// TODO: Return errors as trivia
+
 use std::{fmt::Debug, marker::PhantomData, mem};
 
 use jrsonnet_rowan_parser::{
@@ -9,9 +11,52 @@
 
 pub type ChildTrivia = Vec<Trivia>;
 
-pub struct ChildIterator<I, T> {
-	inner: I,
-	_marker: PhantomData<T>,
+/// Node should have no non-trivia tokens before element
+pub fn trivia_before(node: SyntaxNode, end: Option<&SyntaxElement>) -> ChildTrivia {
+	let mut out = Vec::new();
+	for item in node.children_with_tokens() {
+		if Some(&item) == end {
+			break;
+		}
+
+		if let Some(trivia) = item.as_token().cloned().and_then(Trivia::cast) {
+			out.push(trivia);
+		} else if end.is_none() {
+			break;
+		} else {
+			assert!(
+				TS![, ;].contains(item.kind()) || item.kind() == ERROR,
+				"silently eaten token: {:?}",
+				item.kind()
+			)
+		}
+	}
+	out
+}
+/// Node should have no non-trivia tokens after element
+pub fn trivia_after(node: SyntaxNode, start: Option<&SyntaxElement>) -> ChildTrivia {
+	if start.is_none() {
+		return Vec::new();
+	}
+	let mut iter = node.children_with_tokens().peekable();
+	while iter.peek() != start {
+		// println!("Skipped {}");
+		dbg!(&iter.next());
+	}
+	dbg!(&iter.next());
+	let mut out = Vec::new();
+	for item in iter {
+		if let Some(trivia) = item.as_token().cloned().and_then(Trivia::cast) {
+			out.push(trivia);
+		} else {
+			assert!(
+				TS![, ;].contains(item.kind()) || item.kind() == ERROR,
+				"silently eaten token: {:?}",
+				item.kind()
+			)
+		}
+	}
+	out
 }
 
 pub fn children_between<T: AstNode + Debug>(
@@ -20,9 +65,10 @@
 	end: Option<&SyntaxElement>,
 ) -> (Vec<Child<T>>, ChildTrivia) {
 	let mut iter = node.children_with_tokens().peekable();
-	while iter.peek() == start {
+	while iter.peek() != start {
 		iter.next();
 	}
+	iter.next();
 	children(
 		iter.take_while(|i| Some(i) != end),
 		start.is_none() || end.is_none(),
@@ -31,7 +77,7 @@
 
 pub fn should_start_with_newline(tt: &ChildTrivia) -> bool {
 	// First for previous item end
-	count_newlines_before(&tt) >= 2
+	count_newlines_before(tt) >= 2
 }
 
 fn count_newlines_before(tt: &ChildTrivia) -> usize {
--- a/cmds/jrsonnet-fmt/src/main.rs
+++ b/cmds/jrsonnet-fmt/src/main.rs
@@ -1,6 +1,6 @@
 use std::any::type_name;
 
-use children::children_between;
+use children::{children_between, trivia_before};
 use dprint_core::formatting::{PrintItems, PrintOptions};
 use jrsonnet_rowan_parser::{
 	nodes::{
@@ -13,7 +13,7 @@
 };
 
 use crate::{
-	children::should_start_with_newline,
+	children::{should_start_with_newline, trivia_after},
 	comments::{format_comments, CommentLocation},
 };
 
@@ -463,8 +463,25 @@
 
 impl Printable for SourceFile {
 	fn print(&self) -> PrintItems {
-		assert!(self.expr().is_some());
-		self.expr().print()
+		let mut pi = p!(new:);
+		let before = trivia_before(
+			self.syntax().clone(),
+			self.expr()
+				.map(|e| e.syntax().clone())
+				.map(Into::into)
+				.as_ref(),
+		);
+		let after = trivia_after(
+			self.syntax().clone(),
+			self.expr()
+				.map(|e| e.syntax().clone())
+				.map(Into::into)
+				.as_ref(),
+		);
+		p!(pi: items(format_comments(&before, CommentLocation::AboveItem)));
+		p!(pi: {self.expr()} nl);
+		p!(pi: items(format_comments(&after, CommentLocation::EndOfItems)));
+		pi
 	}
 }
 
@@ -574,6 +591,7 @@
 		} + Template
 
 
+		// Comment after everything
 "#,
 	);
 
--- a/crates/jrsonnet-rowan-parser/src/event.rs
+++ b/crates/jrsonnet-rowan-parser/src/event.rs
@@ -117,6 +117,9 @@
 					eat_start_whitespace = false;
 				}
 				Event::Finish { wrapper } => {
+					if depth == 1 {
+						self.skip_whitespace();
+					}
 					self.builder.finish_node();
 					depth -= 1;
 					let mut idx = idx;
@@ -126,6 +129,9 @@
 						wrapper = if let Event::Finish { wrapper } =
 							mem::replace(&mut self.events[idx], Event::Noop)
 						{
+							if depth == 1 {
+								self.skip_whitespace();
+							}
 							self.builder.finish_node();
 							depth -= 1;
 							wrapper

-- 
gitstuff