diple/highlight.go

package main

import (
	"bytes"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"sync"

	"github.com/alecthomas/chroma/v2/lexers"
	"github.com/alecthomas/chroma/v2/quick"
)

const reviewContextLines = 3
const highlightedDiffCacheLimit = 256

var codeHighlightTheme = "github-dark"
var colorEnabled = true

var hunkHeaderPattern = regexp.MustCompile(`^@@ -(\d+)(?:,\d+)? \+(\d+)(?:,\d+)? @@`)

type highlightedDiffCacheKey struct {
	path, hunk, side, theme string
	startLine, endLine      int
	color                   bool
}

type highlightedDiffCache struct {
	mu      sync.Mutex
	entries map[highlightedDiffCacheKey][]highlightedDiffLine
	order   []highlightedDiffCacheKey
}

var highlightedDiffs = highlightedDiffCache{
	entries: make(map[highlightedDiffCacheKey][]highlightedDiffLine),
}

func (c *highlightedDiffCache) get(key highlightedDiffCacheKey) ([]highlightedDiffLine, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	lines, ok := c.entries[key]
	return append([]highlightedDiffLine(nil), lines...), ok
}

func (c *highlightedDiffCache) put(
	key highlightedDiffCacheKey, lines []highlightedDiffLine,
) []highlightedDiffLine {
	c.mu.Lock()
	defer c.mu.Unlock()
	if cached, ok := c.entries[key]; ok {
		return append([]highlightedDiffLine(nil), cached...)
	}
	if len(c.entries) >= highlightedDiffCacheLimit {
		delete(c.entries, c.order[0])
		c.order = c.order[1:]
	}
	c.entries[key] = append([]highlightedDiffLine(nil), lines...)
	c.order = append(c.order, key)
	return append([]highlightedDiffLine(nil), lines...)
}

func (c *highlightedDiffCache) clear() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.entries = make(map[highlightedDiffCacheKey][]highlightedDiffLine)
	c.order = nil
}

type highlightedDiffLine struct {
	gutter   string
	code     string
	selected bool
}

type parsedDiffLine struct {
	raw      string
	oldLine  int
	newLine  int
	header   bool
	notice   bool
	selected bool
}

func highlightDiff(path, hunk string, startLine, endLine int, side string) []highlightedDiffLine {
	key := highlightedDiffCacheKey{
		path: path, hunk: hunk, side: side, theme: codeHighlightTheme,
		startLine: startLine, endLine: endLine, color: colorEnabled,
	}
	if cached, ok := highlightedDiffs.get(key); ok {
		return cached
	}
	return highlightedDiffs.put(key, highlightDiffUncached(path, hunk, startLine, endLine, side))
}

func highlightDiffUncached(path, hunk string, startLine, endLine int, side string) []highlightedDiffLine {
	if hunk == "" {
		return []highlightedDiffLine{{code: "(GitHub did not return a diff hunk)"}}
	}
	if endLine <= 0 {
		endLine = startLine
	}
	if startLine <= 0 {
		startLine = endLine
	}
	if startLine > endLine {
		startLine, endLine = endLine, startLine
	}

	parsed := parseDiff(hunk, startLine, endLine, side)
	visible := reviewedWindow(parsed)
	padding := commonIndent(visible)
	lexer := lexerForPath(path)
	out := make([]highlightedDiffLine, 0, len(visible))
	for _, line := range visible {
		if line.raw == "⋯" {
			code := "⋯"
			if colorEnabled {
				code = "\x1b[38;5;245m⋯\x1b[0m"
			}
			out = append(out, highlightedDiffLine{code: code})
			continue
		}
		if line.notice {
			code := line.raw
			if colorEnabled {
				code = "\x1b[38;5;245m" + line.raw + "\x1b[0m"
			}
			out = append(out, highlightedDiffLine{code: code})
			continue
		}
		if line.header {
			code := line.raw
			if colorEnabled {
				code = "\x1b[38;5;141m" + line.raw + "\x1b[0m"
			}
			out = append(out, highlightedDiffLine{code: code})
			continue
		}
		out = append(out, renderDiffLine(lexer, line, padding))
	}
	return out
}

