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

Assisted-by: GLM 5.3

parser/expr.go

@@ -16,12 +16,21 @@ import (
 	"sourcedock.dev/petrbalvin/gasm-sdk/token"
 )
 
+// maxFoldDepth bounds the nesting the folder follows, counting parenthesis
+// groups and unary signs alike.  Real expressions nest a handful deep, and
+// beyond the bound the expression simply fails to fold, which lands it in
+// the ordinary operand paths.  Without a bound, a hostile file parks the
+// folder millions of frames down the mutual recursion of foldAdd, foldMul
+// and foldFactor, and the goroutine stack overflows with a fatal error no
+// caller can recover from.
+const maxFoldDepth = 1000
+
 // foldExpr evaluates the constant integer expression at the head of ts and
 // returns its value together with the unconsumed tokens.  ok is false when
 // the tokens do not form an expression, which is the callers' signal to use
 // the ordinary parsing paths.
 func foldExpr(ts []token.Token) (val int64, rest []token.Token, ok bool) {
-	v, rest, ok := foldAdd(ts)
+	v, rest, ok := foldAdd(ts, 0)
 	if !ok {
 		return 0, ts, false
 	}
@@ -30,8 +39,8 @@ func foldExpr(ts []token.Token) (val int64, rest []token.Token, ok bool) {
 
 // foldAdd parses addition-level expressions: +, - and | bind loosest, the
 // Plan 9 convention that makes x<<1|3 read as (x<<1)|3.
-func foldAdd(ts []token.Token) (int64, []token.Token, bool) {
+func foldAdd(ts []token.Token, depth int) (int64, []token.Token, bool) {
-	v, rest, ok := foldMul(ts)
+	v, rest, ok := foldMul(ts, depth)
 	if !ok {
 		return 0, ts, false
 	}
@@ -40,7 +49,7 @@ func foldAdd(ts []token.Token) (int64, []token.Token, bool) {
 		if kind != token.Plus && kind != token.Minus && kind != token.Pipe {
 			return v, rest, true
 		}
-		w, r2, ok := foldMul(rest[1:])
+		w, r2, ok := foldMul(rest[1:], depth)
 		if !ok {
 			return v, rest, true
 		}
@@ -59,43 +68,43 @@ func foldAdd(ts []token.Token) (int64, []token.Token, bool) {
 
 // foldMul parses multiplication-level expressions: *, / and the bit
 // operators &, << and >>.
-func foldMul(ts []token.Token) (int64, []token.Token, bool) {
+func foldMul(ts []token.Token, depth int) (int64, []token.Token, bool) {
-	v, rest, ok := foldFactor(ts)
+	v, rest, ok := foldFactor(ts, depth)
 	if !ok {
 		return 0, ts, false
 	}
 	for len(rest) > 0 {
 		switch rest[0].Kind {
 		case token.Star:
-			w, r2, ok := foldFactor(rest[1:])
+			w, r2, ok := foldFactor(rest[1:], depth)
 			if !ok {
 				return v, rest, true
 			}
 			v *= w
 			rest = r2
 		case token.Slash:
-			w, r2, ok := foldFactor(rest[1:])
+			w, r2, ok := foldFactor(rest[1:], depth)
 			if !ok || w == 0 {
 				return v, rest, true
 			}
 			v /= w
 			rest = r2
 		case token.Ampersand:
-			w, r2, ok := foldFactor(rest[1:])
+			w, r2, ok := foldFactor(rest[1:], depth)
 			if !ok {
 				return v, rest, true
 			}
 			v &= w
 			rest = r2
 		case token.LShift:
-			w, r2, ok := foldFactor(rest[1:])
+			w, r2, ok := foldFactor(rest[1:], depth)
 			if !ok || w < 0 || w >= 64 {
 				return v, rest, true
 			}
 			v <<= uint(w)
 			rest = r2
 		case token.RShift:
-			w, r2, ok := foldFactor(rest[1:])
+			w, r2, ok := foldFactor(rest[1:], depth)
 			if !ok || w < 0 || w >= 64 {
 				return v, rest, true
 			}
@@ -109,9 +118,11 @@ func foldMul(ts []token.Token) (int64, []token.Token, bool) {
 }
 
 // foldFactor parses a number, a parenthesised expression, or a unary sign
-// or complement.
+// or complement.  Each parenthesis group and each unary operator costs one
-func foldFactor(ts []token.Token) (int64, []token.Token, bool) {
+// depth level, so nesting beyond maxFoldDepth fails here instead of growing
-	if len(ts) == 0 {
+// the stack without end.
+func foldFactor(ts []token.Token, depth int) (int64, []token.Token, bool) {
+	if len(ts) == 0 || depth > maxFoldDepth {
 		return 0, ts, false
 	}
 	switch ts[0].Kind {
@@ -122,21 +133,21 @@ func foldFactor(ts []token.Token) (int64, []token.Token, bool) {
 		}
 		return v, ts[1:], true
 	case token.LParen:
-		v, rest, ok := foldAdd(ts[1:])
+		v, rest, ok := foldAdd(ts[1:], depth+1)
 		if !ok || len(rest) == 0 || rest[0].Kind != token.RParen {
 			return 0, ts, false
 		}
 		return v, rest[1:], true
 	case token.Minus:
-		v, rest, ok := foldFactor(ts[1:])
+		v, rest, ok := foldFactor(ts[1:], depth+1)
 		if !ok {
 			return 0, ts, false
 		}
 		return -v, rest, true
 	case token.Plus:
-		return foldFactor(ts[1:])
+		return foldFactor(ts[1:], depth+1)
 	case token.Tilde:
-		v, rest, ok := foldFactor(ts[1:])
+		v, rest, ok := foldFactor(ts[1:], depth+1)
 		if !ok {
 			return 0, ts, false
 		}

parser/fuzz_test.go

@@ -4,17 +4,17 @@
 package parser
 
 import (
+	"errors"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 )
 
-// FuzzParse hammers the parser with arbitrary input.  The contract: no panic,
+// corpusFiles seeds a fuzz target with the repository's kernels, so a plain
-// and always a usable file, whether or not diagnostics were reported.  The
+// `go test` run replays every seed as a regression case and CI exercises them
-// seed corpus carries the repository's kernels, so a plain `go test` run
+// without any fuzzing budget.
-// replays every seed as a regression case and CI exercises them without any
+func corpusFiles(f *testing.F) {
-// fuzzing budget.
-func FuzzParse(f *testing.F) {
 	for _, pattern := range []string{
 		"../testdata/*.s",
 		"../testdata/verify/*.s",
@@ -26,16 +26,93 @@ func FuzzParse(f *testing.F) {
 			}
 		}
 	}
+}
+
+// FuzzParse hammers the parser with arbitrary input.  The contract: no panic,
+// always a usable file whether or not diagnostics were reported, and every
+// diagnostic carries a real position and a non-empty message.
+func FuzzParse(f *testing.F) {
+	corpusFiles(f)
 	f.Add("TEXT ·f(SB), NOSPLIT, $0\n\tRET\n")
 	f.Add("garbage ### ??? ::: \xff\xfe\n")
 	f.Add("#define A(x) x+1\nTEXT ·f(SB), $0\n\tA(2)\n\tRET\n")
 	f.Add("DATA t<>+0(SB)/8, $1\nGLOBL t<>(SB), RODATA, $8\n")
 	f.Add("TEXT ·f(SB), $0\n\tJMP (AX)\n\tCALL (BX)\n\tRET\n")
+	f.Add("\xef\xbb\xbfTEXT ·f(SB), $0\n")                     // BOM
+	f.Add("CALL internal∕runtime∕atomic·Xchg(SB)\n")           // U+2215 path, U+00B7 name
+	f.Add("TEXT ·f(SB), $0\n\tJMP -3(PC)\n\tRET\n")            // negative PC jump
+	f.Add("DATA d<>+0(SB)/8, $0xFFFFFFFFFFFFFFFF\n")           // unsigned 64-bit immediate
+	f.Add("TEXT ·f(SB), $0-0xFFFFFFFFFFFFFFFF\n\tRET\n")       // argument area overflow
+	f.Add("TEXT ·f(SB), $0\n\tMOVD $(1<<0|1<<9), R0\n\tRET\n") // constant expression
+	f.Add("TEXT ·f(SB), $0\n\tADD $-8, R10; SUB $~63, R11\n\tRET\n")
+	// Constant-expression nesting: an immediate and a displacement read
+	// through the recursive folder, so deep parenthesis groups probe its
+	// depth handling on every plain run.  The immediate seed is deep enough
+	// that, before the folder carried a depth bound, it grew the goroutine
+	// stack past its limit with a fatal error.
+	f.Add("TEXT ·f(SB), $0\n\tMOVD $" + strings.Repeat("(", 4_000_000) + "1" + strings.Repeat(")", 4_000_000) + ", R0\n\tRET\n")
+	f.Add("TEXT ·f(SB), $0\n\tMOVQ " + strings.Repeat("(", 2000) + "1" + strings.Repeat(")", 2000) + "(AX), BX\n\tRET\n")
+	f.Add("TEXT ·f(SB), $0\n\tMOVD $" + strings.Repeat("~", 5000) + "1, R0\n\tRET\n")
 
 	f.Fuzz(func(t *testing.T, src string) {
-		file, _ := Parse("fuzz.s", src)
+		file, errs := Parse("fuzz.s", src)
 		if file == nil {
 			t.Fatal("Parse returned a nil file")
 		}
+		for _, err := range errs {
+			var perr Error
+			if !errors.As(err, &perr) {
+				t.Fatalf("diagnostic %v is not a parser Error", err)
+			}
+			if !perr.Pos.IsValid() {
+				t.Fatalf("diagnostic %q carries no position", perr.Msg)
+			}
+			if strings.TrimSpace(perr.Msg) == "" {
+				t.Fatalf("diagnostic at %v carries no message", perr.Pos)
+			}
+		}
+	})
+}
+
+// FuzzParseExpand hammers the preprocessing path, macro expansion and include
+// splicing included, with arbitrary input.  The contract is FuzzParse's: no
+// panic, a usable file, and diagnostics that name a place and a reason.
+func FuzzParseExpand(f *testing.F) {
+	corpusFiles(f)
+	f.Add("#define A 1\nTEXT ·f(SB), $0\n\tMOVD $A, R0\n\tRET\n")
+	f.Add("#define A(x) x+1\nA(2)\n")
+	f.Add("#define A(x) ((x))\nA(A(A(1)))\n")
+	f.Add("#define A\nTEXT ·f(SB), $0\n\tRET\n") // empty body
+	f.Add("#define A A\nA\n")                    // self-reference
+	f.Add("#define A B\n#define B A\nA\n")       // mutual recursion
+	f.Add("#ifdef X\n#else\n#endif\n")
+	f.Add("#ifndef X\none\n#else\ntwo\n#endif\n")
+	f.Add("#undef Y\n")
+	f.Add("#include \"textflag.h\"\n")
+	f.Add("#include \"no/such/header.h\"\n")
+	f.Add("#line 7 \"f.s\"\n")
+	f.Add("#define M(x) \\\n\tADD $x, R0 \\\n\tSUB $x, R1\nM(1)\n") // continuations
+	// A body that repeats its argument multiplies every nesting level; the
+	// expansion must stop at the work budget, not produce the exponential
+	// result it names.
+	f.Add("#define A(x) x x x x x x x x x x\nA(A(A(A(A(A(A(A(A(A(1))))))))))\n")
+
+	f.Fuzz(func(t *testing.T, src string) {
+		file, errs := ParseWithOptions("fuzz.s", src, Options{Expand: true})
+		if file == nil {
+			t.Fatal("ParseWithOptions returned a nil file")
+		}
+		for _, err := range errs {
+			var perr Error
+			if !errors.As(err, &perr) {
+				t.Fatalf("diagnostic %v is not a parser Error", err)
+			}
+			if !perr.Pos.IsValid() {
+				t.Fatalf("diagnostic %q carries no position", perr.Msg)
+			}
+			if strings.TrimSpace(perr.Msg) == "" {
+				t.Fatalf("diagnostic at %v carries no message", perr.Pos)
+			}
+		}
 	})
 }

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
+// recursively, with a depth guard, the way the toolchain's input stack
-// place and rescanned, the way the toolchain's input stack re-reads pushed
+// re-reads pushed tokens: an object macro may name a parameterised one, and
-// tokens: an object macro may name a parameterised one, and the argument
+// the argument list of the expansion may then come from the tokens that
-// list of the expansion may then come from the tokens that follow.
+// 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
+	stack := []tokenFrame{{toks: in}}
-	i := 0
+	out := make([]token.Token, 0, len(in))
 	consecutive := 0
-	for i < len(s) {
+	budget := expandWorkFactor*(len(in)+pp.bodyTokens) + expandWorkFloor
-		t := s[i]
+	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.
+		// unexpanded, naming itself, as in the toolchain.  The parenthesis
-		if i+1 >= len(s) || s[i+1].Kind != token.LParen {
+		// may sit past the end of this body, in the pending frames behind
-			i++
+		// 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)
+		args, fi, off, ok := pp.collectArgs(stack, t)
-		if args == nil {
+		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
+// collectArgs reads the actual argument tokens of an invocation at the top
-// parenthesis is at start.  Commas separate arguments except inside nested
+// of the stack, whose opening parenthesis is the first unread token behind
-// parentheses.  A nil result means the list was unterminated, which is a
+// the invocation.  Commas separate arguments except inside nested
-// diagnostic.
+// parentheses, and the list may run on into the pending frames below, the
-func (pp *preproc) collectArgs(in []token.Token, start int, name token.Token) ([][]token.Token, int) {
+// way the toolchain's input stack keeps reading pushed-back tokens.  It
-	var args [][]token.Token
+// 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++ {
+	started := false
-		t := in[i]
+	for j, s := range slices.Backward(stack) {
-		switch t.Kind {
+		start := s.i
-		case token.LParen:
+		if j == len(stack)-1 {
-			nesting++
+			start++ // past the invocation's name
-			cur = append(cur, t)
+		}
-		case token.RParen:
+		for k := start; k < len(s.toks); k++ {
-			if nesting == 0 {
+			t := s.toks[k]
-				return append(args, cur), i + 1
+			if !started {
-			}
+				// The opening parenthesis itself.
-			nesting--
+				started = true
-			cur = append(cur, t)
-		case token.Comma:
-			if nesting == 0 {
-				args = append(args, cur)
-				cur = nil
 				continue
 			}
-			cur = append(cur, t)
+			switch t.Kind {
-		case token.Comment:
+			case token.LParen:
-			pp.errorf(name.Pos, "unterminated arg list invoking macro %s", name.Text)
+				nesting++
-			return nil, i
+				cur = append(cur, t)
-		default:
+			case token.RParen:
-			cur = append(cur, t)
+				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

parser/preproc_test.go

@@ -8,6 +8,7 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 
 	"sourcedock.dev/petrbalvin/gasm-sdk/ast"
 )
@@ -58,6 +59,64 @@ func wantLines(t *testing.T, got []string, want ...string) {
 	}
 }
 
+// TestManyInvocationsExpandInLinearTime pins the cost contract of macro
+// expansion: a long run of invocations costs work proportional to what it
+// produces.  Splicing each expansion into one flat slice instead made every
+// invocation shift the whole remaining tail, so the 200000 invocations here
+// took minutes; the pending stream is a stack of frames since, measured at
+// roughly 120 ms on this machine, and ten seconds stays a generous ceiling
+// even on a runner an order of magnitude slower.
+func TestManyInvocationsExpandInLinearTime(t *testing.T) {
+	const n = 200000
+	src := "#define A MOVQ AX, BX\nTEXT ·f(SB), NOSPLIT, $0\n\t" +
+		strings.Repeat("A\n", n) + "\tRET\n"
+	start := time.Now()
+	f, errs := ParseWithOptions("t_amd64.s", src, Options{Expand: true})
+	if len(errs) > 0 {
+		t.Fatalf("parse: %v", errs)
+	}
+	var ts []*ast.Text
+	for _, d := range f.Decls {
+		if text, ok := d.(*ast.Text); ok {
+			ts = append(ts, text)
+		}
+	}
+	if len(ts) != 1 {
+		t.Fatalf("got %d TEXT declarations, want 1", len(ts))
+	}
+	if len(ts[0].Body) != n+1 {
+		t.Fatalf("got %d body statements, want %d", len(ts[0].Body), n+1)
+	}
+	if instr, ok := ts[0].Body[0].(*ast.Instr); !ok || instr.Mnemonic.Text != "MOVQ" {
+		t.Fatalf("first body statement is not the expansion: %#v", ts[0].Body[0])
+	}
+	if elapsed := time.Since(start); elapsed > 10*time.Second {
+		t.Fatalf("expansion of %d invocations took %v, want the linear-cost contract", n, elapsed)
+	}
+}
+
+// TestAmplifyingMacrosStopAtBudget pins the work budget of expansion: a body
+// that repeats its argument multiplies every level of nesting, so a short
+// line of nested invocations names exponentially many tokens.  The expansion
+// must stop at the budget with a diagnostic instead of committing the
+// machine to producing them.
+func TestAmplifyingMacrosStopAtBudget(t *testing.T) {
+	src := "#define A(x) x x x x x x x x x x\n" +
+		"A(A(A(A(A(A(A(A(A(A(1))))))))))\n"
+	f, errs := ParseWithOptions("t_amd64.s", src, Options{Expand: true})
+	if len(errs) == 0 {
+		t.Fatal("amplifying macro expansion reported no error")
+	}
+	if !strings.Contains(errs[0].Error(), "work budget") {
+		t.Fatalf("error is not the budget diagnostic: %v", errs[0])
+	}
+	for _, d := range f.Decls {
+		if _, ok := d.(*ast.Text); ok {
+			t.Fatal("amplified line produced a TEXT declaration")
+		}
+	}
+}
+
 func TestObjectMacroExpandsAtUse(t *testing.T) {
 	_, got := expand(t, `
 #define REGTMP CX