feat: add configurable render budget with per-op stats

README.md

@@ -358,3 +358,101 @@ This package absolutely 100% could not have been written without the help of Tho
 Not only did the book make understanding the process of writing lexers, parsers, and asts, but it also provided the basis for the syntax of Plush itself.
 
 If you have yet to read Thorsten's book, I can't recommend it enough. Please go and buy it!
+
+---
+
+## Render Budget
+
+Plush lets you attach a work-unit **budget** to any render to protect against runaway templates — deeply nested loops, recursive partials, or unexpectedly expensive helpers.
+
+A **nil budget = unlimited**, so all existing code is completely unaffected.
+
+### Quick start
+
+```go
+b := plush.NewBudget(10_000)
+ctx := plush.NewContext()
+ctx.Set("products", products)
+ctx.WithBudget(b)
+
+html, err := plush.Render(tmpl, ctx)
+if errors.Is(err, plush.ErrBudgetExceeded) {
+    log.Printf("budget exceeded: used=%d remaining=%d", b.Used(), b.Remaining())
+    return errorPage()
+}
+
+// One-liner convenience wrapper
+html, err = plush.RenderWithBudget(tmpl, 10_000, ctx)
+```
+
+### Default operation costs
+
+| Operation | Default cost |
+|---|---|
+| Loop iteration | 1 |
+| Helper / function call | 5 |
+| Filter call | 3 |
+| Partial / sub-render | 10 |
+| Condition check (`if`) | 1 |
+| Variable assignment | 0 |
+| Object traversal (per segment) | 1 |
+
+### Custom costs
+
+Pass a `BudgetCosts` struct to override any cost:
+
+```go
+costs := plush.ZeroCosts()          // start from all-zero
+costs.LoopIteration = 1
+costs.SubRender     = 25
+
+html, err = plush.RenderWithBudgetConfig(tmpl, 5_000, costs, ctx)
+```
+
+### Per-function costs
+
+Override the cost for individual functions registered in the context:
+
+```go
+costs := plush.DefaultBudgetCosts()
+costs.FunctionCosts = map[string]int64{
+    "expensiveQuery": 50, // charged 50 per call instead of the default 5
+    "cheapHelper":     1,
+}
+
+html, err = plush.RenderWithBudgetConfig(tmpl, 10_000, costs, ctx)
+```
+
+Functions not listed in `FunctionCosts` fall back to the `HelperCall` cost.
+
+### Stats report
+
+After rendering, call `b.Stats()` to see exactly where the budget was spent:
+
+```go
+b := plush.NewBudget(10_000)
+ctx.WithBudget(b)
+plush.Render(tmpl, ctx)
+
+s := b.Stats()
+fmt.Printf("total=%d  loops=%d  calls=%d  conditions=%d\n",
+    s.TotalUsed, s.LoopIterations, s.FunctionCalls, s.ConditionChecks)
+
+for name, units := range s.ByFunction {
+    fmt.Printf("  %s: %d units\n", name, units)
+}
+```
+
+`BudgetStats` fields:
+
+| Field | What it measures |
+|---|---|
+| `TotalUsed` | Sum of all units spent |
+| `LoopIterations` | Units from loop iterations |
+| `FunctionCalls` | Units from all function/helper calls |
+| `FilterCalls` | Units from filter calls |
+| `SubRenders` | Units from partial renders |
+| `ConditionChecks` | Units from `if`/`unless` evaluations |
+| `Assignments` | Units from variable assignments |
+| `ObjectTraversals` | Units from dot-notation traversal |
+| `ByFunction` | Per-function breakdown (map of name → units) |

budget.go

