current commit
261030.1 · A scope box starts at the least indent among its rows
package main

import (
	"slices"
	"strconv"
	"strings"

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

// A scope is a block of code that spans lines, drawn as a sunken box
// behind its text. Monospace text is a grid, so a box is a rectangle of
// cells: from its first row to its last, from the least indent among its
// rows to past its widest line. A leaf is a bracket pair within one line, eg
// (x, y), whose box spans its brackets' cells. Boxes nest without limit.
type scope struct {
	first, last int // rows, counted from 0
	left, right int // cells; right is one past the widest line or the closing bracket
	depth       int // 1 for a scope in no other
	leaf        int // for a leaf, 1 plus the leaves around it on its line; 0 for a block
}

// tabStop is CSS's default tab-size, which code pages keep.
const tabStop = 8

// lineCells measures each line in cells: where its text starts and where
// it ends. A blank line's indent is -1.
func lineCells(lines [][]chroma.Token) (indent, width []int) {
	indent = make([]int, len(lines))
	width = make([]int, len(lines))
	for i, line := range lines {
		col, ind := 0, -1
		for _, tok := range line {
			for j := 0; j < len(tok.Value); j++ {
				switch c := tok.Value[j]; {
				case c == '\n':
				case c == '\t':
					col = (col/tabStop + 1) * tabStop
				case c == ' ':
					col++
				case c&0xC0 == 0x80: // a UTF-8 continuation byte
				default:
					if ind < 0 {
						ind = col
					}
					col++
				}
			}
		}
		indent[i] = ind
		width[i] = col
	}
	return indent, width
}

// findScopes finds the blocks of code: bracket pairs whose ends lie on
// different lines, strings that span lines, and in Python the lines
// indented under a line that ends with a colon. Bracket pairs and strings
// within a line, with something inside them, are leaves, the units vim's
// ci( and ci" act on. Brackets count only in punctuation and operator
// tokens, so a brace inside a string or comment opens nothing; a string
// is a run of string tokens, its quotes and escapes included. Blocks come
// first, outer before inner, then leaves by row, outer before inner.
func findScopes(lines [][]chroma.Token, python bool) []scope {
	indent, width := lineCells(lines)
	type open struct {
		c        byte
		row, col int
	}
	var stack []open
	var found, leaves []scope
	// level counts the brackets open on row, plus one: a leaf's level
	// among the leaves around it.
	level := func(open []open, row int) int {
		n := 1
		for _, e := range open {
			if e.row == row {
				n++
			}
		}
		return n
	}
	// str is the string being read: where it starts and where it has got
	// to, row and cell.
	var str struct {
		on                 bool
		row, col, end, top int
	}
	endString := func() {
		if !str.on {
			return
		}
		str.on = false
		switch {
		case str.end > str.row:
			found = append(found, scope{first: str.row, last: str.end})
		case str.top > str.col+2:
			leaves = append(leaves, scope{first: str.row, last: str.row, left: str.col, right: str.top, leaf: level(stack, str.row)})
		}
	}
	for row, line := range lines {
		col := 0
		for _, tok := range line {
			cat := tok.Type.Category()
			bracketed := cat == chroma.Punctuation || cat == chroma.Operator
			if tok.Type.SubCategory() == chroma.LiteralString {
				if !str.on {
					str.on, str.row, str.col = true, row, col
				}
			} else {
				endString()
			}
			for j := 0; j < len(tok.Value); j++ {
				c := tok.Value[j]
				switch {
				case c == '\n', c&0xC0 == 0x80: // a UTF-8 continuation byte
					continue
				case c == '\t':
					col = (col/tabStop + 1) * tabStop
					continue
				case !bracketed:
				case c == '(' || c == '[' || c == '{':
					stack = append(stack, open{c, row, col})
				case c == ')' || c == ']' || c == '}':
					// A stray closer ends the brackets it matches, if any.
					for k := len(stack) - 1; k >= 0; k-- {
						o := stack[k]
						if o.c != opener(c) {
							continue
						}
						if o.row < row {
							found = append(found, scope{first: o.row, last: row})
						} else if col > o.col+1 {
							leaves = append(leaves, scope{first: row, last: row, left: o.col, right: col + 1, leaf: level(stack[:k], row)})
						}
						stack = stack[:k]
						break
					}
				}
				col++
			}
			if str.on {
				str.end, str.top = row, col
			}
		}
		if python && len(stack) == 0 && endsWithColon(line) {
			last := row
			for r := row + 1; r < len(lines); r++ {
				if indent[r] < 0 {
					continue
				}
				if indent[r] <= indent[row] {
					break
				}
				last = r
			}
			if last > row {
				found = append(found, scope{first: row, last: last})
			}
		}
	}
	endString()

	// Outer scopes first; two brackets that open and close on the same
	// rows, as in foo(func() {, draw one box.
	slices.SortFunc(found, func(a, b scope) int {
		if a.first != b.first {
			return a.first - b.first
		}
		return b.last - a.last
	})
	found = slices.CompactFunc(found, func(a, b scope) bool { return a.first == b.first && a.last == b.last })

	parent := make([]int, len(found))
	var in []int
	for i := range found {
		s := &found[i]
		for len(in) > 0 && found[in[len(in)-1]].last < s.last {
			in = in[:len(in)-1]
		}
		parent[i] = -1
		if len(in) > 0 {
			parent[i] = in[len(in)-1]
		}
		s.depth = len(in) + 1
		// A box starts at the least indent among its rows, so no text
		// lies left of it, eg a raw string's lines at the margin. A box
		// thus never starts left of the box around it.
		s.left = indent[s.first]
		for r := s.first; r <= s.last; r++ {
			if indent[r] >= 0 {
				s.left = min(s.left, indent[r])
			}
			s.right = max(s.right, width[r])
		}
		in = append(in, i)
	}
	// A box reaches one cell past each box inside it.
	for i := len(found) - 1; i >= 0; i-- {
		if p := parent[i]; p >= 0 {
			found[p].right = max(found[p].right, found[i].right+1)
		}
	}

	// A leaf lies one level below the blocks around its row and the
	// leaves around it.
	blocks := make([]int, len(lines)+1)
	for _, b := range found {
		blocks[b.first]++
		blocks[b.last+1]--
	}
	for r := 1; r < len(blocks); r++ {
		blocks[r] += blocks[r-1]
	}
	for i := range leaves {
		leaves[i].depth = blocks[leaves[i].first] + leaves[i].leaf
	}
	slices.SortFunc(leaves, func(a, b scope) int {
		if a.first != b.first {
			return a.first - b.first
		}
		if a.left != b.left {
			return a.left - b.left
		}
		return b.right - a.right
	})

	// Merge by row, a row's blocks before its leaves.
	all := make([]scope, 0, len(found)+len(leaves))
	for len(found) > 0 || len(leaves) > 0 {
		if len(leaves) == 0 || len(found) > 0 && found[0].first <= leaves[0].first {
			all, found = append(all, found[0]), found[1:]
		} else {
			all, leaves = append(all, leaves[0]), leaves[1:]
		}
	}
	return all
}

// opener is the bracket that a closing bracket closes.
func opener(c byte) byte {
	switch c {
	case ')':
		return '('
	case ']':
		return '['
	}
	return '{'
}

// endsWithColon reports whether a line's last token, comments and space
// aside, is a colon.
func endsWithColon(line []chroma.Token) bool {
	for i := len(line) - 1; i >= 0; i-- {
		t := line[i]
		v := strings.TrimRight(t.Value, " \t\r\n")
		if v == "" || t.Type.Category() == chroma.Comment {
			continue
		}
		return t.Type.Category() == chroma.Punctuation && strings.HasSuffix(v, ":")
	}
	return false
}

// writeScopes writes the boxes that start on one row. Each is empty and
// positioned on its own, so it adds no text and moves none; the
// stylesheet turns the cells and rows into lengths. Boxes at odd depths
// are shaded, so shade and page alternate inward.
func writeScopes(b *strings.Builder, scopes []scope) {
	var num [20]byte
	cell := func(name string, n int) {
		b.WriteString(name)
		b.Write(strconv.AppendInt(num[:0], int64(n), 10))
	}
	for _, s := range scopes {
		b.WriteString(`<i class="scope`)
		if s.leaf > 0 {
			b.WriteString(" leaf")
		}
		if s.depth%2 == 1 {
			b.WriteString(" shade")
		}
		cell(`" style="--c:`, s.left)
		cell(`;--w:`, s.right-s.left)
		if s.leaf == 0 {
			cell(`;--r:`, s.last-s.first+1)
		}
		cell(`;--d:`, s.depth)
		b.WriteString(`"></i>`)
	}
}