feat(diffview): add syntax highlighting

Andrey Nering created 4 months ago

Change summary

internal/exp/diffview/chroma.go        |  52 ++++++++++++++
internal/exp/diffview/diffview.go      | 102 +++++++++++++++++++++------
internal/exp/diffview/diffview_test.go |  29 ++++++-
3 files changed, 155 insertions(+), 28 deletions(-)

Detailed changes

internal/exp/diffview/chroma.go

@@ -0,0 +1,52 @@
+package diffview
+
+import (
+	"fmt"
+	"image/color"
+	"io"
+
+	"github.com/alecthomas/chroma/v2"
+	"github.com/charmbracelet/lipgloss/v2"
+)
+
+var _ chroma.Formatter = chromaFormatter{}
+
+// chromaFormatter is a custom formatter for Chroma that uses Lip Gloss for
+// foreground styling, while keeping a forced background color.
+type chromaFormatter struct {
+	bgColor color.Color
+}
+
+// Format implements the chroma.Formatter interface.
+func (c chromaFormatter) Format(w io.Writer, style *chroma.Style, it chroma.Iterator) error {
+	for token := it(); token != chroma.EOF; token = it() {
+		entry := style.Get(token.Type)
+		if entry.IsZero() {
+			if _, err := fmt.Fprint(w, token.Value); err != nil {
+				return err
+			}
+			continue
+		}
+
+		s := lipgloss.NewStyle().
+			Background(c.bgColor)
+
+		if entry.Bold == chroma.Yes {
+			s = s.Bold(true)
+		}
+		if entry.Underline == chroma.Yes {
+			s = s.Underline(true)
+		}
+		if entry.Italic == chroma.Yes {
+			s = s.Italic(true)
+		}
+		if entry.Colour.IsSet() {
+			s = s.Foreground(lipgloss.Color(entry.Colour.String()))
+		}
+
+		if _, err := fmt.Fprint(w, s.Render(token.Value)); err != nil {
+			return err
+		}
+	}
+	return nil
+}

internal/exp/diffview/diffview.go

