Fix comment handling / output formatting

Author: breml

ast/ast.go

@@ -607,7 +607,6 @@ func (b Branch) String() string {
 		s.WriteString(block.String())
 	}
 	s.WriteString(b.ElseBlock.String())
-	s.WriteString("\n")
 	return s.String()
 }
 
@@ -638,6 +637,9 @@ func (ib IfBlock) String() string {
 		if block == nil {
 			continue
 		}
+		if ss.Len() > 0 {
+			ss.WriteString("\n")
+		}
 		ss.WriteString(fmt.Sprintf("%v", block))
 		if ss.Len() > 0 {
 			ss.WriteString("\n")
@@ -650,6 +652,7 @@ func (ib IfBlock) String() string {
 	s.WriteString(prefix(ss.String(), true))
 
 	s.WriteString("}")
+
 	return s.String()
 }
 
@@ -685,6 +688,9 @@ func (eib ElseIfBlock) String() string {
 		if block == nil {
 			continue
 		}
+		if ss.Len() > 0 {
+			ss.WriteString("\n")
+		}
 		ss.WriteString(fmt.Sprint(block))
 		if ss.Len() > 0 {
 			ss.WriteString("\n")
@@ -733,6 +739,9 @@ func (eb ElseBlock) String() string {
 		if block == nil {
 			continue
 		}
+		if ss.Len() > 0 {
+			ss.WriteString("\n")
+		}
 		ss.WriteString(fmt.Sprint(block))
 		if ss.Len() > 0 {
 			ss.WriteString("\n")

diff_test.go

@@ -22,30 +22,42 @@ func printDiff(want, got string) string {
 		switch diff.Type {
 		case diffmatchpatch.DiffInsert:
 			buff.WriteString("\x1b[32m")
-			buff.WriteString("+ ")
-			buff.WriteString(highlightWhitespaces(text, "\x1b[42m"))
+			lines := strings.Split(strings.TrimRight(text, "\n"), "\n")
+			for _, line := range lines {
+				buff.WriteString("+ ")
+				buff.WriteString(highlightWhitespaces(line, "\x1b[42m") + "\n")
+			}
 			buff.WriteString("\x1b[0m")
 			if !strings.HasSuffix(text, "\n") {
 				buff.WriteString("\n\\ No newline at end of file\n")
 			}
 		case diffmatchpatch.DiffDelete:
 			buff.WriteString("\x1b[31m")
-			buff.WriteString("- ")
-			buff.WriteString(highlightWhitespaces(text, "\x1b[41m"))
+			lines := strings.Split(strings.TrimRight(text, "\n"), "\n")
+			for _, line := range lines {
+				buff.WriteString("- ")
+				buff.WriteString(highlightWhitespaces(line, "\x1b[41m") + "\n")
+			}
 			buff.WriteString("\x1b[0m")
 			if !strings.HasSuffix(text, "\n") {
 				buff.WriteString("\n\\ No newline at end of file\n")
 			}
+			if !strings.HasSuffix(text, "\n") {
+				buff.WriteString("\n\\ No newline at end of file\n")
+			}
 		case diffmatchpatch.DiffEqual:
-			buff.WriteString("  ")
-			buff.WriteString(text)
+			lines := strings.Split(strings.TrimRight(text, "\n"), "\n")
+			for _, line := range lines {
+				buff.WriteString("  ")
+				buff.WriteString(line + "\n")
+			}
 		}
 	}
 	return buff.String()
 }
 
-var tailingWhitespace = regexp.MustCompile(`([ \t]+)\n$`)
+var tailingWhitespace = regexp.MustCompile(`([ \t]+)$`)
 
 func highlightWhitespaces(in string, color string) string {
-	return tailingWhitespace.ReplaceAllString(in, color+"$1\x1b[49m\n")
+	return tailingWhitespace.ReplaceAllString(in, color+"$1\x1b[49m")
 }