Refactor error handling for FOR and COLLECT statements: enhance error messages for missing variables, improve error span tracking, and streamline error analysis for better clarity and maintainability.

pkg/compiler/internal/diagnostics/error_listener.go

@@ -38,7 +38,11 @@ func (d *ErrorListener) SyntaxError(_ antlr.Recognizer, offendingSymbol interfac
 		offending = tok
 	}
 
-	d.handler.Add(d.parseError(msg, offending))
+	if err := d.parseError(msg, offending); err != nil {
+		if !d.handler.HasErrorOnLine(line) {
+			d.handler.Add(err)
+		}
+	}
 }
 
 func (d *ErrorListener) parseError(msg string, offending antlr.Token) *CompilationError {
@@ -55,8 +59,9 @@ func (d *ErrorListener) parseError(msg string, offending antlr.Token) *Compilati
 	}
 
 	for _, handler := range []func(*CompilationError) bool{
-		d.extraneousError,
-		d.noViableAltError,
+		d.extraneousInput,
+		d.noViableAlternative,
+		d.mismatchedInput,
 	} {
 		if handler(err) {
 			break
@@ -66,7 +71,7 @@ func (d *ErrorListener) parseError(msg string, offending antlr.Token) *Compilati
 	return err
 }
 
-func (d *ErrorListener) extraneousError(err *CompilationError) (matched bool) {
+func (d *ErrorListener) extraneousInput(err *CompilationError) (matched bool) {
 	if !strings.Contains(err.Message, "extraneous input") {
 		return false
 	}
@@ -88,10 +93,18 @@ func (d *ErrorListener) extraneousError(err *CompilationError) (matched bool) {
 	return true
 }
 
-func (d *ErrorListener) noViableAltError(err *CompilationError) bool {
+func (d *ErrorListener) noViableAlternative(err *CompilationError) bool {
 	if !strings.Contains(err.Message, "viable alternative at input") {
 		return false
 	}
 
 	return AnalyzeSyntaxError(d.src, err, d.history.Last())
 }
+
+func (d *ErrorListener) mismatchedInput(err *CompilationError) bool {
+	if !strings.Contains(err.Message, "mismatched input") {
+		return false
+	}
+
+	return AnalyzeSyntaxError(d.src, err, d.history.Last())
+}

pkg/compiler/internal/diagnostics/handler.go

@@ -9,9 +9,10 @@ import (
 )
 
 type ErrorHandler struct {
-	src       *file.Source
-	errors    []*CompilationError
-	threshold int
+	src             *file.Source
+	errors          []*CompilationError
+	linesWithErrors map[int]bool
+	threshold       int
 }
 
 func NewErrorHandler(src *file.Source, threshold int) *ErrorHandler {
@@ -20,9 +21,10 @@ func NewErrorHandler(src *file.Source, threshold int) *ErrorHandler {
 	}
 
 	return &ErrorHandler{
-		src:       src,
-		errors:    make([]*CompilationError, 0),
-		threshold: threshold,
+		src:             src,
+		errors:          make([]*CompilationError, 0),
+		linesWithErrors: make(map[int]bool),
+		threshold:       threshold,
 	}
 }
 
@@ -60,6 +62,13 @@ func (h *ErrorHandler) Add(err *CompilationError) {
 		err.Source = h.src
 	}
 
+	for _, span := range err.Spans {
+		if err.Source != nil {
+			line, _ := err.Source.LocationAt(span.Span)
+			h.linesWithErrors[line] = true
+		}
+	}
+
 	h.errors = append(h.errors, err)
 
 	if len(h.errors) == h.threshold {
@@ -71,6 +80,10 @@ func (h *ErrorHandler) Add(err *CompilationError) {
 	}
 }
 
+func (h *ErrorHandler) HasErrorOnLine(line int) bool {
+	return h.linesWithErrors[line]
+}
+
 func (h *ErrorHandler) VariableNotUnique(ctx antlr.ParserRuleContext, name string) {
 	// TODO: Add information where the variable was defined
 	h.Add(&CompilationError{

pkg/compiler/internal/diagnostics/helpers.go

@@ -107,3 +107,7 @@ func is(node *TokenNode, expected string) bool {
 
 	return strings.ToUpper(node.GetText()) == expected
 }
+
+func has(msg string, substr string) bool {
+	return strings.Contains(strings.ToLower(msg), strings.ToLower(substr))
+}

pkg/compiler/internal/diagnostics/match_for_loop_errors.go

@@ -46,5 +46,23 @@ func matchForLoopErrors(src *file.Source, err *CompilationError, offending *Toke
 		return true
 	}
 
+	if is(offending, "COLLECT") {
+		msg := err.Message
+
+		if has(msg, "COLLECT =") {
+			span := spanFromTokenSafe(offending.Token(), src)
+			span.Start = span.End
+			span.End = span.Start + 1
+
+			err.Message = "Expected variable before '=' in COLLECT"
+			err.Hint = "COLLECT must group by a variable."
+			err.Spans = []ErrorSpan{
+				NewMainErrorSpan(span, "missing variable"),
+			}
+
+			return true
+		}
+	}
+
 	return false
 }

pkg/file/snippet.go

@@ -23,8 +23,15 @@ func NewSnippetWithCaret(lines []string, span Span, line int) Snippet {
 	}
 
 	srcLine := lines[line-1]
-	startCol := computeVisualOffset(srcLine, span.Start)
-	endCol := computeVisualOffset(srcLine, span.End)
+	lineStartOffset := 0
+
+	// Compute actual start-of-line offset
+	for i := 0; i < line-1; i++ {
+		lineStartOffset += len(lines[i]) + 1 // +1 for \n
+	}
+
+	startCol := computeVisualOffset(srcLine, span.Start-lineStartOffset+1)
+	endCol := computeVisualOffset(srcLine, span.End-lineStartOffset+1)
 
 	caret := ""
 

test/integration/compiler/compiler_errors_syntax_test.go

@@ -67,8 +67,30 @@ func TestSyntaxErrors(t *testing.T) {
 				RETURN i
 		`, E{
 				Kind:    compiler.SyntaxError,
-				Message: "--",
-				Hint:    "Use 'FOR x IN [iterable]' syntax.",
+				Message: "Expected loop variable before 'IN'",
+				Hint:    "FOR must declare a variable.",
 			}, "FOR without variable"),
+		ErrorCase(
+			`
+			LET users = []
+			FOR x IN users
+				COLLECT =
+				RETURN x
+		`, E{
+				Kind:    compiler.SyntaxError,
+				Message: "Expected variable before '=' in COLLECT",
+				Hint:    "COLLECT must group by a variable.",
+			}, "COLLECT with no variable"),
+		ErrorCase(
+			`
+			LET users = []
+			FOR x IN users
+				COLLECT i =
+				RETURN x
+		`, E{
+				Kind:    compiler.SyntaxError,
+				Message: "Expected expression after '=' for variable 'i'",
+				Hint:    "Did you forget to provide a value?",
+			}, "COLLECT with no variable assignment"),
 	})
 }