@@ -2,10 +2,13 @@ package diffview
 
 import (
 	"fmt"
+	"image/color"
 	"os"
 	"strconv"
 	"strings"
 
+	"github.com/alecthomas/chroma/v2"
+	"github.com/alecthomas/chroma/v2/lexers"
 	"github.com/aymanbagabas/go-udiff"
 	"github.com/aymanbagabas/go-udiff/myers"
 	"github.com/charmbracelet/lipgloss/v2"
@@ -43,6 +46,7 @@ type DiffView struct {
 	yOffset      int
 	style        Style
 	tabWidth     int
+	chromaStyle  *chroma.Style
 
 	isComputed bool
 	err        error
@@ -153,6 +157,13 @@ func (dv *DiffView) TabWidth(tabWidth int) *DiffView {
 	return dv
 }
 
+// ChromaStyle sets the chroma style for syntax highlighting.
+// If nil, no syntax highlighting will be applied.
+func (dv *DiffView) ChromaStyle(style *chroma.Style) *DiffView {
+	dv.chromaStyle = style
+	return dv
+}
+
 // String returns the string representation of the DiffView.
 func (dv *DiffView) String() string {
 	dv.replaceTabs()
@@ -361,14 +372,19 @@ outer:
 				break outer
 			}
 
-			content := strings.TrimSuffix(l.Content, "\n")
-			content = ansi.GraphemeWidth.Cut(content, dv.xOffset, len(content))
-			content = ansi.Truncate(content, dv.codeWidth, "…")
+			getContent := func(ls LineStyle) string {
+				content := strings.TrimSuffix(l.Content, "\n")
+				content = dv.hightlightCode(content, ls.Code.GetBackground())
+				content = ansi.GraphemeWidth.Cut(content, dv.xOffset, len(content))
+				content = ansi.Truncate(content, dv.codeWidth, "…")
+				return content
+			}
 
 			leadingEllipsis := dv.xOffset > 0 && strings.TrimSpace(content) != ""
 
 			switch l.Kind {
 			case udiff.Equal:
+				content := getContent(dv.style.EqualLine)
 				if dv.lineNumbers {
 					write(dv.style.EqualLine.LineNumber.Render(pad(beforeLine, dv.beforeNumDigits)))
 					write(dv.style.EqualLine.LineNumber.Render(pad(afterLine, dv.afterNumDigits)))
@@ -379,6 +395,7 @@ outer:
 				beforeLine++
 				afterLine++
 			case udiff.Insert:
+				content := getContent(dv.style.InsertLine)
 				if dv.lineNumbers {
 					write(dv.style.InsertLine.LineNumber.Render(pad(" ", dv.beforeNumDigits)))
 					write(dv.style.InsertLine.LineNumber.Render(pad(afterLine, dv.afterNumDigits)))
@@ -389,6 +406,7 @@ outer:
 				))
 				afterLine++
 			case udiff.Delete:
+				content := getContent(dv.style.DeleteLine)
 				if dv.lineNumbers {
 					write(dv.style.DeleteLine.LineNumber.Render(pad(beforeLine, dv.beforeNumDigits)))
 					write(dv.style.DeleteLine.LineNumber.Render(pad(" ", dv.afterNumDigits)))
@@ -479,22 +497,17 @@ outer:
 				break outer
 			}
 
-			var beforeContent string
-			var afterContent string
-			if l.before != nil {
-				beforeContent = strings.TrimSuffix(l.before.Content, "\n")
-				beforeContent = ansi.GraphemeWidth.Cut(beforeContent, dv.xOffset, len(beforeContent))
-				beforeContent = ansi.Truncate(beforeContent, dv.codeWidth, "…")
+			getContent := func(content string, ls LineStyle) string {
+				content = strings.TrimSuffix(content, "\n")
+				content = dv.hightlightCode(content, ls.Code.GetBackground())
+				content = ansi.GraphemeWidth.Cut(content, dv.xOffset, len(content))
+				content = ansi.Truncate(content, dv.codeWidth, "…")
+				return content
 			}
-			if l.after != nil {
-				afterContent = strings.TrimSuffix(l.after.Content, "\n")
-				afterContent = ansi.GraphemeWidth.Cut(afterContent, dv.xOffset, len(afterContent))
-				afterContent = ansi.Truncate(afterContent, dv.codeWidth+btoi(dv.extraColOnAfter), "…")
+			getLeadingEllipsis := func(content string) bool {
+				return dv.xOffset > 0 && strings.TrimSpace(content) != ""
 			}
 
-			leadingBeforeEllipsis := dv.xOffset > 0 && strings.TrimSpace(beforeContent) != ""
-			leadingAfterEllipsis := dv.xOffset > 0 && strings.TrimSpace(afterContent) != ""
-
 			switch {
 			case l.before == nil:
 				if dv.lineNumbers {
@@ -504,20 +517,24 @@ outer:
 					dv.style.MissingLine.Code.Width(dv.fullCodeWidth).Render("  "),
 				))
 			case l.before.Kind == udiff.Equal:
+				content := getContent(l.before.Content, dv.style.EqualLine)
+				leadingEllipsis := getLeadingEllipsis(content)
 				if dv.lineNumbers {
 					write(dv.style.EqualLine.LineNumber.Render(pad(beforeLine, dv.beforeNumDigits)))
 				}
 				write(beforeFullContentStyle.Render(
-					dv.style.EqualLine.Code.Width(dv.fullCodeWidth).Render(ternary(leadingBeforeEllipsis, " …", "  ") + beforeContent),
+					dv.style.EqualLine.Code.Width(dv.fullCodeWidth).Render(ternary(leadingEllipsis, " …", "  ") + content),
 				))
 				beforeLine++
 			case l.before.Kind == udiff.Delete:
+				content := getContent(l.before.Content, dv.style.DeleteLine)
+				leadingEllipsis := getLeadingEllipsis(content)
 				if dv.lineNumbers {
 					write(dv.style.DeleteLine.LineNumber.Render(pad(beforeLine, dv.beforeNumDigits)))
 				}
 				write(beforeFullContentStyle.Render(
-					dv.style.DeleteLine.Symbol.Render(ternary(leadingBeforeEllipsis, "-…", "- ")) +
-						dv.style.DeleteLine.Code.Width(dv.codeWidth).Render(beforeContent),
+					dv.style.DeleteLine.Symbol.Render(ternary(leadingEllipsis, "-…", "- ")) +
+						dv.style.DeleteLine.Code.Width(dv.codeWidth).Render(content),
 				))
 				beforeLine++
 			}
@@ -531,20 +548,24 @@ outer:
 					dv.style.MissingLine.Code.Width(dv.fullCodeWidth + btoi(dv.extraColOnAfter)).Render("  "),
 				))
 			case l.after.Kind == udiff.Equal:
+				content := getContent(l.after.Content, dv.style.EqualLine)
+				leadingEllipsis := getLeadingEllipsis(content)
 				if dv.lineNumbers {
 					write(dv.style.EqualLine.LineNumber.Render(pad(afterLine, dv.afterNumDigits)))
 				}
 				write(afterFullContentStyle.Render(
-					dv.style.EqualLine.Code.Width(dv.fullCodeWidth + btoi(dv.extraColOnAfter)).Render(ternary(leadingAfterEllipsis, " …", "  ") + afterContent),
+					dv.style.EqualLine.Code.Width(dv.fullCodeWidth + btoi(dv.extraColOnAfter)).Render(ternary(leadingEllipsis, " …", "  ") + content),
 				))
 				afterLine++
 			case l.after.Kind == udiff.Insert:
+				content := getContent(l.after.Content, dv.style.InsertLine)
+				leadingEllipsis := getLeadingEllipsis(content)
 				if dv.lineNumbers {
 					write(dv.style.InsertLine.LineNumber.Render(pad(afterLine, dv.afterNumDigits)))
 				}
 				write(afterFullContentStyle.Render(
-					dv.style.InsertLine.Symbol.Render(ternary(leadingAfterEllipsis, "+…", "+ ")) +
-						dv.style.InsertLine.Code.Width(dv.codeWidth+btoi(dv.extraColOnAfter)).Render(afterContent),
+					dv.style.InsertLine.Symbol.Render(ternary(leadingEllipsis, "+…", "+ ")) +
+						dv.style.InsertLine.Code.Width(dv.codeWidth+btoi(dv.extraColOnAfter)).Render(content),
 				))
 				afterLine++
 			}
@@ -613,3 +634,40 @@ func (dv *DiffView) lineStyleForType(t udiff.OpKind) LineStyle {
 		return dv.style.MissingLine
 	}
 }