@@ -0,0 +1,217 @@
+package plush
+
+import (
+	"errors"
+	"sync"
+	"sync/atomic"
+)
+
+// ErrBudgetExceeded is returned when a render exhausts its budget.
+var ErrBudgetExceeded = errors.New("render budget exceeded")
+
+// BudgetStats is a snapshot of work units consumed per operation category.
+// Retrieve it after rendering via b.Stats().
+type BudgetStats struct {
+	// TotalUsed is the sum of all units spent (equals b.Used()).
+	TotalUsed int64
+	// LoopIterations is total units charged by loop iterations.
+	LoopIterations int64
+	// FunctionCalls is total units charged by all function/helper calls.
+	FunctionCalls int64
+	// FilterCalls is total units charged by filter calls.
+	FilterCalls int64
+	// SubRenders is total units charged by partial/snippet renders.
+	SubRenders int64
+	// ConditionChecks is total units charged by if/unless evaluations.
+	ConditionChecks int64
+	// Assignments is total units charged by variable assignments.
+	Assignments int64
+	// ObjectTraversals is total units charged by dot-notation traversal.
+	ObjectTraversals int64
+	// ByFunction breaks FunctionCalls down by name for calls made via
+	// SpendFunctionCall. Functions without a FunctionCosts override appear
+	// here using the generic HelperCall cost.
+	ByFunction map[string]int64
+}
+
+// Budget tracks render work units during template evaluation.
+// A nil Budget is always unlimited — zero breaking changes.
+type Budget struct {
+	limit   int64
+	counter atomic.Int64
+	costs   BudgetCosts
+
+	// per-category stat counters — all lock-free
+	statLoop      atomic.Int64
+	statFunction  atomic.Int64 // total of all function/helper calls
+	statFilter    atomic.Int64
+	statSubRender atomic.Int64
+	statCondition atomic.Int64
+	statAssign    atomic.Int64
+	statTraversal atomic.Int64
+
+	// per-function breakdown — mutex-protected plain map
+	statFuncsMu  sync.Mutex
+	statFuncsMap map[string]int64
+}
+
+// NewBudget creates a Budget with a limit and default costs.
+func NewBudget(limit int64) *Budget {
+	return &Budget{
+		limit:        limit,
+		costs:        DefaultBudgetCosts(),
+		statFuncsMap: make(map[string]int64),
+	}
+}
+
+// NewBudgetWithCosts creates a Budget with fully custom per-operation costs.
+func NewBudgetWithCosts(limit int64, costs BudgetCosts) *Budget {
+	return &Budget{
+		limit:        limit,
+		costs:        costs,
+		statFuncsMap: make(map[string]int64),
+	}
+}
+
+// WithCosts replaces the cost configuration. Returns self for chaining.
+func (b *Budget) WithCosts(costs BudgetCosts) *Budget {
+	b.costs = costs
+	return b
+}
+
+// Costs returns the active cost configuration.
+func (b *Budget) Costs() BudgetCosts {
+	return b.costs
+}
+
+// Used returns total units consumed so far.
+func (b *Budget) Used() int64 {
+	return b.counter.Load()
+}
+
+// Remaining returns units left before the limit is hit.
+func (b *Budget) Remaining() int64 {
+	r := b.limit - b.counter.Load()
+	if r < 0 {
+		return 0
+	}
+	return r
+}
+
+// Stats returns a snapshot of work units consumed per operation category.
+// Safe to call at any point during or after rendering.
+func (b *Budget) Stats() BudgetStats {
+	if b == nil {
+		return BudgetStats{}
+	}
+	s := BudgetStats{
+		TotalUsed:        b.counter.Load(),
+		LoopIterations:   b.statLoop.Load(),
+		FunctionCalls:    b.statFunction.Load(),
+		FilterCalls:      b.statFilter.Load(),
+		SubRenders:       b.statSubRender.Load(),
+		ConditionChecks:  b.statCondition.Load(),
+		Assignments:      b.statAssign.Load(),
+		ObjectTraversals: b.statTraversal.Load(),
+		ByFunction:       make(map[string]int64),
+	}
+	b.statFuncsMu.Lock()
+	for k, v := range b.statFuncsMap {
+		s.ByFunction[k] = v
+	}
+	b.statFuncsMu.Unlock()
+	return s
+}
+
+// SpendLoop spends the loop iteration cost.
+func (b *Budget) SpendLoop() error {
+	if b == nil {
+		return nil
+	}
+	b.statLoop.Add(b.costs.LoopIteration)
+	return b.spend(b.costs.LoopIteration)
+}
+
+// SpendHelperCall spends the helper call cost.
+func (b *Budget) SpendHelperCall() error {
+	if b == nil {
+		return nil
+	}
+	b.statFunction.Add(b.costs.HelperCall)
+	return b.spend(b.costs.HelperCall)
+}
+
+// SpendFilter spends the filter call cost.
+func (b *Budget) SpendFilter() error {
+	if b == nil {
+		return nil
+	}
+	b.statFilter.Add(b.costs.FilterCall)
+	return b.spend(b.costs.FilterCall)
+}
+
+// SpendSubRender spends the sub-render cost.
+func (b *Budget) SpendSubRender() error {
+	if b == nil {
+		return nil
+	}
+	b.statSubRender.Add(b.costs.SubRender)
+	return b.spend(b.costs.SubRender)
+}
+
+// SpendCondition spends the condition check cost.
+func (b *Budget) SpendCondition() error {
+	if b == nil {
+		return nil
+	}
+	b.statCondition.Add(b.costs.ConditionCheck)
+	return b.spend(b.costs.ConditionCheck)
+}
+
+// SpendAssignment spends the assignment cost.
+func (b *Budget) SpendAssignment() error {
+	if b == nil {
+		return nil
+	}
+	b.statAssign.Add(b.costs.Assignment)
+	return b.spend(b.costs.Assignment)
+}
+
+// SpendFunctionCall spends the cost for a named function call.
+// Uses FunctionCosts[name] if set, otherwise falls back to HelperCall cost.
+func (b *Budget) SpendFunctionCall(name string) error {
+	if b == nil {
+		return nil
+	}
+	cost := b.costs.HelperCall
+	if c, ok := b.costs.FunctionCosts[name]; ok {
+		cost = c
+	}
+	b.statFunction.Add(cost)
+	b.statFuncsMu.Lock()
+	b.statFuncsMap[name] += cost
+	b.statFuncsMu.Unlock()
+	return b.spend(cost)
+}
+
+// SpendObjectTraversal spends ObjectTraversal * segments units.
+// e.g. product.variants.first = 3 segments → costs ObjectTraversal * 3
+func (b *Budget) SpendObjectTraversal(segments int) error {
+	if b == nil {
+		return nil
+	}
+	units := b.costs.ObjectTraversal * int64(segments)
+	b.statTraversal.Add(units)
+	return b.spend(units)
+}
+
+// spend is the internal hot path. Uses atomic add with no locks.
+func (b *Budget) spend(units int64) error {
+	if b == nil || units == 0 {
+		return nil
+	}
+	if b.counter.Add(units) > b.limit {
+		return ErrBudgetExceeded
+	}
+	return nil
+}

