feat: ast.Program.Lines — the source line of every statement

ast/ast.go

@@ -10,7 +10,27 @@ import "math/big"
 type Node interface{ node() }
 
 // Program is the top-level sequence of expressions.
-type Program struct{ Body []Node }
+//
+// Lines gives the 1-based source line each STATEMENT node in the tree began
+// on, keyed by the node itself (every Node is a distinct pointer, so the map is
+// exact). It is how a consumer says WHERE something is: a backtrace, Ruby's
+// __LINE__, Module#const_source_location. A node with no entry reads back as
+// line 0, which is what MRI's rb_iseq_line_no returns for a pc its own table
+// cannot place (iseq.c v3_4_0:2314).
+//
+// The granularity is deliberate. MRI attaches a line to every node and compiles
+// a per-instruction table from it (compile.c v3_4_0: ADD_INSN records
+// nd_line(node); iseq.c:673 rb_iseq_insns_info_encode_positions builds
+// insns_info + positions), but that table is COMPRESSED — one entry per line
+// CHANGE, not per instruction — and every reader resolves a pc to the nearest
+// preceding entry. Statement boundaries are where the line changes in all but
+// pathological source, so this map holds the same information at the
+// granularity a consumer can observe, without a position field on all 55 node
+// types.
+type Program struct {
+	Body  []Node
+	Lines map[Node]int
+}
 
 // IntLit is an integer literal that fits in int64.
 type IntLit struct{ Value int64 }

lines_test.go

