fix(parser): bound folding depth and macro expansion work

Assisted-by: GLM 5.3

parser/preproc.go

@@ -51,7 +51,9 @@ func ParseWithOptions(path, src string, opts Options) (*ast.File, []error) {
 	if opts.Expand {
 		pp := &preproc{opts: opts, macros: map[string]*macroDef{}}
 		for name, value := range opts.Predefines {
-			pp.macros[name] = &macroDef{name: name, body: lexer.Tokenize(value)}
+			body := lexer.Tokenize(value)
+			pp.macros[name] = &macroDef{name: name, body: body}
+			pp.bodyTokens += len(body)
 		}
 		lines = pp.fileLines(path, tokens, token.Position{})
 		errs = pp.errs
@@ -67,6 +69,17 @@ func ParseWithOptions(path, src string, opts Options) (*ast.File, []error) {
 // assembler gives up after 100 nested invocations without producing a token.
 const maxExpansionDepth = 100
 
+// expandWorkFactor and expandWorkFloor size the per-line expansion work
+// budget: a generous multiple of the line and of everything the macro table
+// can inject into it.  No real file approaches the budget, because real
+// macros multiply their input by a modest factor, while amplification that
+// grows exponentially with nesting pays for every token it produces and
+// stops at the budget instead of running the machine for hours.
+const (
+	expandWorkFactor = 64
+	expandWorkFloor  = 4096
+)
+
 // textflagHeader names the one header gasm does not splice: its flag macros
 // (NOSPLIT, RODATA, …) are consumed by name throughout gasm's parser,
 // encoders and linter, and expanding them to their numeric constants would
@@ -91,6 +104,7 @@ type preproc struct {
 	errs       []error
 	stack      []string // absolute paths of files being read, innermost last
 	ifdefStack []bool   // one entry per open #ifdef/#ifndef, its truth
+	bodyTokens int      // total length of every defined macro body
 }
 
 // enabled reports whether the position being read is inside a live
@@ -253,13 +267,15 @@ func (pp *preproc) define(line []token.Token) {
 	if i := slices.IndexFunc(body, func(t token.Token) bool { return t.Kind == token.Comment }); i >= 0 {
 		body = body[:i]
 	}
-	if _, exists := pp.macros[name.Text]; exists {
+	if old, exists := pp.macros[name.Text]; exists {
 		// The toolchain refuses redefinition, so a file the oracle accepts
 		// never redefines; failing here keeps that contract visible.
 		pp.errorf(name.Pos, "redefinition of macro %s", name.Text)
+		pp.bodyTokens -= len(old.body)
 	}
-	pp.macros[name.Text] = &macroDef{name: name.Text, args: args, body: pp.bodyWithBreaks(body)}
-
+	stored := pp.bodyWithBreaks(body)
+	pp.bodyTokens += len(stored)
+	pp.macros[name.Text] = &macroDef{name: name.Text, args: args, body: stored}
 }
 
 // bodyWithBreaks records the statement boundaries the continuations carry.
@@ -285,10 +301,12 @@ func (pp *preproc) undef(line []token.Token) {
 	if len(line) < 3 || line[2].Kind != token.Ident {
 		return
 	}
-	if _, ok := pp.macros[line[2].Text]; !ok {
+	def, ok := pp.macros[line[2].Text]
+	if !ok {
 		pp.errorf(line[2].Pos, "#undef for undefined macro %s", line[2].Text)
 		return
 	}
+	pp.bodyTokens -= len(def.body)
 	delete(pp.macros, line[2].Text)
 }
 
@@ -342,26 +360,73 @@ func (pp *preproc) resolve(name, dir string) (string, bool) {
 	return "", false
 }
 
+// tokenFrame is one layer of the pending token stream during expansion: a
+// slice of tokens plus how far it has been read.
+type tokenFrame struct {
+	toks []token.Token
+	i    int
+}
+
 // expandTokens expands every macro invocation in a token sequence,
-// recursively, with a depth guard.  A body is spliced into the sequence in
-// place and rescanned, the way the toolchain's input stack re-reads pushed
-// tokens: an object macro may name a parameterised one, and the argument
-// list of the expansion may then come from the tokens that follow.
+// recursively, with a depth guard, the way the toolchain's input stack
+// re-reads pushed tokens: an object macro may name a parameterised one, and
+// the argument list of the expansion may then come from the tokens that
+// follow.  The pending stream is a stack of frames, so an expansion pushes
+// its body as the next thing to read instead of splicing it into one flat
+// slice: a long run of invocations costs work proportional to what it
+// produces, never the square of the line.  Work is billed against a budget,
+// because a body that repeats its argument multiplies every nesting level;
+// amplification that outgrows the budget is an error, not an hours-long
+// machine commitment.
 func (pp *preproc) expandTokens(in []token.Token) []token.Token {
-	s := in
-	i := 0
+	stack := []tokenFrame{{toks: in}}
+	out := make([]token.Token, 0, len(in))
 	consecutive := 0
-	for i < len(s) {
-		t := s[i]
+	budget := expandWorkFactor*(len(in)+pp.bodyTokens) + expandWorkFloor
+	spent := func(n int) bool {
+		budget -= n
+		return budget < 0
+	}
+
+	// nextAfter returns the first unread token behind the invocation at the
+	// top of the stack, together with the frame and offset it sits at.
+	nextAfter := func() (fi, off int, tok token.Token, ok bool) {
+		for j, s := range slices.Backward(stack) {
+			start := s.i
+			if j == len(stack)-1 {
+				start++
+			}
+			if start < len(s.toks) {
+				return j, start, s.toks[start], true
+			}
+		}
+		return 0, 0, token.Token{}, false
+	}
+
+	for len(stack) > 0 {
+		top := &stack[len(stack)-1]
+		if top.i >= len(top.toks) {
+			stack = stack[:len(stack)-1]
+			continue
+		}
+		t := top.toks[top.i]
 		if t.Kind != token.Ident {
-			i++
+			out = append(out, t)
+			top.i++
 			consecutive = 0
+			if spent(1) {
+				break
+			}
 			continue
 		}
 		def, suffix := pp.macroFor(t.Text)
 		if def == nil {
-			i++
+			out = append(out, t)
+			top.i++
 			consecutive = 0
+			if spent(1) {
+				break
+			}
 			continue
 		}
 		// The guard mirrors the toolchain's: 100 nested invocations in a
@@ -372,6 +437,7 @@ func (pp *preproc) expandTokens(in []token.Token) []token.Token {
 			return nil
 		}
 		if def.args == nil {
+			top.i++
 			body := restamp(def.body, t.Pos)
 			if suffix != "" {
 				// The macro was reached only through a compound spelling
@@ -379,27 +445,42 @@ func (pp *preproc) expandTokens(in []token.Token) []token.Token {
 				// to travel with the expansion.
 				body = appendSelector(body, suffix, t.Pos)
 			}
-			s = append(s[:i], append(body, s[i+1:]...)...)
+			stack = append(stack, tokenFrame{toks: body})
+			if spent(len(body)) {
+				break
+			}
 			continue
 		}
 		// A parameterised macro invoked without its parentheses stands
-		// unexpanded, naming itself, as in the toolchain.
-		if i+1 >= len(s) || s[i+1].Kind != token.LParen {
-			i++
+		// unexpanded, naming itself, as in the toolchain.  The parenthesis
+		// may sit past the end of this body, in the pending frames behind
+		// it, exactly where the toolchain's input stack would find it.
+		if _, _, nxt, ok := nextAfter(); !ok || nxt.Kind != token.LParen {
+			out = append(out, t)
+			top.i++
 			consecutive = 0
+			if spent(1) {
+				break
+			}
 			continue
 		}
-		args, next := pp.collectArgs(s, i+1, t)
-		if args == nil {
+		args, fi, off, ok := pp.collectArgs(stack, t)
+		if !ok {
 			return nil
 		}
+		if spent(collectedTokens(args)) {
+			break
+		}
 		// A zero-argument macro may be invoked as NAME().
 		if len(def.args) == 0 && len(args) == 1 && len(args[0]) == 0 {
 			args = nil
 		}
 		if len(args) != len(def.args) {
 			pp.errorf(t.Pos, "wrong arg count for macro %s: got %d, want %d", t.Text, len(args), len(def.args))
-			i = next
+			// Skip the invocation: drop the frames it consumed and resume
+			// right after its closing parenthesis.
+			stack = stack[:fi+1]
+			stack[fi].i = off
 			consecutive = 0
 			continue
 		}
@@ -424,9 +505,31 @@ func (pp *preproc) expandTokens(in []token.Token) []token.Token {
 			}
 			sub = append(sub, bt)
 		}
-		s = append(s[:i], append(sub, s[next:]...)...)
+		// The invocation consumed every frame down to its closing
+		// parenthesis; resume there, with the substitution read first.
+		stack = stack[:fi+1]
+		stack[fi].i = off
+		stack = append(stack, tokenFrame{toks: sub})
+		if spent(len(sub)) {
+			break
+		}
 	}
-	return s
+	if budget < 0 {
+		pp.errorf(in[0].Pos, "macro expansion exceeds the work budget of %d tokens",
+			expandWorkFactor*(len(in)+pp.bodyTokens)+expandWorkFloor)
+		return nil
+	}
+	return out
+}
+
+// collectedTokens counts the tokens an argument list carried off the pending
+// stream, so the expansion's work budget pays for reading them.
+func collectedTokens(args [][]token.Token) int {
+	n := 0
+	for _, a := range args {
+		n += len(a)
+	}
+	return n
 }
 
 // macroFor finds the macro a use names.  The lexer folds NAME.selector into
@@ -498,42 +601,57 @@ func pasteSelector(val []token.Token, suffix string) []token.Token {
 	return append(out, token.Token{Kind: token.Ident, Text: suffix})
 }
 
-// collectArgs reads the actual argument tokens of an invocation; the opening
-// parenthesis is at start.  Commas separate arguments except inside nested
-// parentheses.  A nil result means the list was unterminated, which is a
-// diagnostic.
-func (pp *preproc) collectArgs(in []token.Token, start int, name token.Token) ([][]token.Token, int) {
-	var args [][]token.Token
+// collectArgs reads the actual argument tokens of an invocation at the top
+// of the stack, whose opening parenthesis is the first unread token behind
+// the invocation.  Commas separate arguments except inside nested
+// parentheses, and the list may run on into the pending frames below, the
+// way the toolchain's input stack keeps reading pushed-back tokens.  It
+// returns the arguments together with the frame and offset of the first
+// token after the closing parenthesis; ok is false when the list never
+// closes, which is a diagnostic.
+func (pp *preproc) collectArgs(stack []tokenFrame, name token.Token) (args [][]token.Token, fi, off int, ok bool) {
 	var cur []token.Token
 	nesting := 0
-	for i := start + 1; i < len(in); i++ {
-		t := in[i]
-		switch t.Kind {
-		case token.LParen:
-			nesting++
-			cur = append(cur, t)
-		case token.RParen:
-			if nesting == 0 {
-				return append(args, cur), i + 1
-			}
-			nesting--
-			cur = append(cur, t)
-		case token.Comma:
-			if nesting == 0 {
-				args = append(args, cur)
-				cur = nil
+	started := false
+	for j, s := range slices.Backward(stack) {
+		start := s.i
+		if j == len(stack)-1 {
+			start++ // past the invocation's name
+		}
+		for k := start; k < len(s.toks); k++ {
+			t := s.toks[k]
+			if !started {
+				// The opening parenthesis itself.
+				started = true
 				continue
 			}
-			cur = append(cur, t)
-		case token.Comment:
-			pp.errorf(name.Pos, "unterminated arg list invoking macro %s", name.Text)
-			return nil, i
-		default:
-			cur = append(cur, t)
+			switch t.Kind {
+			case token.LParen:
+				nesting++
+				cur = append(cur, t)
+			case token.RParen:
+				if nesting == 0 {
+					return append(args, cur), j, k + 1, true
+				}
+				nesting--
+				cur = append(cur, t)
+			case token.Comma:
+				if nesting == 0 {
+					args = append(args, cur)
+					cur = nil
+					continue
+				}
+				cur = append(cur, t)
+			case token.Comment:
+				pp.errorf(name.Pos, "unterminated arg list invoking macro %s", name.Text)
+				return nil, 0, 0, false
+			default:
+				cur = append(cur, t)
+			}
 		}
 	}
 	pp.errorf(name.Pos, "unterminated arg list invoking macro %s", name.Text)
-	return nil, len(in)
+	return nil, 0, 0, false
 }
 
 // restamp copies body tokens to the invocation's position, so diagnostics