ui: parse strings for ansi styles

Parse UI strings for ANSI styles. If there are styles in the string, use
those as the display style in tcell. This is in preparation for a
template function which can apply arbitrary styles to UI elements.

Signed-off-by: Tim Culverhouse <tim@timculverhouse.com>
Acked-by: Robin Jarry <robin@jarry.cc>

2 files changed, 662 insertions, 2 deletions

diff --git a/lib/parse/ansi.go b/lib/parse/ansi.go
index e1d603ea..a9a46fdd 100644
--- a/lib/parse/ansi.go
+++ b/lib/parse/ansi.go
@@ -3,15 +3,21 @@ package parse
 import (
 	"bufio"
 	"bytes"
+	"errors"
 	"fmt"
 	"io"
 	"os"
 	"regexp"
+	"strconv"
+	"strings"
 
 	"git.sr.ht/~rjarry/aerc/log"
+	"github.com/gdamore/tcell/v2"
+	"github.com/gdamore/tcell/v2/terminfo"
+	"github.com/mattn/go-runewidth"
 )
 
-var ansi = regexp.MustCompile("\x1B\\[[0-?]*[ -/]*[@-~]")
+var AnsiReg = regexp.MustCompile("\x1B\\[[0-?]*[ -/]*[@-~]")
 
 // StripAnsi strips ansi escape codes from the reader
 func StripAnsi(r io.Reader) io.Reader {
@@ -20,7 +26,7 @@ func StripAnsi(r io.Reader) io.Reader {
 	scanner.Buffer(nil, 1024*1024*1024)
 	for scanner.Scan() {
 		line := scanner.Bytes()
-		line = ansi.ReplaceAll(line, []byte(""))
+		line = AnsiReg.ReplaceAll(line, []byte(""))
 		_, err := buf.Write(line)
 		if err != nil {
 			log.Warnf("failed write ", err)
@@ -35,3 +41,395 @@ func StripAnsi(r io.Reader) io.Reader {
 	}
 	return buf
 }
+
+// StyledRune is a rune and it's associated style. The rune has already been
+// measured using go-runewidth
+type StyledRune struct {
+	Value rune
+	Width int
+	Style tcell.Style
+}
+
+// RuneBuffer is a buffer of runes styled with tcell.Style objects
+type RuneBuffer struct {
+	buf []*StyledRune
+}
+
+// Returns the internal slice of styled runes
+func (rb *RuneBuffer) Runes() []*StyledRune {
+	return rb.buf
+}
+
+// Write writes a rune and it's associated style to the RuneBuffer
+func (rb *RuneBuffer) Write(r rune, style tcell.Style) {
+	w := runewidth.RuneWidth(r)
+	rb.buf = append(rb.buf, &StyledRune{r, w, style})
+}
+
+// Prepend inserts the rune at the beginning of the rune buffer
+func (rb *RuneBuffer) Prepend(r rune, style tcell.Style) {
+	w := runewidth.RuneWidth(r)
+	rb.buf = append([]*StyledRune{{r, w, style}}, rb.buf...)
+}
+
+// String outputs a styled-string using TERM=xterm-256color
+func (rb *RuneBuffer) String() string {
+	return rb.string(0, false, 0)
+}
+
+// string returns a string no longer than n runes. If 'left' is true, the left
+// side of the text is truncated. Pass 0 to return the full string
+func (rb *RuneBuffer) string(n int, left bool, char rune) string {
+	// Use xterm-256color to generate the string. Ultimately all output will
+	// be re-parsed as 'xterm-256color' and tcell will handle the final
+	// output sequences based on the user's TERM
+	ti, err := terminfo.LookupTerminfo("xterm-256color")
+	if err != nil {
+		// Who knows what happened
+		return ""
+	}
+	var (
+		s        = strings.Builder{}
+		style    = tcell.StyleDefault
+		hasStyle = false
+		// w will track the length we have written, or would have
+		// written in the case of left truncate
+		w      = 0
+		offset = 0
+	)
+
+	if left {
+		offset = rb.Len() - n
+	}
+
+	for _, r := range rb.buf {
+		if style != r.Style {
+			hasStyle = true
+			style = r.Style
+			s.WriteString(ti.AttrOff)
+			fg, bg, attrs := style.Decompose()
+
+			switch {
+			case fg.IsRGB() && bg.IsRGB() && ti.SetFgBgRGB != "":
+				fr, fg, fb := fg.RGB()
+				br, bg, bb := bg.RGB()
+				s.WriteString(ti.TParm(
+					ti.SetFgBgRGB,
+					int(fr),
+					int(fg),
+					int(fb),
+					int(br),
+					int(bg),
+					int(bb),
+				))
+			case fg.IsRGB() && ti.SetFgRGB != "":
+				// RGB
+				r, g, b := fg.RGB()
+				s.WriteString(ti.TParm(ti.SetFgRGB, int(r), int(g), int(b)))
+			case bg.IsRGB() && ti.SetBgRGB != "":
+				// RGB
+				r, g, b := bg.RGB()
+				s.WriteString(ti.TParm(ti.SetBgRGB, int(r), int(g), int(b)))
+
+				// Indexed
+			case fg.Valid() && bg.Valid() && ti.SetFgBg != "":
+				s.WriteString(ti.TParm(ti.SetFgBg, int(fg&0xff), int(bg&0xff)))
+			case fg.Valid() && ti.SetFg != "":
+				s.WriteString(ti.TParm(ti.SetFg, int(fg&0xff)))
+			case bg.Valid() && ti.SetBg != "":
+				s.WriteString(ti.TParm(ti.SetBg, int(bg&0xff)))
+			}
+
+			if attrs&tcell.AttrBold != 0 {
+				s.WriteString(ti.Bold)
+			}
+			if attrs&tcell.AttrUnderline != 0 {
+				s.WriteString(ti.Underline)
+			}
+			if attrs&tcell.AttrReverse != 0 {
+				s.WriteString(ti.Reverse)
+			}
+			if attrs&tcell.AttrBlink != 0 {
+				s.WriteString(ti.Blink)
+			}
+			if attrs&tcell.AttrDim != 0 {
+				s.WriteString(ti.Dim)
+			}
+			if attrs&tcell.AttrItalic != 0 {
+				s.WriteString(ti.Italic)
+			}
+			if attrs&tcell.AttrStrikeThrough != 0 {
+				s.WriteString(ti.StrikeThrough)
+			}
+		}
+
+		w += r.Width
+		if left && w <= offset {
+			if w == offset && char != 0 {
+				s.WriteRune(char)
+			}
+			continue
+		}
+		s.WriteRune(r.Value)
+		if n != 0 && !left && w == n {
+			if char != 0 {
+				s.WriteRune(char)
+			}
+			break
+		}
+	}
+	if hasStyle {
+		s.WriteString(ti.AttrOff)
+	}
+	return s.String()
+}
+
+// Len is the length of the string, without ansi sequences
+func (rb *RuneBuffer) Len() int {
+	l := 0
+	for _, r := range rb.buf {
+		l += r.Width
+	}
+	return l
+}
+
+// Truncates to a width of n, optionally append a character to the string.
+// Appending via Truncate allows the character to retain the same style as the
+// string at the truncated location
+func (rb *RuneBuffer) Truncate(n int, char rune) string {
+	return rb.string(n, false, char)
+}
+
+// Truncates a width of n off the beginning of the string, optionally append a
+// character to the string. Appending via Truncate allows the character to
+// retain the same style as the string at the truncated location
+func (rb *RuneBuffer) TruncateHead(n int, char rune) string {
+	return rb.string(n, true, char)
+}
+
+// Applies a style to the buffer. Any currently applied styles will not be
+// overwritten
+func (rb *RuneBuffer) ApplyStyle(style tcell.Style) {
+	for _, sr := range rb.buf {
+		if sr.Style == tcell.StyleDefault {
+			sr.Style = style
+		}
+	}
+}
+
+// ApplyAttrs applies the style, and if another style is present ORs the
+// attributes
+func (rb *RuneBuffer) ApplyAttrs(style tcell.Style) {
+	for _, sr := range rb.buf {
+		if sr.Style == tcell.StyleDefault {
+			sr.Style = style
+			continue
+		}
+		_, _, srAttrs := sr.Style.Decompose()
+		_, _, attrs := style.Decompose()
+		sr.Style = sr.Style.Attributes(srAttrs | attrs)
+	}
+}
+
+// Applies a style to a string. Any currently applied styles will not be overwritten
+func ApplyStyle(style tcell.Style, str string) string {
+	rb := ParseANSI(str)
+	for _, sr := range rb.buf {
+		if sr.Style == tcell.StyleDefault {
+			sr.Style = style
+		}
+	}
+	return rb.String()
+}
+
+// Parses a styled string into a RuneBuffer
+func ParseANSI(s string) *RuneBuffer {
+	p := &parser{
+		buf:      &RuneBuffer{},
+		curStyle: tcell.StyleDefault,
+	}
+	rdr := strings.NewReader(s)
+
+	for {
+		r, _, err := rdr.ReadRune()
+		if err == io.EOF {
+			break
+		}
+		switch r {
+		case 0x1b:
+			p.handleSeq(rdr)
+		default:
+			p.buf.Write(r, p.curStyle)
+		}
+	}
+	return p.buf
+}
+
+// A parser parses a string into a RuneBuffer
+type parser struct {
+	buf      *RuneBuffer
+	curStyle tcell.Style
+}
+
+func (p *parser) handleSeq(rdr io.RuneReader) {
+	r, _, err := rdr.ReadRune()
+	if errors.Is(err, io.EOF) {
+		return
+	}
+	switch r {
+	case '[': // CSI
+		p.handleCSI(rdr)
+	case ']': // OSC
+	case '(': // Designate G0 charset
+		p.swallow(rdr, 1)
+	}
+}
+
+func (p *parser) handleCSI(rdr io.RuneReader) {
+	var (
+		params []int
+		param  []rune
+		hasErr bool
+		er     error
+	)
+outer:
+	for {
+		r, _, err := rdr.ReadRune()
+		if errors.Is(err, io.EOF) {
+			return
+		}
+		switch {
+		case r >= 0x30 && r <= 0x39:
+			param = append(param, r)
+		case r == ':' || r == ';':
+			var ps int
+			if len(param) > 0 {
+				ps, er = strconv.Atoi(string(param))
+				if er != nil {
+					hasErr = true
+					continue
+				}
+			}
+			params = append(params, ps)
+			param = []rune{}
+		case r == 'm':
+			var ps int
+			if len(param) > 0 {
+				ps, er = strconv.Atoi(string(param))
+				if er != nil {
+					hasErr = true
+					continue
+				}
+			}
+			params = append(params, ps)
+			break outer
+		}
+	}
+	if hasErr {
+		// leave the cursor unchanged
+		return
+	}
+	for i := 0; i < len(params); i++ {
+		param := params[i]
+		switch param {
+		case 0:
+			p.curStyle = tcell.StyleDefault
+		case 1:
+			p.curStyle = p.curStyle.Bold(true)
+		case 2:
+			p.curStyle = p.curStyle.Dim(true)
+		case 3:
+			p.curStyle = p.curStyle.Italic(true)
+		case 4:
+			p.curStyle = p.curStyle.Underline(true)
+		case 5:
+			p.curStyle = p.curStyle.Blink(true)
+		case 6:
+			// rapid blink, not supported by tcell. fallback to slow
+			// blink
+			p.curStyle = p.curStyle.Blink(true)
+		case 7:
+			p.curStyle = p.curStyle.Reverse(true)
+		case 8:
+			// Hidden. not supported by tcell
+		case 9:
+			p.curStyle = p.curStyle.StrikeThrough(true)
+		case 21:
+			p.curStyle = p.curStyle.Bold(false)
+		case 22:
+			p.curStyle = p.curStyle.Dim(false)
+		case 23:
+			p.curStyle = p.curStyle.Italic(false)
+		case 24:
+			p.curStyle = p.curStyle.Underline(false)
+		case 25:
+			p.curStyle = p.curStyle.Blink(false)
+		case 26:
+			// rapid blink, not supported by tcell. fallback to slow
+			// blink
+			p.curStyle = p.curStyle.Blink(false)
+		case 27:
+			p.curStyle = p.curStyle.Reverse(false)
+		case 28:
+			// Hidden. unsupported by tcell
+		case 29:
+			p.curStyle = p.curStyle.StrikeThrough(false)
+		case 30, 31, 32, 33, 34, 35, 36, 37:
+			p.curStyle = p.curStyle.Foreground(tcell.PaletteColor(param - 30))
+		case 38:
+			if i+2 < len(params) && params[i+1] == 5 {
+				p.curStyle = p.curStyle.Foreground(tcell.PaletteColor(params[i+2]))
+				i += 2
+			}
+			if i+4 < len(params) && params[i+1] == 2 {
+				switch len(params) {
+				case 6:
+					r := int32(params[i+3])
+					g := int32(params[i+4])
+					b := int32(params[i+5])
+					p.curStyle = p.curStyle.Foreground(tcell.NewRGBColor(r, g, b))
+					i += 5
+				default:
+					r := int32(params[i+2])
+					g := int32(params[i+3])
+					b := int32(params[i+4])
+					p.curStyle = p.curStyle.Foreground(tcell.NewRGBColor(r, g, b))
+					i += 4
+				}
+			}
+		case 40, 41, 42, 43, 44, 45, 46, 47:
+			p.curStyle = p.curStyle.Background(tcell.PaletteColor(param - 40))
+		case 48:
+			if i+2 < len(params) && params[i+1] == 5 {
+				p.curStyle = p.curStyle.Background(tcell.PaletteColor(params[i+2]))
+				i += 2
+			}
+			if i+4 < len(params) && params[i+1] == 2 {
+				switch len(params) {
+				case 6:
+					r := int32(params[i+3])
+					g := int32(params[i+4])
+					b := int32(params[i+5])
+					p.curStyle = p.curStyle.Background(tcell.NewRGBColor(r, g, b))
+					i += 5
+				default:
+					r := int32(params[i+2])
+					g := int32(params[i+3])
+					b := int32(params[i+4])
+					p.curStyle = p.curStyle.Background(tcell.NewRGBColor(r, g, b))
+					i += 4
+				}
+			}
+		case 90, 91, 92, 93, 94, 95, 96, 97:
+			p.curStyle = p.curStyle.Foreground(tcell.PaletteColor(param - 90 + 8))
+		case 100, 101, 102, 103, 104, 105, 106, 107:
+			p.curStyle = p.curStyle.Background(tcell.PaletteColor(param - 100 + 8))
+		}
+	}
+}
+
+func (p *parser) swallow(rdr io.RuneReader, n int) {
+	for i := 0; i < n; i++ {
+		rdr.ReadRune() //nolint:errcheck // we are throwing these reads away
+	}
+}
diff --git a/lib/parse/ansi_test.go b/lib/parse/ansi_test.go
new file mode 100644
index 00000000..f916412b
--- /dev/null
+++ b/lib/parse/ansi_test.go
@@ -0,0 +1,262 @@
+package parse_test
+
+import (
+	"os"
+	"testing"
+
+	"git.sr.ht/~rjarry/aerc/lib/parse"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestParser(t *testing.T) {
+	//nolint:errcheck // we'll fail the test if this fails
+	_ = os.Setenv("COLORTERM", "truecolor")
+	tests := []struct {
+		name           string
+		input          string
+		expectedString string
+		expectedLen    int
+	}{
+		{
+			name:           "no style",
+			input:          "hello, world",
+			expectedString: "hello, world",
+			expectedLen:    12,
+		},
+		{
+			name:           "bold",
+			input:          "\x1b[1mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[1mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "dim",
+			input:          "\x1b[2mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[2mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "bold and dim",
+			input:          "\x1b[1;2mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[1m\x1b[2mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "italic",
+			input:          "\x1b[3mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[3mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "underline",
+			input:          "\x1b[4mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[4mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "blink",
+			input:          "\x1b[5mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[5mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "fast blink",
+			input:          "\x1b[6mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[5mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "reverse",
+			input:          "\x1b[7mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[7mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "hidden",
+			input:          "\x1b[8mhello, world",
+			expectedString: "hello, world",
+			expectedLen:    12,
+		},
+		{
+			name:           "strikethrough",
+			input:          "\x1b[9mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[9mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "bold hello, normal world",
+			input:          "\x1b[1mhello, \x1b[21mworld",
+			expectedString: "\x1b(B\x1b[m\x1b[1mhello, \x1b(B\x1b[mworld\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "bold hello, normal world v2",
+			input:          "\x1b[1mhello, \x1b[mworld",
+			expectedString: "\x1b(B\x1b[m\x1b[1mhello, \x1b(B\x1b[mworld\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "8 bit color: foreground",
+			input:          "\x1b[30mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[30mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "8 bit color: background",
+			input:          "\x1b[41mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[41mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "8 bit color: foreground and background",
+			input:          "\x1b[31;41mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[31;41mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "16 bit color: foreground",
+			input:          "\x1b[90mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[90mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "16 bit color: background",
+			input:          "\x1b[101mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[101mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "16 bit color: foreground and background",
+			input:          "\x1b[91;101mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[91;101mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "256 color: foreground",
+			input:          "\x1b[38;5;2mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[32mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "256 color: foreground",
+			input:          "\x1b[38;5;132mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[38;5;132mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "256 color: background",
+			input:          "\x1b[48;5;132mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[48;5;132mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "256 color: foreground and background",
+			input:          "\x1b[38;5;20;48;5;20mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[38;5;20;48;5;20mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "256 color: background",
+			input:          "\x1b[48;5;2mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[42mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "true color: foreground",
+			input:          "\x1b[38;2;0;0;0mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[38;2;0;0;0mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "true color: foreground with color space",
+			input:          "\x1b[38;2;;0;0;0mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[38;2;0;0;0mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "true color: foreground with color space and colons",
+			input:          "\x1b[38:2::0:0:0mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[38;2;0;0;0mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "true color: background",
+			input:          "\x1b[48;2;0;0;0mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[48;2;0;0;0mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "true color: background with color space",
+			input:          "\x1b[48;2;;0;0;0mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[48;2;0;0;0mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+		{
+			name:           "true color: foreground and background",
+			input:          "\x1b[38;2;200;200;200;48;2;0;0;0mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[38;2;200;200;200;48;2;0;0;0mhello, world\x1b(B\x1b[m",
+			expectedLen:    12,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			buf := parse.ParseANSI(test.input)
+			assert.Equal(t, test.expectedString, buf.String())
+			assert.Equal(t, test.expectedLen, buf.Len())
+		})
+	}
+}
+
+func TestTruncate(t *testing.T) {
+	tests := []struct {
+		name           string
+		input          string
+		expectedString string
+	}{
+		{
+			name:           "no style, truncate at 5",
+			input:          "hello, world",
+			expectedString: "hello",
+		},
+		{
+			name:           "bold, truncate at 5",
+			input:          "\x1b[1mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[1mhello\x1b(B\x1b[m",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			buf := parse.ParseANSI(test.input)
+			assert.Equal(t, test.expectedString, buf.Truncate(5, 0))
+		})
+	}
+}
+
+func TestTruncateHead(t *testing.T) {
+	tests := []struct {
+		name           string
+		input          string
+		expectedString string
+		expectedLen    int
+	}{
+		{
+			name:           "no style, truncate head at 5",
+			input:          "hello, world",
+			expectedString: "world",
+		},
+		{
+			name:           "bold, truncate head at 5",
+			input:          "\x1b[1mhello, world",
+			expectedString: "\x1b(B\x1b[m\x1b[1mworld\x1b(B\x1b[m",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			buf := parse.ParseANSI(test.input)
+			assert.Equal(t, test.expectedString, buf.TruncateHead(5, 0))
+		})
+	}
+}