@@ -0,0 +1,218 @@
+// Copyright (c) the go-ruby-parser authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+package parser
+
+import (
+	"testing"
+
+	"github.com/go-ruby-parser/parser/ast"
+)
+
+// lineOf walks the statement list and returns the recorded line of the nth
+// top-level statement.
+func lineOf(t *testing.T, lines map[ast.Node]int, n ast.Node) int {
+	t.Helper()
+	l, ok := lines[n]
+	if !ok {
+		t.Fatalf("no line recorded for %T", n)
+	}
+	return l
+}
+
+// TestParsedLinesTopLevel pins the line of each top-level statement,
+// including the blank lines and comments MRI skips over.
+func TestParsedLinesTopLevel(t *testing.T) {
+	src := "a = 1\n\n# comment\nb = 2\nc = 3; d = 4\n"
+	prog, err := Parse(src)
+	if err != nil {
+		t.Fatal(err)
+	}
+	lines := prog.Lines
+	want := []int{1, 4, 5, 5}
+	if len(prog.Body) != len(want) {
+		t.Fatalf("got %d statements, want %d", len(prog.Body), len(want))
+	}
+	for i, w := range want {
+		if got := lineOf(t, lines, prog.Body[i]); got != w {
+			t.Errorf("statement %d: line %d, want %d", i, got, w)
+		}
+	}
+}
+
+// TestParsedLinesNested checks that a body nested inside a def, a class and
+// a block is stamped with its own line, not the enclosing construct's.
+func TestParsedLinesNested(t *testing.T) {
+	src := "def foo\n  bar\nend\nclass C\n  def baz\n    qux\n  end\nend\n[1].each do |v|\n  p v\nend\n"
+	prog, err := Parse(src)
+	if err != nil {
+		t.Fatal(err)
+	}
+	lines := prog.Lines
+	def, ok := prog.Body[0].(*ast.MethodDef)
+	if !ok {
+		t.Fatalf("statement 0 is %T, want *ast.MethodDef", prog.Body[0])
+	}
+	if got := lineOf(t, lines, def); got != 1 {
+		t.Errorf("def line %d, want 1", got)
+	}
+	if got := lineOf(t, lines, def.Body[0]); got != 2 {
+		t.Errorf("def body line %d, want 2", got)
+	}
+
+	cls, ok := prog.Body[1].(*ast.ClassDef)
+	if !ok {
+		t.Fatalf("statement 1 is %T, want *ast.ClassDef", prog.Body[1])
+	}
+	if got := lineOf(t, lines, cls); got != 4 {
+		t.Errorf("class line %d, want 4", got)
+	}
+	inner, ok := cls.Body[0].(*ast.MethodDef)
+	if !ok {
+		t.Fatalf("class body 0 is %T, want *ast.MethodDef", cls.Body[0])
+	}
+	if got := lineOf(t, lines, inner); got != 5 {
+		t.Errorf("inner def line %d, want 5", got)
+	}
+	if got := lineOf(t, lines, inner.Body[0]); got != 6 {
+		t.Errorf("inner def body line %d, want 6", got)
+	}
+
+	call, ok := prog.Body[2].(*ast.Call)
+	if !ok {
+		t.Fatalf("statement 2 is %T, want *ast.Call", prog.Body[2])
+	}
+	if got := lineOf(t, lines, call); got != 9 {
+		t.Errorf("each call line %d, want 9", got)
+	}
+	if got := lineOf(t, lines, call.Block.Body[0]); got != 10 {
+		t.Errorf("block body line %d, want 10", got)
+	}
+}
+
+// TestParsedLinesFirstWriteWins covers the pass-through construct: `begin;
+// foo; end` hands foo's node straight back, and the stamp must stay foo's line
+// (the inner one, written first) rather than the begin's.
+func TestParsedLinesFirstWriteWins(t *testing.T) {
+	src := "x = 0\nbegin\n  foo\nend\n"
+	prog, err := Parse(src)
+	if err != nil {
+		t.Fatal(err)
+	}
+	lines := prog.Lines
+	begin, ok := prog.Body[1].(*ast.Begin)
+	if !ok {
+		t.Fatalf("statement 1 is %T, want *ast.Begin", prog.Body[1])
+	}
+	if got := lineOf(t, lines, begin.Body[0]); got != 3 {
+		t.Errorf("begin body line %d, want 3", got)
+	}
+}
+
+// TestParsedLinesMultiline pins the rule a consumer depends on: a statement
+// spanning several lines carries the line it STARTED on, which is the line MRI
+// attributes to the call node too.
+func TestParsedLinesMultiline(t *testing.T) {
+	src := "foo(1,\n    2,\n    3)\n"
+	prog, err := Parse(src)
+	if err != nil {
+		t.Fatal(err)
+	}
+	lines := prog.Lines
+	if got := lineOf(t, lines, prog.Body[0]); got != 1 {
+		t.Errorf("multi-line call line %d, want 1", got)
+	}
+}
+
+// TestParsedLinesModifier checks a statement rewritten by applyModifiers:
+// `foo if bar` becomes an *ast.If, and that If is what carries the line.
+func TestParsedLinesModifier(t *testing.T) {
+	src := "\n\nfoo if bar\n"
+	prog, err := Parse(src)
+	if err != nil {
+		t.Fatal(err)
+	}
+	lines := prog.Lines
+	if _, ok := prog.Body[0].(*ast.If); !ok {
+		t.Fatalf("statement 0 is %T, want *ast.If", prog.Body[0])
+	}
+	if got := lineOf(t, lines, prog.Body[0]); got != 3 {
+		t.Errorf("modifier-if line %d, want 3", got)
+	}
+}
+
+// TestParsedLinesKeywordStatements covers the non-default arms of
+// parseStatement1 (return/break/next/retry/alias/undef), each of which is
+// stamped through the same wrapper.
+func TestParsedLinesKeywordStatements(t *testing.T) {
+	src := "def m\n  return 1\nend\n[1].each do\n  break\nend\n[1].each do\n  next\nend\nbegin\n  retry\nrescue\nend\nalias a b\nundef c\n"
+	prog, err := Parse(src)
+	if err != nil {
+		t.Fatal(err)
+	}
+	lines := prog.Lines
+	def := prog.Body[0].(*ast.MethodDef)
+	if got := lineOf(t, lines, def.Body[0]); got != 2 {
+		t.Errorf("return line %d, want 2", got)
+	}
+	brk := prog.Body[1].(*ast.Call)
+	if got := lineOf(t, lines, brk.Block.Body[0]); got != 5 {
+		t.Errorf("break line %d, want 5", got)
+	}
+	nxt := prog.Body[2].(*ast.Call)
+	if got := lineOf(t, lines, nxt.Block.Body[0]); got != 8 {
+		t.Errorf("next line %d, want 8", got)
+	}
+	bg := prog.Body[3].(*ast.Begin)
+	if got := lineOf(t, lines, bg.Body[0]); got != 11 {
+		t.Errorf("retry line %d, want 11", got)
+	}
+	if got := lineOf(t, lines, prog.Body[4]); got != 14 {
+		t.Errorf("alias line %d, want 14", got)
+	}
+	if got := lineOf(t, lines, prog.Body[5]); got != 15 {
+		t.Errorf("undef line %d, want 15", got)
+	}
+}
+
+// TestParsedLinesError checks that a malformed input yields no program and no
+// map, like Parse's error contract.
+func TestParsedLinesError(t *testing.T) {
+	prog, err := Parse("def\n")
+	if err == nil {
+		t.Fatal("want a parse error")
+	}
+	if prog != nil {
+		t.Errorf("got prog=%v, want nil", prog)
+	}
+}
+
+// TestParsedLinesEmptyProgram pins that an empty source still yields a usable
+// (non-nil, empty) map rather than one a consumer must nil-check.
+func TestParsedLinesEmptyProgram(t *testing.T) {
+	prog, err := Parse("\n# just a comment\n")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if prog.Lines == nil {
+		t.Error("Lines is nil for an empty program")
+	}
+	if len(prog.Lines) != 0 {
+		t.Errorf("Lines has %d entries, want 0", len(prog.Lines))
+	}
+}
+
+// TestParsedLinesInternalError drives the recover path of the shared parse
+// core through ParsedLines, so the nil-map assignment there is exercised.
+func TestParsedLinesInternalError(t *testing.T) {
+	parseHook = func() { panic("boom") }
+	defer func() { parseHook = nil }()
+	prog, err := Parse("a = 1\n")
+	if err == nil {
+		t.Fatal("want an internal parse error")
+	}
+	if prog != nil {
+		t.Errorf("got prog=%v, want nil", prog)
+	}
+}

