kitty-mirror/tools/highlight/impl.go

package highlight

import (
	"fmt"
	"io"
	"path/filepath"
	"strings"
	"sync"

	"github.com/alecthomas/chroma/v2"
	"github.com/alecthomas/chroma/v2/lexers"
	"github.com/alecthomas/chroma/v2/styles"
	"github.com/kovidgoyal/kitty/tools/utils"
)

var _ = fmt.Print

var default_style = sync.OnceValue(func() *chroma.Style {
	// Default style generated by python style.py default pygments.styles.default.DefaultStyle
	// with https://raw.githubusercontent.com/alecthomas/chroma/master/_tools/style.py
	return styles.Register(chroma.MustNewStyle("default", chroma.StyleEntries{
		chroma.TextWhitespace:        "#bbbbbb",
		chroma.Comment:               "italic #3D7B7B",
		chroma.CommentPreproc:        "noitalic #9C6500",
		chroma.Keyword:               "bold #008000",
		chroma.KeywordPseudo:         "nobold",
		chroma.KeywordType:           "nobold #B00040",
		chroma.Operator:              "#666666",
		chroma.OperatorWord:          "bold #AA22FF",
		chroma.NameBuiltin:           "#008000",
		chroma.NameFunction:          "#0000FF",
		chroma.NameClass:             "bold #0000FF",
		chroma.NameNamespace:         "bold #0000FF",
		chroma.NameException:         "bold #CB3F38",
		chroma.NameVariable:          "#19177C",
		chroma.NameConstant:          "#880000",
		chroma.NameLabel:             "#767600",
		chroma.NameEntity:            "bold #717171",
		chroma.NameAttribute:         "#687822",
		chroma.NameTag:               "bold #008000",
		chroma.NameDecorator:         "#AA22FF",
		chroma.LiteralString:         "#BA2121",
		chroma.LiteralStringDoc:      "italic",
		chroma.LiteralStringInterpol: "bold #A45A77",
		chroma.LiteralStringEscape:   "bold #AA5D1F",
		chroma.LiteralStringRegex:    "#A45A77",
		chroma.LiteralStringSymbol:   "#19177C",
		chroma.LiteralStringOther:    "#008000",
		chroma.LiteralNumber:         "#666666",
		chroma.GenericHeading:        "bold #000080",
		chroma.GenericSubheading:     "bold #800080",
		chroma.GenericDeleted:        "#A00000",
		chroma.GenericInserted:       "#008400",
		chroma.GenericError:          "#E40000",
		chroma.GenericEmph:           "italic",
		chroma.GenericStrong:         "bold",
		chroma.GenericPrompt:         "bold #000080",
		chroma.GenericOutput:         "#717171",
		chroma.GenericTraceback:      "#04D",
		chroma.Error:                 "border:#FF0000",
		chroma.Background:            " bg:#f8f8f8",
	}))
})

// Clear the background colour.
func clear_background(style *chroma.Style) *chroma.Style {
	builder := style.Builder()
	bg := builder.Get(chroma.Background)
	bg.Background = 0
	bg.NoInherit = true
	builder.AddEntry(chroma.Background, bg)
	style, _ = builder.Build()
	return style
}

