Remove jsonOptions from AST nodes and terms

This was originally added to allow Regal to get a serialized AST that included
location data wherever that was possible. Regal is however no longer using OPA's
JSON serialization but its own custom encoder. Attaching these options to every
AST node and term comes with a cost attached, and without any known users of this
feature vs. the many users who care about resource utilization, this feels like
an easy choice.

While it seems unlikely to be users depending on this functionality — in case
someone needs it, the options for serializing AST nodes to JSON can now be set
globally instead. Global state is always awkward, but since JSON marshalling
methods only have access to the node being marshalled and of course, global
state, there's not a whole lot of options if we intend to keep this feature.

Signed-off-by: Anders Eknert <[email protected]>

cmd/parse.go

@@ -90,7 +90,7 @@ func parse(args []string, params *parseParams, stdout io.Writer, stderr io.Write
 		RegoVersion:       params.regoVersion(),
 	}
 	if exposeLocation {
-		parserOpts.JSONOptions = &astJSON.Options{
+		astJSON.SetOptions(astJSON.Options{
 			MarshalOptions: astJSON.MarshalOptions{
 				IncludeLocationText: true,
 				IncludeLocation: astJSON.NodeToggle{
@@ -108,7 +108,8 @@ func parse(args []string, params *parseParams, stdout io.Writer, stderr io.Write
 					AnnotationsRef: true,
 				},
 			},
-		}
+		})
+		defer astJSON.SetOptions(astJSON.Defaults())
 	}
 
 	result, err := loader.RegoWithOpts(args[0], parserOpts)

internal/bundle/inspect/inspect.go

@@ -43,18 +43,20 @@ func FileForRegoVersion(regoVersion ast.RegoVersion, path string, includeAnnotat
 }
 
 func bundleOrDirInfoForRegoVersion(regoVersion ast.RegoVersion, path string, includeAnnotations bool) (*Info, error) {
+	json.SetOptions(json.Options{
+		MarshalOptions: json.MarshalOptions{
+			IncludeLocation: json.NodeToggle{
+				// Annotation location data is only included if includeAnnotations is set
+				AnnotationsRef: includeAnnotations,
+			},
+		},
+	})
+	defer json.SetOptions(json.Defaults())
+
 	b, err := loader.NewFileLoader().
 		WithRegoVersion(regoVersion).
 		WithSkipBundleVerification(true).
 		WithProcessAnnotation(true). // Always process annotations, for enriching namespace listing
-		WithJSONOptions(&json.Options{
-			MarshalOptions: json.MarshalOptions{
-				IncludeLocation: json.NodeToggle{
-					// Annotation location data is only included if includeAnnotations is set
-					AnnotationsRef: includeAnnotations,
-				},
-			},
-		}).
 		AsBundle(path)
 	if err != nil {
 		return nil, err
@@ -209,18 +211,20 @@ func (bi *Info) getBundleDataWasmAndSignatures(name string) error {
 }
 
 func fileInfoForRegoVersion(regoVersion ast.RegoVersion, path string, includeAnnotations bool) (*Info, error) {
+	json.SetOptions(json.Options{
+		MarshalOptions: json.MarshalOptions{
+			IncludeLocation: json.NodeToggle{
+				// Annotation location data is only included if includeAnnotations is set
+				AnnotationsRef: includeAnnotations,
+			},
+		},
+	})
+	defer json.SetOptions(json.Defaults())
+
 	res, err := loader.NewFileLoader().
 		WithRegoVersion(regoVersion).
 		WithSkipBundleVerification(true).
 		WithProcessAnnotation(true). // Always process annotations, for enriching namespace listing
-		WithJSONOptions(&json.Options{
-			MarshalOptions: json.MarshalOptions{
-				IncludeLocation: json.NodeToggle{
-					// Annotation location data is only included if includeAnnotations is set
-					AnnotationsRef: includeAnnotations,
-				},
-			},
-		}).
 		All([]string{path})
 	if err != nil {
 		return nil, err

internal/bundle/inspect/inspect_test.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/open-policy-agent/opa/internal/file/archive"
 	"github.com/open-policy-agent/opa/v1/ast"
+	astJSON "github.com/open-policy-agent/opa/v1/ast/json"
 	"github.com/open-policy-agent/opa/v1/bundle"
 	"github.com/open-policy-agent/opa/v1/util/test"
 )
@@ -94,6 +95,16 @@ p = 1`,
 			t.Fatalf("expected %d annotation, but got: %d", exp, got)
 		}
 
+		astJSON.SetOptions(astJSON.Options{
+			MarshalOptions: astJSON.MarshalOptions{
+				IncludeLocation: astJSON.NodeToggle{
+					// Annotation location data is only included if includeAnnotations is set
+					AnnotationsRef: true,
+				},
+			},
+		})
+		defer astJSON.SetOptions(astJSON.Defaults())
+
 		bs, err := json.Marshal(info.Annotations[0])
 		if err != nil {
 			t.Fatalf("Unexpected error: %v", err)

v1/ast/annotations.go

@@ -38,9 +38,8 @@ type (
 		Custom           map[string]interface{}       `json:"custom,omitempty"`
 		Location         *Location                    `json:"location,omitempty"`
 
-		comments    []*Comment
-		node        Node
-		jsonOptions astJSON.Options
+		comments []*Comment
+		node     Node
 	}
 
 	// SchemaAnnotation contains a schema declaration for the document identified by the path.
@@ -77,8 +76,6 @@ type (
 		Annotations *Annotations `json:"annotations,omitempty"`
 		Location    *Location    `json:"location,omitempty"` // The location of the node the annotations are applied to
 
-		jsonOptions astJSON.Options
-
 		node Node // The node the annotations are applied to
 	}
 
@@ -181,13 +178,6 @@ func (a *Annotations) GetTargetPath() Ref {
 	}
 }
 
-func (a *Annotations) setJSONOptions(opts astJSON.Options) {
-	a.jsonOptions = opts
-	if a.Location != nil {
-		a.Location.JSONOptions = opts
-	}
-}
-
 func (a *Annotations) MarshalJSON() ([]byte, error) {
 	if a == nil {
 		return []byte(`{"scope":""}`), nil
@@ -229,7 +219,7 @@ func (a *Annotations) MarshalJSON() ([]byte, error) {
 		data["custom"] = a.Custom
 	}
 
-	if a.jsonOptions.MarshalOptions.IncludeLocation.Annotations {
+	if astJSON.GetOptions().MarshalOptions.IncludeLocation.Annotations {
 		if a.Location != nil {
 			data["location"] = a.Location
 		}
@@ -249,7 +239,6 @@ func NewAnnotationsRef(a *Annotations) *AnnotationsRef {
 		Path:        a.GetTargetPath(),
 		Annotations: a,
 		node:        a.node,
-		jsonOptions: a.jsonOptions,
 	}
 }
 
@@ -282,10 +271,20 @@ func (ar *AnnotationsRef) MarshalJSON() ([]byte, error) {
 		data["annotations"] = ar.Annotations
 	}
 
-	if ar.jsonOptions.MarshalOptions.IncludeLocation.AnnotationsRef {
+	if astJSON.GetOptions().MarshalOptions.IncludeLocation.AnnotationsRef {
 		if ar.Location != nil {
 			data["location"] = ar.Location
 		}
+
+		// The location set for the schema ref terms is wrong (always set to
+		// row 1) and not really useful anyway.. so strip it out bedfore marshalling
+		for _, schema := range ar.Annotations.Schemas {
+			if schema.Path != nil {
+				for _, term := range schema.Path {
+					term.Location = nil
+				}
+			}
+		}
 	}
 
 	return json.Marshal(data)

v1/ast/json/json.go

@@ -1,5 +1,26 @@
+// Copyright 2023 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+// This package provides options for JSON marshalling of AST nodes, and location
+// data in particular. Since location data occupies a significant portion of the
+// AST when included, it is excluded by default. The options provided here allow
+// changing that behavior — either for all nodes or for specific types. Since
+// JSONMarshaller implementations have access only to the node being marshaled,
+// our options are to either attach these settings to *all* nodes in the AST, or
+// to provide them via global state. The former is perhaps a little more elegant,
+// and is what we went with initially. The cost of attaching these settings to
+// every node however turned out to be non-negligible, and given that the number
+// of users who have an interest in AST serialization are likely to be few, we
+// have since switched to using global state, as provided here. Note that this
+// is mostly to provide an equivalent feature to what we had before, should
+// anyone depend on that. Users who need fine-grained control over AST
+// serialization are recommended to use external libraries for that purpose,
+// such as `github.com/json-iterator/go`.
 package json
 
+import "sync"
+
 // Options defines the options for JSON operations,
 // currently only marshaling can be configured
 type Options struct {
@@ -34,3 +55,52 @@ type NodeToggle struct {
 	Annotations    bool
 	AnnotationsRef bool
 }
+
+// configuredJSONOptions synchronizes access to the global JSON options
+type configuredJSONOptions struct {
+	options Options
+	lock    sync.RWMutex
+}
+
+var options = &configuredJSONOptions{
+	options: Defaults(),
+}
+
+// SetOptions sets the global options for marshalling AST nodes to JSON
+func SetOptions(opts Options) {
+	options.lock.Lock()
+	defer options.lock.Unlock()
+	options.options = opts
+}
+
+// GetOptions returns (a copy of) the global options for marshalling AST nodes to JSON
+func GetOptions() Options {
+	options.lock.RLock()
+	defer options.lock.RUnlock()
+	return options.options
+}
+
+// Defaults returns the default JSON options, which is to exclude location
+// information in serialized JSON AST.
+func Defaults() Options {
+	return Options{
+		MarshalOptions: MarshalOptions{
+			IncludeLocation: NodeToggle{
+				Term:           false,
+				Package:        false,
+				Comment:        false,
+				Import:         false,
+				Rule:           false,
+				Head:           false,
+				Expr:           false,
+				SomeDecl:       false,
+				Every:          false,
+				With:           false,
+				Annotations:    false,
+				AnnotationsRef: false,
+			},
+			IncludeLocationText: false,
+			ExcludeLocationFile: false,
+		},
+	}
+}

v1/ast/location/location.go

@@ -18,9 +18,6 @@ type Location struct {
 	Col    int    `json:"col"`  // The column in the row.
 	Offset int    `json:"-"`    // The byte offset for the location in the source.
 
-	// JSONOptions specifies options for marshaling and unmarshalling of locations
-	JSONOptions astJSON.Options
-
 	Tabs []int `json:"-"` // The column offsets of tabs in the source.
 }
 
@@ -98,7 +95,8 @@ func (loc *Location) Compare(other *Location) int {
 
 func (loc *Location) MarshalJSON() ([]byte, error) {
 	// structs are used here to preserve the field ordering of the original Location struct
-	if loc.JSONOptions.MarshalOptions.ExcludeLocationFile {
+	jsonOptions := astJSON.GetOptions().MarshalOptions
+	if jsonOptions.ExcludeLocationFile {
 		data := struct {
 			Row  int    `json:"row"`
 			Col  int    `json:"col"`
@@ -108,7 +106,7 @@ func (loc *Location) MarshalJSON() ([]byte, error) {
 			Col: loc.Col,
 		}
 
-		if loc.JSONOptions.MarshalOptions.IncludeLocationText {
+		if jsonOptions.IncludeLocationText {
 			data.Text = loc.Text
 		}
 
@@ -126,7 +124,7 @@ func (loc *Location) MarshalJSON() ([]byte, error) {
 		File: loc.File,
 	}
 
-	if loc.JSONOptions.MarshalOptions.IncludeLocationText {
+	if jsonOptions.IncludeLocationText {
 		data.Text = loc.Text
 	}
 

v1/ast/location/location_test.go

@@ -90,8 +90,9 @@ func TestLocationCompare(t *testing.T) {
 
 func TestLocationMarshal(t *testing.T) {
 	testCases := map[string]struct {
-		loc *Location
-		exp string
+		loc     *Location
+		options astJSON.Options
+		exp     string
 	}{
 		"default json options": {
 			loc: &Location{
@@ -108,10 +109,10 @@ func TestLocationMarshal(t *testing.T) {
 				File: "file",
 				Row:  1,
 				Col:  1,
-				JSONOptions: astJSON.Options{
-					MarshalOptions: astJSON.MarshalOptions{
-						IncludeLocationText: true,
-					},
+			},
+			options: astJSON.Options{
+				MarshalOptions: astJSON.MarshalOptions{
+					IncludeLocationText: true,
 				},
 			},
 			exp: `{"file":"file","row":1,"col":1,"text":"dGV4dA=="}`,
@@ -121,10 +122,10 @@ func TestLocationMarshal(t *testing.T) {
 				File: "file",
 				Row:  1,
 				Col:  1,
-				JSONOptions: astJSON.Options{
-					MarshalOptions: astJSON.MarshalOptions{
-						ExcludeLocationFile: true,
-					},
+			},
+			options: astJSON.Options{
+				MarshalOptions: astJSON.MarshalOptions{
+					ExcludeLocationFile: true,
 				},
 			},
 			exp: `{"row":1,"col":1}`,
@@ -133,6 +134,9 @@ func TestLocationMarshal(t *testing.T) {
 
 	for id, tc := range testCases {
 		t.Run(id, func(t *testing.T) {
+			astJSON.SetOptions(tc.options)
+			defer astJSON.SetOptions(astJSON.Defaults())
+
 			bs, err := json.Marshal(tc.loc)
 			if err != nil {
 				t.Fatal(err)