diple/suggestions.go

package main

import (
	"strings"
	"sync"

	"github.com/charmbracelet/x/ansi"
)

const suggestionRenderCacheLimit = 256

var renderedSuggestions = newSuggestionRenderCache(suggestionRenderCacheLimit)

type parsedCommentBody struct {
	Prose       string
	Suggestions []string
}

type codeRange struct {
	Start int
	End   int
}

type suggestionRenderCacheKey struct {
	path         string
	removed      string
	removedLines int
	replacement  string
	width        int
}

type suggestionRenderCacheEntry struct {
	removed []detailLine
	added   []detailLine
}

type suggestionRenderCache struct {
	mu      sync.Mutex
	limit   int
	entries map[suggestionRenderCacheKey]suggestionRenderCacheEntry
	order   []suggestionRenderCacheKey
}

func newSuggestionRenderCache(limit int) *suggestionRenderCache {
	return &suggestionRenderCache{
		limit: limit, entries: make(map[suggestionRenderCacheKey]suggestionRenderCacheEntry),
	}
}

func (c *suggestionRenderCache) get(
	key suggestionRenderCacheKey,
) (suggestionRenderCacheEntry, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	entry, ok := c.entries[key]
	return cloneSuggestionRenderEntry(entry), ok
}

func (c *suggestionRenderCache) put(
	key suggestionRenderCacheKey, entry suggestionRenderCacheEntry,
) suggestionRenderCacheEntry {
	c.mu.Lock()
	defer c.mu.Unlock()
	if cached, ok := c.entries[key]; ok {
		return cloneSuggestionRenderEntry(cached)
	}
	if len(c.entries) >= c.limit {
		delete(c.entries, c.order[0])
		c.order = c.order[1:]
	}
	c.entries[key] = cloneSuggestionRenderEntry(entry)
	c.order = append(c.order, key)
	return cloneSuggestionRenderEntry(entry)
}

func (c *suggestionRenderCache) clear() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.entries = make(map[suggestionRenderCacheKey]suggestionRenderCacheEntry)
	c.order = nil
}

func cloneSuggestionRenderEntry(entry suggestionRenderCacheEntry) suggestionRenderCacheEntry {
	return suggestionRenderCacheEntry{
		removed: append([]detailLine(nil), entry.removed...),
		added:   append([]detailLine(nil), entry.added...),
	}
}

func renderSuggestion(
	path string, reviewed []string, replacement string, width int,
) suggestionRenderCacheEntry {
	key := suggestionRenderCacheKey{
		path: path, removed: strings.Join(reviewed, "\n"), removedLines: len(reviewed),
		replacement: replacement, width: width,
	}
	if cached, ok := renderedSuggestions.get(key); ok {
		return cached
	}
	removed, added := normalizeSuggestion(reviewed, replacement)
	removedRanges, addedRanges := suggestionChangedRanges(removed, added)
	entry := suggestionRenderCacheEntry{}
	for index, source := range removed {
		entry.removed = append(entry.removed, wrapSuggestionLine(
			path, source, '-', removedRanges[index], width,
		)...)
	}
	for index, source := range added {
		entry.added = append(entry.added, wrapSuggestionLine(
			path, source, '+', addedRanges[index], width,
		)...)
	}
	return renderedSuggestions.put(key, entry)
}

func parseCommentBody(body string) parsedCommentBody {
	var (
		result       parsedCommentBody
		prose        []string
		suggestion   []string
		openingLine  string
		inSuggestion bool
	)
	for _, line := range strings.Split(body, "\n") {
		trimmed := strings.TrimSpace(line)
		if !inSuggestion && isSuggestionFence(trimmed) {
			inSuggestion = true
			openingLine = line
			suggestion = nil
			continue
		}
		if inSuggestion && trimmed == "```" {
			result.Suggestions = append(result.Suggestions, strings.Join(suggestion, "\n"))
			inSuggestion = false
			openingLine = ""
			continue
		}
		if inSuggestion {
			suggestion = append(suggestion, line)
		} else {
			prose = append(prose, line)
		}
	}
	if inSuggestion {
		prose = append(prose, openingLine)
		prose = append(prose, suggestion...)
	}
	result.Prose = strings.TrimSpace(strings.Join(prose, "\n"))
	return result
}

func isSuggestionFence(line string) bool {
	if !strings.HasPrefix(line, "```suggestion") {
		return false
	}
	suffix := strings.TrimPrefix(line, "```suggestion")
	return suffix == "" || strings.HasPrefix(suffix, ":")
}

func reviewedSourceLines(thread ReviewThread, comment ReviewComment) []string {
	start, end := commentReviewAnchor(thread, comment)
	parsed := parseDiff(comment.DiffHunk, start, end, thread.DiffSide)
	lines := make([]string, 0, end-start+1)
	for _, line := range parsed {
		if !line.selected {
			continue
		}
		lines = append(lines, sourceFromDiffLine(line.raw))
	}
	return lines
}

