codeaware.go

package tok

import (
	"regexp"
	"strings"
)

// Span identifies a protected region of source code that must survive
// compression. Spans are line-oriented (1-based, inclusive) and carry the
// extracted symbol name when one is available.
//
// Spans are the unit of "code intelligence" exchanged between a SymbolProvider
// and the compression pipeline. The default provider derives them from the
// chunker's symbol regexes, but an embedder can supply spans from any source
// (e.g. an LSP server) via a custom SymbolProvider.
type Span struct {
	StartLine int    `json:"start_line"` // 1-based, inclusive
	EndLine   int    `json:"end_line"`   // 1-based, inclusive
	Symbol    string `json:"symbol,omitempty"`
	Lines     []string
}

// SymbolProvider identifies the semantically significant spans of a piece of
// source code. It is the seam that lets hawk later inject a real LSP-backed
// provider in place of tok's dependency-light regex default.
//
// Implementations must be safe for concurrent use; they receive the full code
// text and a language name (e.g. "go", "python"), and return the spans that
// compression must preserve. An empty lang signals best-effort detection.
type SymbolProvider interface {
	Symbols(code, lang string) []Span
}

// RegexSymbolProvider is the default, dependency-light SymbolProvider. It reuses
// the chunker's language boundary and symbol patterns to find declaration and
// signature lines without any LSP integration.
type RegexSymbolProvider struct{}

// DefaultSymbolProvider is the SymbolProvider used by WithCodeAware when no
// other provider is supplied via WithSymbolProvider.
var DefaultSymbolProvider SymbolProvider = RegexSymbolProvider{}

// Symbols implements SymbolProvider using the chunker's boundary/symbol regexes.
//
// A line is treated as a protected signature when it matches the language
// boundary pattern (function/type/class/export/etc.). For each such line the
// returned Span covers the signature itself plus any continuation lines that
// complete a multi-line signature (an unbalanced "(" run, typical of long
// parameter lists), so the full signature is preserved as a unit.
func (RegexSymbolProvider) Symbols(code, lang string) []Span {
	if code == "" {
		return nil
	}

	l := strings.ToLower(strings.TrimSpace(lang))
	if l == "" {
		l = detectLanguage(code)
	}

	boundary := boundaryForLang(l)
	if boundary == nil {
		return nil
	}
	symbolRe := symbolReForLang(l)

	lines := strings.Split(code, "\n")
	var spans []Span
	for i := 0; i < len(lines); i++ {
		trimmed := strings.TrimSpace(lines[i])
		if trimmed == "" || !boundary.MatchString(trimmed) {
			continue
		}

		start := i
		// Extend across an unbalanced-paren multi-line signature.
		end := i
		if depth := parenDepth(lines[i]); depth > 0 {
			for end+1 < len(lines) && depth > 0 {
				end++
				depth += parenDepth(lines[end])
			}
		}

		symbol := ""
		if symbolRe != nil {
			if m := symbolRe.FindStringSubmatch(trimmed); m != nil {
				for _, g := range m[1:] {
					if g != "" {
						symbol = g
						break
					}
				}
			}
		}

		spans = append(spans, Span{
			StartLine: start + 1,
			EndLine:   end + 1,
			Symbol:    symbol,
			Lines:     append([]string(nil), lines[start:end+1]...),
		})
		i = end
	}

	return spans
}

// parenDepth returns the net change in parenthesis nesting for a line: open
// parens minus close parens. Used to detect multi-line signatures.
func parenDepth(line string) int {
	return strings.Count(line, "(") - strings.Count(line, ")")
}

// codeProtector captures the protected signature lines for an input so they can
// be re-injected after compression.
type codeProtector struct {
	spans []Span
}

// newCodeProtector builds a protector for code using the given provider (or the
// default when nil). It returns nil when there is nothing to protect.
func newCodeProtector(code, lang string, provider SymbolProvider) *codeProtector {
	if code == "" {
		return nil
	}
	if provider == nil {
		provider = DefaultSymbolProvider
	}
	spans := provider.Symbols(code, lang)
	if len(spans) == 0 {
		return nil
	}
	return &codeProtector{spans: spans}
}

// protectedLines flattens all span lines into the set of trimmed, non-empty
// signature lines that must appear in the output.
func (p *codeProtector) protectedLines() []string {
	if p == nil {
		return nil
	}
	var out []string
	for _, s := range p.spans {
		for _, ln := range s.Lines {
			t := strings.TrimSpace(ln)
			if t != "" {
				out = append(out, ln)
			}
		}
	}
	return out
}

// guard re-injects any protected line that the compression pipeline dropped.
// A protected line is considered present if its trimmed form is a substring of
// the (whitespace-normalized) compressed output, which tolerates downstream
// reflowing/indentation changes. Missing lines are appended verbatim under a
// marker so the signatures are never lost.
func (p *codeProtector) guard(original, compressed string) string {
	if p == nil {
		return compressed
	}

	haystack := normalizeWhitespace(compressed)
	var missing []string
	seen := make(map[string]struct{})
	for _, ln := range p.protectedLines() {
		needle := normalizeWhitespace(ln)
		if needle == "" {
			continue
		}
		if _, dup := seen[needle]; dup {
			continue
		}
		if strings.Contains(haystack, needle) {
			continue
		}
		seen[needle] = struct{}{}
		missing = append(missing, strings.TrimRight(ln, " \t"))
	}

	if len(missing) == 0 {
		return compressed
	}

	var b strings.Builder
	b.WriteString(compressed)
	if compressed != "" && !strings.HasSuffix(compressed, "\n") {
		b.WriteByte('\n')
	}
	b.WriteString(codeAwareMarker)
	b.WriteByte('\n')
	b.WriteString(strings.Join(missing, "\n"))
	return b.String()
}

// codeAwareMarker labels re-injected signature lines in the output.
const codeAwareMarker = "// [tok:code-aware preserved signatures]"

var wsRe = regexp.MustCompile(`\s+`)

// normalizeWhitespace collapses runs of whitespace to a single space and trims,
// so signature containment checks are robust to indentation/reflow changes.
func normalizeWhitespace(s string) string {
	return strings.TrimSpace(wsRe.ReplaceAllString(s, " "))
}