func parseDiff(hunk string, startLine, endLine int, side string) []parsedDiffLine {
	oldLine, newLine := 0, 0
	lines := make([]parsedDiffLine, 0, strings.Count(hunk, "\n")+1)
	for _, raw := range strings.Split(hunk, "\n") {
		if match := hunkHeaderPattern.FindStringSubmatch(raw); match != nil {
			oldLine, _ = strconv.Atoi(match[1])
			newLine, _ = strconv.Atoi(match[2])
			lines = append(lines, parsedDiffLine{raw: raw, header: true})
			continue
		}

		line := parsedDiffLine{raw: raw}
		switch {
		case strings.HasPrefix(raw, "+"):
			line.newLine = newLine
			newLine++
		case strings.HasPrefix(raw, "-"):
			line.oldLine = oldLine
			oldLine++
		case strings.HasPrefix(raw, `\`):
			// "\ No newline at end of file" has no source coordinate.
		default:
			line.oldLine, line.newLine = oldLine, newLine
			oldLine++
			newLine++
		}
		coordinate := line.newLine
		if side == "LEFT" {
			coordinate = line.oldLine
		}
		line.selected = coordinate > 0 && coordinate >= startLine && coordinate <= endLine
		lines = append(lines, line)
	}
	return lines
}

func reviewedWindow(lines []parsedDiffLine) []parsedDiffLine {
	first, last := -1, -1
	for i, line := range lines {
		if line.selected {
			if first == -1 {
				first = i
			}
			last = i
		}
	}
	if first == -1 {
		out := make([]parsedDiffLine, 0, 2)
		for _, line := range lines {
			if line.header {
				out = append(out, line)
				break
			}
		}
		return append(out, parsedDiffLine{
			raw:    "(original reviewed lines unavailable in GitHub's historical diff)",
			notice: true,
		})
	}

	start := max(0, first-reviewContextLines)
	end := min(len(lines), last+reviewContextLines+1)
	header := -1
	for i := first; i >= 0; i-- {
		if lines[i].header {
			header = i
			break
		}
	}

	out := make([]parsedDiffLine, 0, end-start+3)
	if header >= 0 && header < start {
		out = append(out, lines[header])
	}
	if start > 0 && start != header {
		out = append(out, parsedDiffLine{raw: "⋯"})
	}
	out = append(out, lines[start:end]...)
	if end < len(lines) {
		out = append(out, parsedDiffLine{raw: "⋯"})
	}
	return out
}

func renderDiffLine(lexer string, line parsedDiffLine, padding int) highlightedDiffLine {
	marker, source := " ", line.raw
	lineNumber := ""
	markerStyle := "\x1b[38;5;245m"
	switch {
	case strings.HasPrefix(line.raw, "+"):
		marker, source = "+", strings.TrimPrefix(line.raw, "+")
		lineNumber = coordinateText(line.newLine)
		markerStyle = "\x1b[38;5;114m"
	case strings.HasPrefix(line.raw, "-"):
		marker, source = "-", strings.TrimPrefix(line.raw, "-")
		lineNumber = coordinateText(line.oldLine)
		markerStyle = "\x1b[38;5;203m"
	default:
		source = strings.TrimPrefix(line.raw, " ")
		lineNumber = coordinateText(line.newLine)
		if lineNumber == "" {
			lineNumber = coordinateText(line.oldLine)
		}
	}

	source = strings.ReplaceAll(source, "\t", "    ")
	source = trimIndent(source, padding)
	gutter := fmt.Sprintf("%5s %s ", lineNumber, marker)
	if colorEnabled {
		gutter = fmt.Sprintf(
			"\x1b[38;5;245m%5s\x1b[0m %s%s\x1b[0m ",
			lineNumber, markerStyle, marker,
		)
	}
	return highlightedDiffLine{
		gutter:   gutter,
		code:     highlightedSource(lexer, source),
		selected: line.selected,
	}
}

func commonIndent(lines []parsedDiffLine) int {
	padding := -1
	for _, line := range lines {
		if line.header || line.notice || line.raw == "⋯" || strings.HasPrefix(line.raw, `\`) {
			continue
		}
		source := line.raw
		if strings.HasPrefix(source, "+") || strings.HasPrefix(source, "-") || strings.HasPrefix(source, " ") {
			source = source[1:]
		}
		source = strings.ReplaceAll(source, "\t", "    ")
		if strings.TrimSpace(source) == "" {
			continue
		}
		indent := len(source) - len(strings.TrimLeft(source, " "))
		if padding == -1 || indent < padding {
			padding = indent
		}
	}
	return max(0, padding)
}

func trimIndent(source string, padding int) string {
	for padding > 0 && strings.HasPrefix(source, " ") {
		source = source[1:]
		padding--
	}
	return source
}

func coordinateText(line int) string {
	if line <= 0 {
		return ""
	}
	return strconv.Itoa(line)
}

func highlightedSource(lexer, source string) string {
	if !colorEnabled {
		return source
	}
	var highlighted bytes.Buffer
	if err := quick.Highlight(&highlighted, source, lexer, "terminal16m", codeHighlightTheme); err != nil {
		return source
	}
	return strings.TrimSuffix(highlighted.String(), "\n")
}

func lexerForPath(path string) string {
	lexer := lexers.Match(path)
	if lexer == nil {
		return ""
	}
	return lexer.Config().Name
}