func sourceFromDiffLine(raw string) string {
	if strings.HasPrefix(raw, "+") || strings.HasPrefix(raw, "-") || strings.HasPrefix(raw, " ") {
		raw = raw[1:]
	}
	return strings.ReplaceAll(raw, "\t", "    ")
}

func normalizeSuggestion(removed []string, replacement string) ([]string, []string) {
	added := []string(nil)
	if replacement != "" {
		added = strings.Split(replacement, "\n")
		for i := range added {
			added[i] = strings.ReplaceAll(added[i], "\t", "    ")
		}
	}
	padding := commonSourceIndent(removed, added)
	return trimSourceIndent(removed, padding), trimSourceIndent(added, padding)
}

func commonSourceIndent(groups ...[]string) int {
	padding := -1
	for _, lines := range groups {
		for _, line := range lines {
			if strings.TrimSpace(line) == "" {
				continue
			}
			indent := len(line) - len(strings.TrimLeft(line, " "))
			if padding == -1 || indent < padding {
				padding = indent
			}
		}
	}
	return max(0, padding)
}

func trimSourceIndent(lines []string, padding int) []string {
	result := make([]string, len(lines))
	for i, line := range lines {
		result[i] = trimIndent(line, padding)
	}
	return result
}

func suggestionChangedRanges(removed, added []string) ([]codeRange, []codeRange) {
	removedRanges := make([]codeRange, len(removed))
	addedRanges := make([]codeRange, len(added))
	matches := matchingLinePairs(removed, added)
	matches = append(matches, [2]int{len(removed), len(added)})

	oldStart, newStart := 0, 0
	for _, match := range matches {
		oldEnd, newEnd := match[0], match[1]
		blockSize := max(oldEnd-oldStart, newEnd-newStart)
		for offset := 0; offset < blockSize; offset++ {
			oldIndex, newIndex := oldStart+offset, newStart+offset
			switch {
			case oldIndex < oldEnd && newIndex < newEnd:
				removedRanges[oldIndex], addedRanges[newIndex] = changedRange(
					removed[oldIndex], added[newIndex],
				)
			case oldIndex < oldEnd:
				removedRanges[oldIndex] = fullCodeRange(removed[oldIndex])
			case newIndex < newEnd:
				addedRanges[newIndex] = fullCodeRange(added[newIndex])
			}
		}
		if oldEnd < len(removed) && newEnd < len(added) {
			oldStart, newStart = oldEnd+1, newEnd+1
		} else {
			oldStart, newStart = oldEnd, newEnd
		}
	}
	return removedRanges, addedRanges
}

func matchingLinePairs(removed, added []string) [][2]int {
	dp := make([][]int, len(removed)+1)
	for i := range dp {
		dp[i] = make([]int, len(added)+1)
	}
	for i := len(removed) - 1; i >= 0; i-- {
		for j := len(added) - 1; j >= 0; j-- {
			if removed[i] == added[j] {
				dp[i][j] = dp[i+1][j+1] + 1
			} else {
				dp[i][j] = max(dp[i+1][j], dp[i][j+1])
			}
		}
	}

	var matches [][2]int
	for i, j := 0, 0; i < len(removed) && j < len(added); {
		switch {
		case removed[i] == added[j]:
			matches = append(matches, [2]int{i, j})
			i++
			j++
		case dp[i+1][j] >= dp[i][j+1]:
			i++
		default:
			j++
		}
	}
	return matches
}

func changedRange(oldLine, newLine string) (codeRange, codeRange) {
	oldRunes, newRunes := []rune(oldLine), []rune(newLine)
	prefix := 0
	for prefix < len(oldRunes) && prefix < len(newRunes) && oldRunes[prefix] == newRunes[prefix] {
		prefix++
	}
	oldSuffix, newSuffix := len(oldRunes), len(newRunes)
	for oldSuffix > prefix && newSuffix > prefix &&
		oldRunes[oldSuffix-1] == newRunes[newSuffix-1] {
		oldSuffix--
		newSuffix--
	}
	return codeRange{
			Start: ansi.StringWidth(string(oldRunes[:prefix])),
			End:   ansi.StringWidth(string(oldRunes[:oldSuffix])),
		}, codeRange{
			Start: ansi.StringWidth(string(newRunes[:prefix])),
			End:   ansi.StringWidth(string(newRunes[:newSuffix])),
		}
}

func fullCodeRange(line string) codeRange {
	return codeRange{End: ansi.StringWidth(line)}
}

func commentReviewAnchor(thread ReviewThread, comment ReviewComment) (int, int) {
	end := comment.OriginalLine
	start := comment.OriginalStartLine
	if end == 0 {
		end = comment.Line
		start = comment.StartLine
	}
	if end > 0 {
		if start == 0 {
			start = end
		}
		return start, end
	}
	start = thread.StartLine
	if start == 0 {
		start = thread.Line
	}
	return start, thread.Line
}