+
+func (dv *DiffView) hightlightCode(source string, bgColor color.Color) string {
+	if dv.chromaStyle == nil {
+		return source
+	}
+
+	l := dv.getChromaLexer(source)
+	f := dv.getChromaFormatter(bgColor)
+
+	it, err := l.Tokenise(nil, source)
+	if err != nil {
+		return source
+	}
+
+	var b strings.Builder
+	if err := f.Format(&b, dv.chromaStyle, it); err != nil {
+		return source
+	}
+	return b.String()
+}
+
+func (dv *DiffView) getChromaLexer(source string) chroma.Lexer {
+	l := lexers.Match(dv.before.path)
+	if l == nil {
+		l = lexers.Analyse(source)
+	}
+	if l == nil {
+		l = lexers.Fallback
+	}
+	return chroma.Coalesce(l)
+}
+
+func (dv *DiffView) getChromaFormatter(gbColor color.Color) chroma.Formatter {
+	return chromaFormatter{
+		bgColor: gbColor,
+	}
+}

internal/exp/diffview/diffview_test.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/alecthomas/chroma/v2/styles"
 	"github.com/charmbracelet/x/ansi"
 	"github.com/charmbracelet/x/exp/golden"
 	"github.com/opencode-ai/opencode/internal/exp/diffview"