budget_config.go

@@ -0,0 +1,58 @@
+package plush
+
+// BudgetCosts defines the work-unit cost for each operation type.
+type BudgetCosts struct {
+	// LoopIteration is spent once per for-loop iteration.
+	// Default: 1
+	LoopIteration int64
+
+	// HelperCall is spent each time a registered helper is invoked.
+	// Default: 5
+	HelperCall int64
+
+	// FilterCall is spent per filter applied (sort, map, where).
+	// Default: 3
+	FilterCall int64
+
+	// SubRender is spent each time a partial/snippet is rendered.
+	// Default: 10
+	SubRender int64
+
+	// ConditionCheck is spent per if/unless/case evaluation.
+	// Default: 1
+	ConditionCheck int64
+
+	// Assignment is spent per variable assignment.
+	// Default: 0 (free — rarely the bottleneck)
+	Assignment int64
+
+	// ObjectTraversal is spent per dot-notation segment accessed.
+	// e.g. product.variants.first = 3 segments = 3 units
+	// Default: 1
+	ObjectTraversal int64
+
+	// FunctionCosts overrides the default HelperCall cost for specific named
+	// functions. The key is the function name as registered in the context.
+	// If a name is present here, its cost is used instead of HelperCall.
+	// e.g. costs.FunctionCosts = map[string]int64{"expensiveQuery": 50}
+	FunctionCosts map[string]int64
+}
+
+// DefaultBudgetCosts returns recommended production defaults.
+func DefaultBudgetCosts() BudgetCosts {
+	return BudgetCosts{
+		LoopIteration:   1,
+		HelperCall:      5,
+		FilterCall:      3,
+		SubRender:       10,
+		ConditionCheck:  1,
+		Assignment:      0,
+		ObjectTraversal: 1,
+	}
+}
+
+// ZeroCosts returns all-zero costs.
+// Useful for isolating one operation type in tests.
+func ZeroCosts() BudgetCosts {
+	return BudgetCosts{}
+}