fix: locate sequence matches by position, not by searching

Signed-off-by: Joseph Kato <joseph@jdkato.io>

Author: jdkato

go.mod

@@ -14,7 +14,7 @@ require (
 	github.com/expr-lang/expr v1.17.7
 	github.com/gobwas/glob v0.2.3
 	github.com/jdkato/go-tree-sitter-julia v0.1.0
-	github.com/jdkato/prose/v3 v3.0.0
+	github.com/jdkato/prose/v3 v3.0.1
 	github.com/litao91/goldmark-mathjax v0.0.0-20210217064022-a43cf739a50f
 	github.com/mitchellh/mapstructure v1.5.0
 	github.com/niklasfasching/go-org v1.7.0

go.sum

@@ -63,8 +63,8 @@ github.com/huandu/xstrings v1.5.0 h1:2ag3IFq9ZDANvthTwTiqSSZLjDc+BedvHPAp5tJy2TI
 github.com/huandu/xstrings v1.5.0/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq4ovT0aE=
 github.com/jdkato/go-tree-sitter-julia v0.1.0 h1:z+6zTbd6PHMKAge7GJx9QIwPQX2NOKb4Pj5jteJvaYY=
 github.com/jdkato/go-tree-sitter-julia v0.1.0/go.mod h1:lXNEZorcvU63DcANEklLMbDRjwam4VQ44MIV1Cck0w8=
-github.com/jdkato/prose/v3 v3.0.0 h1:MomxPb0GtTWXboechFvI429uBZoalYMgQ1hpTfI6Tbg=
-github.com/jdkato/prose/v3 v3.0.0/go.mod h1:VRQyR/EyA6bX457JTCIF67n04ycHek4v6v0veGiJvUk=
+github.com/jdkato/prose/v3 v3.0.1 h1:tIRDKg5i8HLBtrj3SdVHD948cDlz7b4mDupZI8206Uw=
+github.com/jdkato/prose/v3 v3.0.1/go.mod h1:VRQyR/EyA6bX457JTCIF67n04ycHek4v6v0veGiJvUk=
 github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
 github.com/klauspost/cpuid/v2 v2.0.10/go.mod h1:g2LTdtYhdyuGPqyWyv7qRAmj1WBqxuObKfj5c0PQa7c=
 github.com/klauspost/cpuid/v2 v2.0.12 h1:p9dKCg8i4gmOxtv35DvrYoWqYzQrvEVdjQ762Y0OqZE=

internal/check/sequence.go

@@ -14,9 +14,18 @@ import (
 
 // NLPToken represents a token of text with NLP-related attributes.
 type NLPToken struct {
-	Pattern  string
-	Tag      string
-	Skip     int
+	Pattern string
+	Tag     string
+	Skip    int
+
+	// Target narrows the alert to this token alone.
+	//
+	// Without it a match spans every token in the sequence. Marking one lets a
+	// rule require surrounding context while pointing at only the part the
+	// writer should change -- "flag the space, but only between these two
+	// words".
+	Target bool
+
 	re       *rx.Regexp
 	Negate   bool
 	optional bool
@@ -133,14 +142,36 @@ func tokensMatch(token NLPToken, word tag.Token) bool {
 	return true
 }
 
-func sequenceMatches(idx int, chk Sequence, target NLPToken, words []tag.Token, history []int) ([]string, int) {
+// match describes one sequence hit: the matched words' text, the index of the
+// anchor word, and the range of word indices the match covers.
+//
+// The index range is what lets Run report where the match actually is. The
+// text alone is not enough: the same sequence can occur more than once, and
+// re-joining the words does not reproduce the source when the spacing is
+// irregular.
+type match struct {
+	text  []string
+	index int
+	lo    int
+	hi    int
+
+	// wordAt maps a position in the expanded token slice to the word it
+	// matched, so a targeted token can be resolved back to its span.
+	wordAt map[int]int
+}
+
+func (m match) ok() bool { return len(m.text) > 0 && m.lo >= 0 && m.hi >= m.lo }
+
+func sequenceMatches(idx int, chk Sequence, target NLPToken, words []tag.Token, history []int) match {
 	var text []string
 
 	toks := chk.Tokens
 
 	sizeT := len(toks)
 	sizeW := len(words)
 	index := 0
+	lo, hi := -1, -1
+	wordAt := map[int]int{}
 
 	for jdx, tok := range words {
 		if tokensMatch(target, tok) && !core.IntInSlice(jdx, history) {
@@ -157,12 +188,14 @@ func sequenceMatches(idx int, chk Sequence, target NLPToken, words []tag.Token,
 				// side to check -- hence, `idx > 0`.
 				for i := 1; idx-i >= 0; i++ {
 					if jdx-i < 0 {
-						return []string{}, index
+						return match{index: index, lo: -1, hi: -1}
 					}
 					tok := toks[idx-i]
 
 					word := words[jdx-i]
 					text = append([]string{word.Text}, text...)
+					lo = jdx - i
+					wordAt[idx-i] = jdx - i
 
 					// NOTE: We have to perform this conversion because the token slice is made
 					// with the right-hand orientation in mind. For example,
@@ -176,7 +209,7 @@ func sequenceMatches(idx int, chk Sequence, target NLPToken, words []tag.Token,
 
 					mat := tokensMatch(tok, word)
 					if !mat && !tok.optional {
-						return []string{}, index
+						return match{index: index, lo: -1, hi: -1}
 					} else if mat && tok.optional {
 						break
 					}
@@ -189,16 +222,21 @@ func sequenceMatches(idx int, chk Sequence, target NLPToken, words []tag.Token,
 				// side to check.
 				for i := 0; idx+i < sizeT; i++ {
 					if jdx+i >= sizeW {
-						return []string{}, index
+						return match{index: index, lo: -1, hi: -1}
 					}
 					tok := toks[idx+i]
 
 					word := words[jdx+i]
 					text = append(text, word.Text)
+					if lo < 0 || jdx+i < lo {
+						lo = jdx + i
+					}
+					hi = jdx + i
+					wordAt[idx+i] = jdx + i
 
 					mat := tokensMatch(tok, word)
 					if !mat && !tok.optional {
-						return []string{}, index
+						return match{index: index, lo: -1, hi: -1}
 					} else if mat && tok.optional {
 						break
 					}
@@ -208,7 +246,7 @@ func sequenceMatches(idx int, chk Sequence, target NLPToken, words []tag.Token,
 		}
 	}
 
-	return text, index
+	return match{text: text, index: index, lo: lo, hi: hi, wordAt: wordAt}
 }
 
 func stepsToString(steps []string) string {
@@ -243,6 +281,67 @@ func stepsToString(steps []string) string {
 	return strings.TrimSpace(sb.String())
 }
 
+// locate returns the span of a match within txt, plus the matched text.
+//
+// When the tokens carry offsets, the span is taken straight from them, so it
+// is exact even when the sequence occurs more than once or the source spacing
+// is irregular. Without offsets we fall back to rebuilding the text and
+// searching for it, which is what Vale did before tokens had positions; that
+// path returns nil rather than a negative span when the search fails.
+func (s Sequence) locate(txt string, words []tag.Token, m match, positioned bool) ([]int, string) {
+	if positioned && m.hi < len(words) {
+		lo, hi := m.lo, m.hi
+		if tlo, thi, ok := s.targetRange(m); ok {
+			lo, hi = tlo, thi
+		}
+
+		start := words[lo].Start
+		end := words[hi].Start + len(words[hi].Text)
+		if start >= 0 && end <= len(txt) && start < end {
+			return []int{start, end}, txt[start:end]
+		}
+	}
+
+	seq := stepsToString(m.text)
+	if ssp := strings.Index(txt, seq); ssp >= 0 {
+		return []int{ssp, ssp + len(seq)}, seq
+	}
+	return nil, seq
+}
+
+// targetRange returns the span of words covered by the rule's `target`
+// tokens.
+//
+// Several tokens may be marked, in which case the range runs from the first to
+// the last -- a target of two words reports both, not just one. Unmarked
+// tokens between them are included, since the result has to be a single
+// contiguous span.
+func (s Sequence) targetRange(m match) (int, int, bool) {
+	lo, hi := -1, -1
+	for i := range s.Tokens {
+		if !s.Tokens[i].Target {
+			continue
+		}
+		w, ok := m.wordAt[i]
+		if !ok {
+			// A targeted token that matched nothing -- an optional or skipped
+			// one. Narrowing to a range we cannot fully resolve would report
+			// the wrong span, so fall back to the whole match.
+			return 0, 0, false
+		}
+		if lo < 0 || w < lo {
+			lo = w
+		}
+		if w > hi {
+			hi = w
+		}
+	}
+	if lo < 0 {
+		return 0, 0, false
+	}
+	return lo, hi, true
+}
+
 // Run looks for the user-defined sequence of tokens.
 func (s Sequence) Run(blk nlp.Block, f *core.File, _ *core.Config) ([]core.Alert, error) {
 	var alerts []core.Alert
@@ -252,26 +351,43 @@ func (s Sequence) Run(blk nlp.Block, f *core.File, _ *core.Config) ([]core.Alert
 	// This is *always* sentence-scoped.
 	words := nlp.TextToTokens(blk.Text, &f.NLP)
 
+	// A remote NLP endpoint returns text and tags only, so we have no offsets
+	// to work from and have to fall back to locating the match by its text.
+	positioned := f.NLP.Endpoint == ""
+
 	txt := blk.Text
 	for idx, tok := range s.Tokens {
 		if !tok.Negate && tok.Pattern != "" {
 			// We're looking for our "anchor" ...
 			for _, loc := range tok.re.FindAllStringIndex(txt, -1) {
 				// These are all possible violations in `txt`:
-				steps, index := sequenceMatches(idx, s, tok, words, history)
-				history = append(history, index)
+				m := sequenceMatches(idx, s, tok, words, history)
+				history = append(history, m.index)
+
+				if m.ok() {
+					span, seq := s.locate(txt, words, m, positioned)
+					if span == nil {
+						// We matched but cannot say where; reporting a bogus
+						// span is worse than reporting nothing.
+						continue
+					}
 
-				if len(steps) > 0 {
-					seq := stepsToString(steps)
-					ssp := strings.Index(txt, seq)
+					// When the block knows where it sits in the document, hand
+					// back an absolute offset. Otherwise the span is
+					// block-relative and has to be located by searching, which
+					// resolves every repeat of a sentence to the first one.
+					absolute := blk.Offset >= 0
+					if absolute {
+						span = []int{blk.Offset + span[0], blk.Offset + span[1]}
+					}
 
 					a := core.Alert{
 						Check: s.Name, Severity: s.Level, Link: s.Link,
-						Span: []int{ssp, ssp + len(seq)}, Hide: false,
+						Span: span, Hide: false, HasByteOffsets: absolute,
 						Match: seq, Action: s.Action}
 
 					a.Message, a.Description = formatMessages(s.Message,
-						s.Description, steps...)
+						s.Description, m.text...)
 					a.Offset = offset
 
 					alerts = append(alerts, a)

internal/nlp/prose.go

@@ -1,7 +1,6 @@
 package nlp
 
 import (
-	"strings"
 	"sync"
 
 	"github.com/jdkato/prose/v3/segment"
@@ -47,38 +46,55 @@ var SentenceTokenizer sentenceTokenizer
 // a plain check-then-assign here is a data race.
 var tagger = sync.OnceValues(tag.New)
 
-// wordTokenizer splits a sentence into words for tagging.
-var wordTokenizer = sync.OnceValue(tokenize.NewTreebankWordTokenizer)
+// wordTokenizer splits a sentence into positioned words for tagging.
+//
+// prose's tokenizer rather than the Treebank one: Treebank rewrites the text
+// as it splits (quotes become “ and ”), so its tokens are not substrings of
+// the source and cannot carry offsets. Callers such as the `sequence` check
+// need to know where a token actually is.
+var wordTokenizer = sync.OnceValue(func() *tokenize.Tokenizer {
+	return tokenize.New()
+})
 
-// doTag assigns part-of-speech tags to `words`.
-func doTag(words []string) []tag.Token {
+// tagText splits text into sentences, tags each one, and returns the tokens
+// with offsets relative to text.
+//
+// Tagging is per sentence because the tagger conditions on the previous two
+// tags; letting that context run across a sentence boundary would condition
+// the first word of each sentence on the last word of the one before it.
+func tagText(text string) []tag.Token {
 	t, err := tagger()
 	if err != nil {
 		panic("nlp: loading the part-of-speech model: " + err.Error())
 	}
-	return t.Tag(words)
-}
 
-// textToWords convert raw text into a slice of words.
-func textToWords(text string, nlp bool) []string {
-	words := []string{}
-	for _, s := range SentenceTokenizer.Segment(text) {
-		if nlp {
-			words = append(words, wordTokenizer().Tokenize(s)...)
-		} else {
-			words = append(words, strings.Fields(s)...)
+	var tokens []tag.Token
+	for _, sent := range punktSegmenter().Segment(text) {
+		found := wordTokenizer().Tokenize(sent.Text)
+		t.TagTokens(found)
+
+		// Tokenize reported offsets within the sentence; shift them so they
+		// address the text the caller passed in.
+		for i := range found {
+			found[i].Start += sent.Start
 		}
+		tokens = append(tokens, found...)
 	}
 
-	return words
+	return tokens
 }
 
-// TextToTokens converts a string to a slice of tokens.
+// TextToTokens converts a string to a slice of tagged tokens.
+//
+// Tokens from the built-in tagger carry their byte offset within text, so
+// text[tok.Start:tok.Start+len(tok.Text)] == tok.Text. Tokens from a remote
+// NLP endpoint do not: that API returns text and tags only, so Start is zero
+// throughout and callers needing positions must locate the tokens themselves.
 func TextToTokens(text string, nlp *Info) []tag.Token {
 	// Determine if (and how) we need to do POS tagging.
 	if nlp == nil || nlp.Endpoint == "" {
 		// Fall back to our internal library (English-only).
-		return doTag(textToWords(text, true))
+		return tagText(text)
 	}
 	result, err := pos(text, nlp.Lang, nlp.Endpoint)
 	if err != nil {