@@ -87,12 +88,22 @@ var (
 			After("main.go", TestMultipleHunksAfter).
 			Width(120)
 	}
+	NoSyntaxHighlightFunc = func(dv *diffview.DiffView) *diffview.DiffView {
+		return dv.
+			Before("main.go", TestMultipleHunksBefore).
+			After("main.go", TestMultipleHunksAfter).
+			ChromaStyle(nil)
+	}
 
 	LightModeFunc = func(dv *diffview.DiffView) *diffview.DiffView {
-		return dv.Style(diffview.DefaultLightStyle)
+		return dv.
+			Style(diffview.DefaultLightStyle).
+			ChromaStyle(styles.Get("catppuccin-latte"))
 	}
 	DarkModeFunc = func(dv *diffview.DiffView) *diffview.DiffView {
-		return dv.Style(diffview.DefaultDarkStyle)
+		return dv.
+			Style(diffview.DefaultDarkStyle).
+			ChromaStyle(styles.Get("catppuccin-macchiato"))
 	}
 
 	LayoutFuncs = TestFuncs{
@@ -107,6 +118,7 @@ var (
 		"Narrow":             NarrowFunc,
 		"SmallWidth":         SmallWidthFunc,
 		"LargeWidth":         LargeWidthFunc,
+		"NoSyntaxHighlight":  NoSyntaxHighlightFunc,
 	}
 	ThemeFuncs = TestFuncs{
 		"LightMode": LightModeFunc,
@@ -123,8 +135,8 @@ func TestDiffView(t *testing.T) {
 						t.Run(themeName, func(t *testing.T) {
 							dv := diffview.New()
 							dv = layoutFunc(dv)
-							dv = behaviorFunc(dv)
 							dv = themeFunc(dv)
+							dv = behaviorFunc(dv)
 
 							output := dv.String()
 							golden.RequireEqual(t, []byte(output))
@@ -149,7 +161,8 @@ func TestDiffViewTabs(t *testing.T) {
 			dv := diffview.New().
 				Before("main.go", TestTabsBefore).
 				After("main.go", TestTabsAfter).
-				Style(diffview.DefaultLightStyle)
+				Style(diffview.DefaultLightStyle).
+				ChromaStyle(styles.Get("catppuccin-latte"))
 			dv = layoutFunc(dv)
 
 			output := dv.String()
@@ -171,7 +184,8 @@ func TestDiffViewWidth(t *testing.T) {
 						Before("main.go", TestMultipleHunksBefore).
 						After("main.go", TestMultipleHunksAfter).
 						Width(width).
-						Style(diffview.DefaultLightStyle)
+						Style(diffview.DefaultLightStyle).
+						ChromaStyle(styles.Get("catppuccin-latte"))
 					dv = layoutFunc(dv)
 
 					output := dv.String()
@@ -193,7 +207,8 @@ func TestDiffViewHeight(t *testing.T) {
 						Before("main.go", TestMultipleHunksBefore).
 						After("main.go", TestMultipleHunksAfter).
 						Height(height).
-						Style(diffview.DefaultLightStyle)
+						Style(diffview.DefaultLightStyle).
+						ChromaStyle(styles.Get("catppuccin-latte"))
 					dv = layoutFunc(dv)
 
 					output := dv.String()
@@ -215,6 +230,7 @@ func TestDiffViewXOffset(t *testing.T) {
 						Before("main.go", TestDefaultBefore).
 						After("main.go", TestDefaultAfter).
 						Style(diffview.DefaultLightStyle).
+						ChromaStyle(styles.Get("catppuccin-latte")).
 						Width(60).
 						XOffset(xOffset)
 					dv = layoutFunc(dv)
@@ -238,6 +254,7 @@ func TestDiffViewYOffset(t *testing.T) {
 						Before("main.go", TestMultipleHunksBefore).
 						After("main.go", TestMultipleHunksAfter).
 						Style(diffview.DefaultLightStyle).
+						ChromaStyle(styles.Get("catppuccin-latte")).
 						Height(5).
 						YOffset(yOffset)
 					dv = layoutFunc(dv)