parser.go

@@ -146,6 +146,9 @@ type Parser struct {
 	// next line (`f(`⏎`  k:`⏎`    v)`) is then a continued pair, not a value-omitted
 	// shorthand; at top level (depth 0) a newline after `key:` ends the command.
 	bracketDepth int
+	// lines collects the 1-based source line each statement node began on; it
+	// becomes ast.Program.Lines.
+	lines map[ast.Node]int
 }
 
 // parseHook, when non-nil, runs at the start of Parse. It exists only so a
@@ -157,8 +160,19 @@ var parseHook func()
 // yields a parse error, and any unexpected internal panic is also surfaced as a
 // parse error rather than propagating to the caller.
 func Parse(src string) (prog *ast.Program, err error) {
+
 	toks := lexer.New(src).Tokenize()
-	p := &Parser{toks: toks, scopes: []*scope{newScope(true)}, argParenEnd: -1}
+	// One statement per source line is the dense case, and the lexer emits one
+	// NEWLINE per line (and per `;`), so the terminator count is a tight upper
+	// bound on the entries: sizing the map by it removes the rehash growth that
+	// otherwise makes the table cost several times what it stores.
+	terms := 0
+	for _, t := range toks {
+		if t.Type == token.NEWLINE {
+			terms++
+		}
+	}
+	p := &Parser{toks: toks, scopes: []*scope{newScope(true)}, argParenEnd: -1, lines: make(map[ast.Node]int, terms)}
 	defer func() {
 		if r := recover(); r != nil {
 			if pe, ok := r.(parseError); ok {
@@ -176,7 +190,7 @@ func Parse(src string) (prog *ast.Program, err error) {
 	}
 	body := p.parseStatements(map[token.Type]bool{})
 	p.expect(token.EOF)
-	return &ast.Program{Body: body}, nil
+	return &ast.Program{Body: body, Lines: p.lines}, nil
 }
 
 // --- token cursor ---
@@ -347,7 +361,27 @@ func (p *Parser) parseStatements(stop map[token.Type]bool) []ast.Node {
 	return body
 }
 
+// parseStatement parses one statement and records the line it started on into
+// Program.Lines.
+//
+// The stamp is first-write-wins, and that is what makes it precise: an inner
+// statement is stamped when ITS parseStatement returns, before the enclosing one
+// finishes, so a construct that yields its body's node straight through
+// (`begin; foo; end`) keeps foo's own line rather than the begin's. MRI reaches
+// the same answer from the other side — the node it compiles IS foo's node, and
+// nd_line(foo) is foo's line.
 func (p *Parser) parseStatement() ast.Node {
+	line := p.cur().Line
+	n := p.parseStatement1()
+	if n != nil {
+		if _, seen := p.lines[n]; !seen {
+			p.lines[n] = line
+		}
+	}
+	return n
+}
+
+func (p *Parser) parseStatement1() ast.Node {
 	switch p.cur().Type {
 	case token.RETURN:
 		return p.applyModifiers(p.parseReturn())