git.delta.rocks / jrsonnet / refs/commits / 016538aa1e6d

--- 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
 "#,
 	);
 

117					eat_start_whitespace = false;117					eat_start_whitespace = false;
118				}118				}
119				Event::Finish { wrapper } => {119				Event::Finish { wrapper } => {
120					if depth == 1 {
121						self.skip_whitespace();
122					}
120					self.builder.finish_node();123					self.builder.finish_node();
121					depth -= 1;124					depth -= 1;
122					let mut idx = idx;125					let mut idx = idx;
126						wrapper = if let Event::Finish { wrapper } =129						wrapper = if let Event::Finish { wrapper } =
127							mem::replace(&mut self.events[idx], Event::Noop)130							mem::replace(&mut self.events[idx], Event::Noop)
128						{131						{
132							if depth == 1 {
133								self.skip_whitespace();
134							}
129							self.builder.finish_node();135							self.builder.finish_node();
130							depth -= 1;136							depth -= 1;
131							wrapper137							wrapper