fix: cache countLines. add comments. (#34)

Author: Hoblovski

lang/collect/export.go

@@ -40,11 +40,12 @@ func (c *Collector) fileLine(loc Location) uniast.FileLine {
 		rel = filepath.Base(loc.URI.File())
 	}
 	text := c.cli.GetFile(loc.URI).Text
+	file_uri := string(loc.URI)
 	return uniast.FileLine{
 		File:        rel,
 		Line:        loc.Range.Start.Line,
-		StartOffset: lsp.PositionOffset(text, loc.Range.Start),
-		EndOffset:   lsp.PositionOffset(text, loc.Range.End),
+		StartOffset: lsp.PositionOffset(file_uri, text, loc.Range.Start),
+		EndOffset:   lsp.PositionOffset(file_uri, text, loc.Range.End),
 	}
 }
 

lang/cxx/spec.go

@@ -166,7 +166,7 @@ func (c *CxxSpec) FunctionSymbol(sym lsp.DocumentSymbol) (int, []int, []int, []i
 	// TODO: attributes may contain parens. also inline structs.
 
 	endRelOffset := 0
-	lines := utils.CountLinesCached(sym.Text)
+	lines := utils.CountLinesPooled(sym.Text)
 	phase := 0
 	for i, tok := range sym.Tokens {
 		switch phase {

lang/lsp/utils.go

@@ -36,7 +36,7 @@ import (
 )
 
 func GetDistance(text string, start Position, pos Position) int {
-	lines := utils.CountLinesCached(text)
+	lines := utils.CountLinesPooled(text)
 	defer utils.PutCount(lines)
 	// find the line of the position
 	return (*lines)[pos.Line-start.Line] + pos.Character - start.Character
@@ -59,13 +59,15 @@ func RelativePostionWithLines(lines []int, textPos Position, pos Position) int {
 	return lines[l] + pos.Character - textPos.Character
 }
 
-func PositionOffset(text string, pos Position) int {
+func PositionOffset(file_uri string, text string, pos Position) int {
 	if pos.Line < 0 || pos.Character < 0 {
 		log.Error("invalid text position: %+v", pos)
 		return -1
 	}
-	lines := utils.CountLinesCached(text)
-	defer utils.PutCount(lines)
+	lines := utils.CountLinesCached(file_uri, text)
+
+	// lines := utils.CountLinesPooled(text)
+	// defer utils.PutCount(lines)
 
 	return RelativePostionWithLines(*lines, Position{Line: 0, Character: 0}, pos)
 }

lang/lsp/utils_test.go

@@ -73,7 +73,7 @@ func TestPositionOffset(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			result := PositionOffset(tt.text, tt.pos)
+			result := PositionOffset(tt.name, tt.text, tt.pos)
 			if result != tt.expected {
 				t.Errorf("PositionOffset() = %v, expected %v", result, tt.expected)
 			}

lang/rust/spec.go

@@ -261,7 +261,7 @@ func (c *RustSpec) FunctionSymbol(sym lsp.DocumentSymbol) (int, []int, []int, []
 	if where == -1 {
 		where = len(tokens) - 1
 	}
-	lines := utils.CountLinesCached(sym.Text)
+	lines := utils.CountLinesPooled(sym.Text)
 
 	// find the typeParam's type token between "fn" and "("
 	var typeParams []int

lang/utils/strings.go

@@ -23,12 +23,40 @@ var countPool = sync.Pool{
 	},
 }
 
+var countCache = sync.Map{}
+
+// USE AND ONLY USE in pairs with CountLinesPooled.
 func PutCount(count *[]int) {
 	*count = (*count)[:0]
 	countPool.Put(count)
 }
 
-func CountLinesCached(text string) *[]int {
+// The cached version of CountLines.
+// Avoids redundant computations for the same text.
+// Use when the same text is processed multiple times, such as contents of a file.
+//
+// The key MUST uniquely identify the text i.e. for any two invocations
+//
+//	CountLinesCached(key1, text1) and CountLinesCached(key2, text2),
+//
+// if key1 == key2, then text1 must be equal to text2 (and also vice versa).
+func CountLinesCached(key string, text string) *[]int {
+	if cached, ok := countCache.Load(key); ok {
+		res := cached.([]int)
+		return &res
+	}
+
+	lines := CountLines(text)
+	countCache.Store(key, lines)
+	return &lines
+}
+
+// The pooled version of CountLines.
+// Eases burden on allocation and GC.
+// Use when invocation on small text is frequent.
+//
+// MUST manually invoke `PutCount` when done with the result to return the slice to the pool.
+func CountLinesPooled(text string) *[]int {
 	tmp := countPool.Get().(*[]int)
 	*tmp = append(*tmp, 0)
 	for i, c := range text {
@@ -39,6 +67,9 @@ func CountLinesCached(text string) *[]int {
 	return tmp
 }
 
+// CountLines calculates the starting offsets of lines in a given text.
+// Each offset marks the byte position of the character immediately following a newline character,
+// or 0 for the very beginning of the text.
 func CountLines(text string) []int {
 	var ret []int
 	ret = append(ret, 0)