func ansi_formatter(w io.Writer, style *chroma.Style, sanitize func(string) string, it chroma.Iterator) (err error) {
	const SGR_PREFIX = "\033["
	const SGR_SUFFIX = "m"
	style = clear_background(style)
	before, after := make([]byte, 0, 64), make([]byte, 0, 64)
	nl := []byte{'\n'}
	write_sgr := func(which []byte) (err error) {
		if len(which) > 1 {
			if _, err = w.Write(utils.UnsafeStringToBytes(SGR_PREFIX)); err != nil {
				return err
			}
			if _, err = w.Write(which[:len(which)-1]); err != nil {
				return err
			}
			if _, err = w.Write(utils.UnsafeStringToBytes(SGR_SUFFIX)); err != nil {
				return err
			}
		}
		return
	}
	write := func(text string) (err error) {
		if err = write_sgr(before); err != nil {
			return err
		}
		if _, err = w.Write(utils.UnsafeStringToBytes(text)); err != nil {
			return err
		}
		if err = write_sgr(after); err != nil {
			return err
		}
		return
	}

	for token := it(); token != chroma.EOF; token = it() {
		entry := style.Get(token.Type)
		before, after = before[:0], after[:0]
		if !entry.IsZero() {
			if entry.Bold == chroma.Yes {
				before = append(before, '1', ';')
				after = append(after, '2', '2', '1', ';')
			}
			if entry.Underline == chroma.Yes {
				before = append(before, '4', ';')
				after = append(after, '2', '4', ';')
			}
			if entry.Italic == chroma.Yes {
				before = append(before, '3', ';')
				after = append(after, '2', '3', ';')
			}
			if entry.Colour.IsSet() {
				before = append(before, fmt.Sprintf("38:2:%d:%d:%d;", entry.Colour.Red(), entry.Colour.Green(), entry.Colour.Blue())...)
				after = append(after, '3', '9', ';')
			}
		}
		// independently format each line in a multiline token, needed for the diff kitten highlighting to work, also
		// pagers like less reset SGR formatting at line boundaries
		text := sanitize(token.Value)
		for text != "" {
			idx := strings.IndexByte(text, '\n')
			if idx < 0 {
				if err = write(text); err != nil {
					return err
				}
				break
			}
			if err = write(text[:idx]); err != nil {
				return err
			}
			if _, err = w.Write(nl); err != nil {
				return err
			}
			text = text[idx+1:]
		}
	}
	return nil
}

func resolved_chroma_style(srd StyleResolveData) *chroma.Style {
	name := srd.StyleName()
	var style *chroma.Style
	if name == "default" {
		style = default_style()
	} else {
		style = styles.Get(name)
	}
	if style == nil {
		if srd.UseLightColors() {
			style = default_style()
		} else {
			style = styles.Get("monokai")
			if style == nil {
				style = styles.Get("github-dark")
			}
		}
		if style == nil {
			style = styles.Fallback
		}
	}
	return style
}

type highlighter struct {
	tokens_map map[string][]chroma.Token
	lock       sync.Mutex
	sanitize   func(string) string
}

func (h *highlighter) Sanitize(x string) string { return h.sanitize(x) }

func (h *highlighter) HighlightFile(path string, srd StyleResolveData) (highlighted_string string, err error) {
	defer func() {
		if r := recover(); r != nil {
			text, _ := utils.Format_stacktrace_on_panic(r)
			err = fmt.Errorf("%s", text)
		}
	}()
	filename_for_detection := filepath.Base(path)
	ext := filepath.Ext(filename_for_detection)
	if ext != "" {
		ext = strings.ToLower(ext[1:])
		r := srd.SyntaxAliases()[ext]
		if r != "" {
			filename_for_detection = "file." + r
		}
	}
	text, err := srd.TextForPath(path)
	if err != nil {
		return "", err
	}
	h.lock.Lock()
	tokens := h.tokens_map[path]
	h.lock.Unlock()
	if tokens == nil {
		lexer := lexers.Match(filename_for_detection)
		if lexer == nil {
			lexer = lexers.Analyse(text)
		}
		if lexer == nil {
			return "", fmt.Errorf("Cannot highlight %#v: %w", path, ErrNoLexer)
		}
		lexer = chroma.Coalesce(lexer)
		iterator, err := lexer.Tokenise(nil, text)
		if err != nil {
			return "", err
		}
		tokens = iterator.Tokens()
		h.lock.Lock()
		h.tokens_map[path] = tokens
		h.lock.Unlock()
	}
	w := strings.Builder{}
	w.Grow(len(text) * 2)
	err = ansi_formatter(&w, resolved_chroma_style(srd), h.sanitize, chroma.Literator(tokens...))
	return w.String(), err
}