From 09c6a68804ac16e8b89034de9340601a08b32124 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Thu, 23 Mar 2023 11:49:35 +0530
Subject: [PATCH] Fix syntax highlighting of multiline tokens

---
 tools/cmd/diff/highlight.go | 27 ++++++++++++++++++++++++---
 tools/utils/strings_test.go |  1 +
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/tools/cmd/diff/highlight.go b/tools/cmd/diff/highlight.go
index 882626875..a5138132d 100644
--- a/tools/cmd/diff/highlight.go
+++ b/tools/cmd/diff/highlight.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
 	"path/filepath"
 	"strings"
 
@@ -18,6 +19,8 @@ import (
 )
 
 var _ = fmt.Print
+var _ = os.WriteFile
+
 var ErrNoLexer = errors.New("No lexer available for this format")
 var DefaultStyle = (&utils.Once[*chroma.Style]{Run: func() *chroma.Style {
 	// Default style generated by python style.py default pygments.styles.default.DefaultStyle
@@ -83,6 +86,7 @@ const SGR_SUFFIX = "m"
 func ansi_formatter(w io.Writer, style *chroma.Style, it chroma.Iterator) error {
 	style = clear_background(style)
 	before, after := make([]byte, 0, 64), make([]byte, 0, 64)
+	nl := []byte{'\n'}
 	write_sgr := func(which []byte) {
 		if len(which) > 1 {
 			w.Write(utils.UnsafeStringToBytes(SGR_PREFIX))
@@ -90,6 +94,12 @@ func ansi_formatter(w io.Writer, style *chroma.Style, it chroma.Iterator) error
 			w.Write(utils.UnsafeStringToBytes(SGR_SUFFIX))
 		}
 	}
+	write := func(text string) {
+		write_sgr(before)
+		w.Write(utils.UnsafeStringToBytes(text))
+		write_sgr(after)
+	}
+
 	for token := it(); token != chroma.EOF; token = it() {
 		entry := style.Get(token.Type)
 		before, after = before[:0], after[:0]
@@ -111,9 +121,19 @@ func ansi_formatter(w io.Writer, style *chroma.Style, it chroma.Iterator) error
 				after = append(after, '3', '9', ';')
 			}
 		}
-		write_sgr(before)
-		w.Write(utils.UnsafeStringToBytes(token.Value))
-		write_sgr(after)
+		// independently format each line in a multiline token, needed for the diff kitten highlighting to work, also
+		// pagers like less reset SGR formatting at line boundaries
+		text := token.Value
+		for text != "" {
+			idx := strings.IndexByte(text, '\n')
+			if idx < 0 {
+				write(text)
+				break
+			}
+			write(text[:idx])
+			w.Write(nl)
+			text = text[idx+1:]
+		}
 	}
 	return nil
 }
@@ -172,6 +192,7 @@ func highlight_file(path string) (highlighted string, err error) {
 	w := strings.Builder{}
 	w.Grow(len(text) * 2)
 	err = formatter.Format(&w, style, iterator)
+	// os.WriteFile(path+".highlighted", []byte(w.String()), 0o600)
 	return w.String(), err
 }
 
diff --git a/tools/utils/strings_test.go b/tools/utils/strings_test.go
index d6694fc4d..0c3851947 100644
--- a/tools/utils/strings_test.go
+++ b/tools/utils/strings_test.go
@@ -21,6 +21,7 @@ func TestStringScanner(t *testing.T) {
 		"a\r\r\nb\r\nc\n",
 		"\n1",
 		"",
+		"\n",
 	} {
 		actual := Splitlines(text)
 		expected := make([]string, 0, len(actual))