lazygit/pkg/utils/lines.go

package utils

import (
	"bytes"
	"strings"

	"github.com/rivo/uniseg"
)

// SplitLines takes a multiline string and splits it on newlines
// currently we are also stripping \r's which may have adverse effects for
// windows users (but no issues have been raised yet)
func SplitLines(multilineString string) []string {
	multilineString = strings.ReplaceAll(multilineString, "\r", "")
	if multilineString == "" || multilineString == "\n" {
		return make([]string, 0)
	}
	lines := strings.Split(multilineString, "\n")
	if lines[len(lines)-1] == "" {
		return lines[:len(lines)-1]
	}
	return lines
}

func SplitNul(str string) []string {
	if str == "" {
		return make([]string, 0)
	}
	str = strings.TrimSuffix(str, "\x00")
	return strings.Split(str, "\x00")
}

// NormalizeLinefeeds - Removes all Windows and Mac style line feeds
func NormalizeLinefeeds(str string) string {
	str = strings.ReplaceAll(str, "\r\n", "\n")
	str = strings.ReplaceAll(str, "\r", "")
	return str
}

// EscapeSpecialChars - Replaces all special chars like \n with \\n
func EscapeSpecialChars(str string) string {
	return strings.NewReplacer(
		"\n", "\\n",
		"\r", "\\r",
		"\t", "\\t",
		"\b", "\\b",
		"\f", "\\f",
		"\v", "\\v",
	).Replace(str)
}

func dropCR(data []byte) []byte {
	if len(data) > 0 && data[len(data)-1] == '\r' {
		return data[0 : len(data)-1]
	}
	return data
}

// ScanLinesAndTruncateWhenLongerThanBuffer returns a split function that can be
// used with bufio.Scanner.Split(). It is very similar to bufio.ScanLines,
// except that it will truncate lines that are longer than the scanner's read
// buffer (whereas bufio.ScanLines will return an error in that case, which is
// often difficult to handle).
//
// If you are using your own buffer for the scanner, you must set maxBufferSize
// to the same value as the max parameter that you passed to scanner.Buffer().
// Otherwise, maxBufferSize must be set to bufio.MaxScanTokenSize.
func ScanLinesAndTruncateWhenLongerThanBuffer(maxBufferSize int) func(data []byte, atEOF bool) (int, []byte, error) {
	skipOverRemainderOfLongLine := false

	return func(data []byte, atEOF bool) (int, []byte, error) {
		if atEOF && len(data) == 0 {
			// Done
			return 0, nil, nil
		}
		if i := bytes.IndexByte(data, '\n'); i >= 0 {
			if skipOverRemainderOfLongLine {
				skipOverRemainderOfLongLine = false
				return i + 1, nil, nil
			}
			return i + 1, dropCR(data[0:i]), nil
		}
		if atEOF {
			if skipOverRemainderOfLongLine {
				return len(data), nil, nil
			}

			return len(data), dropCR(data), nil
		}

		// Buffer is full, so we can't get more data
		if len(data) >= maxBufferSize {
			if skipOverRemainderOfLongLine {
				return len(data), nil, nil
			}

			skipOverRemainderOfLongLine = true
			return len(data), data, nil
		}

		// Request more data.
		return 0, nil, nil
	}
}

type escapeSequenceState int

const (
	notInEscapeSequence escapeSequenceState = iota
	afterEscapeChar
	inCSISequence
)

// Recognizes ANSI escape sequences while iterating over the characters of a
// string. Only CSI sequences (ESC [ ... final byte) are recognized; these are
// the color and style codes that lazygit puts into view content.
type escapeSequenceTracker struct {
	state escapeSequenceState
}

// Feeds the next character to the tracker. Returns true if the character is
// part of an escape sequence, and so takes up no space on screen.
func (self *escapeSequenceTracker) consume(chr rune) bool {
	switch self.state {
	case afterEscapeChar:
		if chr == '[' {
			self.state = inCSISequence
		} else {
			self.state = notInEscapeSequence
		}
		return true
	case inCSISequence:
		// Parameter and intermediate bytes are in the range 0x20-0x3F, the
		// final byte is in the range 0x40-0x7E
		if chr >= 0x40 && chr <= 0x7e {
			self.state = notInEscapeSequence
		}
		return true
	default:
		if chr == '\x1b' {
			self.state = afterEscapeChar
			return true
		}
		return false
	}
}

// The width that the string takes up on screen, leaving out escape sequences
func widthIgnoringEscapeSequences(s string) int {
	tracker := escapeSequenceTracker{}
	width := 0
	for _, chr := range s {
		if !tracker.consume(chr) {
			width += uniseg.StringWidth(string(chr))
		}
	}
	return width
}

// Wrap lines to a given width, and return:
// - the wrapped lines
// - the line indices of the wrapped lines, indexed by the original line indices
// - the line indices of the original lines, indexed by the wrapped line indices
// If wrap is false, the text is returned as is.
// Escape sequences stay in the wrapped lines, but don't count towards the
// width, and a line is never broken inside one.
// This code needs to behave the same as `gocui.lineWrap` does.
func WrapViewLinesToWidth(wrap bool, editable bool, text string, width int, tabWidth int) ([]string, []int, []int) {
	if !editable {
		text = strings.TrimSuffix(text, "\n")
	}
	lines := strings.Split(text, "\n")
	if !wrap {
		indices := make([]int, len(lines))
		for i := range lines {
			indices[i] = i
		}
		return lines, indices, indices
	}

	wrappedLines := make([]string, 0, len(lines))
	wrappedLineIndices := make([]int, 0, len(lines))
	originalLineIndices := make([]int, 0, len(lines))

	if tabWidth < 1 {
		tabWidth = 4
	}

	for originalLineIdx, line := range lines {
		wrappedLineIndices = append(wrappedLineIndices, len(wrappedLines))

		// convert tabs to spaces
		for i := 0; i < len(line); i++ {
			if line[i] == '\t' {
				numSpaces := tabWidth - (i % tabWidth)
				line = line[:i] + strings.Repeat(" ", numSpaces) + line[i+1:]
				i += numSpaces - 1
			}
		}

		appendWrappedLine := func(str string) {
			wrappedLines = append(wrappedLines, str)
			originalLineIndices = append(originalLineIndices, originalLineIdx)
		}

		n := 0
		offset := 0
		lastWhitespaceIndex := -1
		escapeSequences := escapeSequenceTracker{}
		for i, currChr := range line {
			if escapeSequences.consume(currChr) {
				continue
			}

			rw := uniseg.StringWidth(string(currChr))
			n += rw

			if n > width {
				if currChr == ' ' {
					appendWrappedLine(line[offset:i])
					offset = i + 1
					n = 0
				} else if currChr == '-' {
					appendWrappedLine(line[offset:i])
					offset = i
					n = rw
				} else if lastWhitespaceIndex != -1 {
					if line[lastWhitespaceIndex] == '-' {
						appendWrappedLine(line[offset : lastWhitespaceIndex+1])
					} else {
						appendWrappedLine(line[offset:lastWhitespaceIndex])
					}
					offset = lastWhitespaceIndex + 1
					n = widthIgnoringEscapeSequences(line[offset : i+1])
				} else {
					appendWrappedLine(line[offset:i])
					offset = i
					n = rw
				}
				lastWhitespaceIndex = -1
			} else if currChr == ' ' || currChr == '-' {
				lastWhitespaceIndex = i
			}
		}

		appendWrappedLine(line[offset:])
	}

	return wrappedLines, wrappedLineIndices, originalLineIndices
}