From 000055dff02aa611a30633c1de36ddd027cd41e9 Mon Sep 17 00:00:00 2001
From: Ville Vesilehto <ville@vesilehto.fi>
Date: Sat, 10 Jan 2026 08:58:18 +0200
Subject: [PATCH] perf(ast): reduce allocations in With.MarshalJSON

Replace map[string]any with a struct for JSON serialization.

Signed-off-by: Ville Vesilehto <ville@vesilehto.fi>
---
 v1/ast/policy.go            | 18 ++++++++++++------
 v1/ast/policy_bench_test.go | 13 +++++++++++++
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/v1/ast/policy.go b/v1/ast/policy.go
index 8aa12a37bb..4ce31953f3 100644
--- a/v1/ast/policy.go
+++ b/v1/ast/policy.go
@@ -1729,16 +1729,22 @@ func (w *With) SetLoc(loc *Location) {
 	w.Location = loc
 }
 
+// withJSON is used for JSON serialization of With to avoid map allocation overhead.
+// Field order is alphabetical to match previous map-based output.
+type withJSON struct {
+	Location *Location `json:"location,omitempty"`
+	Target   *Term     `json:"target"`
+	Value    *Term     `json:"value"`
+}
+
 func (w *With) MarshalJSON() ([]byte, error) {
-	data := map[string]any{
-		"target": w.Target,
-		"value":  w.Value,
+	data := withJSON{
+		Target: w.Target,
+		Value:  w.Value,
 	}
 
 	if astJSON.GetOptions().MarshalOptions.IncludeLocation.With {
-		if w.Location != nil {
-			data["location"] = w.Location
-		}
+		data.Location = w.Location
 	}
 
 	return json.Marshal(data)
diff --git a/v1/ast/policy_bench_test.go b/v1/ast/policy_bench_test.go
index fc2a5d4095..288619a8e1 100644
--- a/v1/ast/policy_bench_test.go
+++ b/v1/ast/policy_bench_test.go
@@ -157,3 +157,16 @@ func BenchmarkRuleMarshalJSON(b *testing.B) {
 		})
 	}
 }
+
+func BenchmarkWithMarshalJSON(b *testing.B) {
+	module := MustParseModule(`
+		package test
+		allow if { input.x with input as {"x": true} }
+	`)
+
+	with := module.Rules[0].Body[0].With[0]
+
+	for b.Loop() {
+		_, _ = json.Marshal(with)
+	}
+}