Finally got the printing of comments right. Also added a test to see if the delomboking a file is an idempotent operation

3 files changed, 96 insertions, 46 deletions

diff --git a/src/delombok/lombok/delombok/Comment.java b/src/delombok/lombok/delombok/Comment.java
index 170ffe22..5a129788 100644
--- a/src/delombok/lombok/delombok/Comment.java
+++ b/src/delombok/lombok/delombok/Comment.java
@@ -29,24 +29,31 @@ public final class Comment {
 		DIRECT_AFTER_PREVIOUS,
 		AFTER_PREVIOUS
 	}
+
+	public enum EndConnection {
+		DIRECT_AFTER_COMMENT,
+		AFTER_COMMENT,
+		ON_NEXT_LINE
+	}
 	
 	final int pos;
 	final int prevEndPos;
 	final String content;
 	final int endPos;
 	final StartConnection start;
-	final boolean newLineAfter;
+	final EndConnection end;
 	
-	public Comment(int prevEndPos, int pos, int endPos, String content, StartConnection start, boolean newLineAfter) {
+	public Comment(int prevEndPos, int pos, int endPos, String content, StartConnection start, EndConnection end) {
 		this.pos = pos;
 		this.prevEndPos = prevEndPos;
 		this.endPos = endPos;
 		this.content = content;
 		this.start = start;
-		this.newLineAfter = newLineAfter;
+		this.end = end;
 	}
 	
-	public String summary() {
-		return String.format("%d: %s", pos, content);
+	@Override
+	public String toString() {
+		return String.format("%d: %s (%s,%s)", pos, content, start, end);
 	}
 }
diff --git a/src/delombok/lombok/delombok/CommentCollectingScanner.java b/src/delombok/lombok/delombok/CommentCollectingScanner.java
index 8c77a0d3..2552cbf7 100644
--- a/src/delombok/lombok/delombok/CommentCollectingScanner.java
+++ b/src/delombok/lombok/delombok/CommentCollectingScanner.java
@@ -23,6 +23,7 @@ package lombok.delombok;
 
 import java.nio.CharBuffer;
 
+import lombok.delombok.Comment.EndConnection;
 import lombok.delombok.Comment.StartConnection;
 import lombok.delombok.CommentPreservingParser.Comments;
 
@@ -31,6 +32,7 @@ import com.sun.tools.javac.util.Context;
 
 public class CommentCollectingScanner extends Scanner {
 	private final Comments comments;
+	private int endComment = 0;
 
 	/** A factory for creating scanners. */
 	public static class Factory extends Scanner.Factory {
@@ -78,17 +80,33 @@ public class CommentCollectingScanner extends Scanner {
 	
 	@Override
 	protected void processComment(CommentStyle style) {
-		int prevEndPos = prevEndPos();
+		int prevEndPos = Math.max(prevEndPos(), endComment);
 		int pos = pos();
 		int endPos = endPos();
+		endComment = endPos;
 		String content = new String(getRawCharacters(pos, endPos));
-		boolean lineBreakAfter = isNewLine(getRawCharacters(endPos, endPos + 1)[0]);
 		StartConnection start = determineStartConnection(prevEndPos, pos);
-		Comment comment = new Comment(prevEndPos, pos, endPos, content, start, lineBreakAfter);
-		
+		EndConnection end = determineEndConnection(endPos); 
+
+		Comment comment = new Comment(prevEndPos, pos, endPos, content, start, end);
 		comments.add(comment);
 	}
 	
+	private EndConnection determineEndConnection(int pos) {
+		boolean first = true;
+		for (int i = pos;; i++) {
+			char c = getRawCharacters(i, i + 1)[0];
+			if (isNewLine(c)) {
+				return EndConnection.ON_NEXT_LINE;
+			}
+			if (Character.isWhitespace(c)) {
+				first = false;
+				continue;
+			}
+			return first ? EndConnection.DIRECT_AFTER_COMMENT : EndConnection.AFTER_COMMENT;
+		}
+	}
+
 	private StartConnection determineStartConnection(int from, int to) {
 		if (from == to) {
 			return StartConnection.DIRECT_AFTER_PREVIOUS;
diff --git a/src/delombok/lombok/delombok/PrettyCommentsPrinter.java b/src/delombok/lombok/delombok/PrettyCommentsPrinter.java
index 695e94da..1d04bb5f 100644
--- a/src/delombok/lombok/delombok/PrettyCommentsPrinter.java
+++ b/src/delombok/lombok/delombok/PrettyCommentsPrinter.java
@@ -38,6 +38,7 @@ import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.Map;
 
+import lombok.delombok.Comment.EndConnection;
 import lombok.delombok.Comment.StartConnection;
 
 import com.sun.source.tree.Tree;
@@ -155,10 +156,14 @@ public class PrettyCommentsPrinter extends JCTree.Visitor {
 	
 	private List<Comment> comments;
 	private final JCCompilationUnit cu;
-	private boolean newLine = true;
-	private boolean indent = false;
+	private boolean onNewLine = true;
+	private boolean aligned = false;
 	private boolean inParams = false;
 	
+	private boolean needsSpace = false;
+	private boolean needsNewLine = false;
+	private boolean needsAlign = false;
+	
     public PrettyCommentsPrinter(Writer out, JCCompilationUnit cu, List<Comment> comments) {
         this.out = out;
 		this.comments = comments;
@@ -170,65 +175,70 @@ public class PrettyCommentsPrinter extends JCTree.Visitor {
 	}
     
     private void consumeComments(int till) throws IOException {
-    	boolean shouldIndent = !newLine;
+    	boolean prevNewLine = onNewLine;
     	boolean found = false;
     	Comment head = comments.head;
 		while (comments.nonEmpty() && head.pos < till) {
-			found = true;
 			printComment(head);
 			comments = comments.tail;
 			head = comments.head;
 		}
-		if (found && !newLine && !indent) {
-			print(" ");
-		}
-		if (newLine && shouldIndent) {
-			align();
+		if (!onNewLine && prevNewLine) {
+			println();
 		}
 	}
 
     private void consumeTrailingComments(int from) throws IOException {
-    	boolean shouldIndent = !newLine;
+    	boolean prevNewLine = onNewLine;
 		Comment head = comments.head;
-		while (comments.nonEmpty() && head.prevEndPos == from) {
+		boolean stop = false;
+		while (comments.nonEmpty() && head.prevEndPos == from && !stop && !(head.start == StartConnection.ON_NEXT_LINE || head.start == StartConnection.START_OF_LINE)) {
 			from = head.endPos;
 			printComment(head);
+			stop = (head.end == EndConnection.ON_NEXT_LINE);
 			comments = comments.tail;
 			head = comments.head;
 		}
-		if (newLine && shouldIndent) {
-			align();
+		if (!onNewLine && prevNewLine) {
+			println();
 		}
 	}
-    
-    private void printComment(Comment comment) throws IOException {
+
+	private void printComment(Comment comment) throws IOException {
     	prepareComment(comment.start);
     	print(comment.content);
-    	if (comment.newLineAfter) {
-    		println();
-    	}
+    	switch (comment.end) {
+		case ON_NEXT_LINE:
+			if (!aligned) {
+				needsNewLine = true;
+				needsAlign = true;
+			}
+			break;
+		case AFTER_COMMENT:
+			needsSpace = true;
+			break;
+		case DIRECT_AFTER_COMMENT:
+			// do nothing
+			break;
+		}
     }
 
 	private void prepareComment(StartConnection start) throws IOException {
 		switch (start) {
 			case DIRECT_AFTER_PREVIOUS:
+				needsSpace = false;
 				break;
 			case AFTER_PREVIOUS:
-				if (!newLine && !indent) {
-					print(" ");
-				}
+				needsSpace = true;
 				break;
 			case START_OF_LINE:
-				if (!newLine) {
-					println();
-				}
+				needsNewLine = true;
+				needsAlign = false;
 				break;
 			case ON_NEXT_LINE:
-				if (!newLine && !indent) {
-					println();
-				}
-				if (!indent) {
-					align();
+				if (!aligned) {
+					needsNewLine = true;
+					needsAlign = true;
 				}
 				break;
 		}
@@ -254,8 +264,9 @@ public class PrettyCommentsPrinter extends JCTree.Visitor {
     /** Align code to be indented to left margin.
      */
     void align() throws IOException {
-    	newLine = false;
-    	indent = true;
+    	onNewLine = false;
+    	aligned = true;
+    	needsAlign = false;
         for (int i = 0; i < lmargin; i++) out.write("\t");
     }
 
@@ -292,16 +303,30 @@ public class PrettyCommentsPrinter extends JCTree.Visitor {
     /** Print string, replacing all non-ascii character with unicode escapes.
      */
     public void print(Object s) throws IOException {
-    	newLine = false;
-    	indent = false;
-        out.write(Convert.escapeUnicode(s.toString()));
+    	boolean align = needsAlign;
+    	if (needsNewLine && !onNewLine) {
+    		println();
+    	}
+    	if (align && !aligned) {
+    		align();
+    	}
+    	if (needsSpace && !onNewLine && !aligned) {
+    		out.write(' ');
+    	}
+    	needsSpace = false;
+        
+    	out.write(Convert.escapeUnicode(s.toString()));
+       
+        onNewLine = false;
+        aligned = false;
     }
 
     /** Print new line.
      */
     public void println() throws IOException {
-    	newLine = true;
-    	indent = false;
+    	onNewLine = true;
+    	aligned = false;
+    	needsNewLine = false;
         out.write(lineSep);
     }
 
@@ -337,7 +362,6 @@ public class PrettyCommentsPrinter extends JCTree.Visitor {
                	tree.accept(this);
                	int endPos = endPos(tree);
 				consumeTrailingComments(endPos);
-//               	consumeComments(endPos);
             }
         } catch (UncheckedIOException ex) {
             IOException e = new IOException(ex.getMessage());
@@ -566,6 +590,7 @@ public class PrettyCommentsPrinter extends JCTree.Visitor {
     public void visitTopLevel(JCCompilationUnit tree) {
         try {
             printUnit(tree, null);
+            consumeComments(Integer.MAX_VALUE);
         } catch (IOException e) {
             throw new UncheckedIOException(e);
         }