Env copy on write (#1405)

* Program plan optimizations

| Benchmark Case | Before (ns/op) | After (ns/op) | Δ Time | Before (B/op) | After (B/op) | Δ Memory | Before (allocs) | After (allocs) | Δ Allocs |
| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
| BenchmarkProgramPlan/Default | 8,153 | 930 | **-88.6%** | 8,320 | 1,416 | **-83.0%** | 36 | 26 | **-27.8%** |
| BenchmarkProgramPlan/OptimizeUnneeded | 7,344 | 1,150 | **-84.3%** | 8,784 | 1,512 | **-82.8%** | 50 | 32 | **-36.0%** |
| BenchmarkProgramPlan/OptimizeNeeded | 8,370 | 2,164 | **-74.1%** | 10,224 | 2,976 | **-70.9%** | 67 | 52 | **-22.4%** |

* Minor refactor of the initialization logic to reduce program size

* Copy-on-write semantics for types.Registry and cel.Env internals

* Ensure shared declarations aren't copied unless necessary within the checker

* Capture NewEnv setup benchmarks as well

* Fix race-related issue with copy-on-write mutability check

* Eliminate dead-code from former Copy() approach. More tests

* Bug fix to support disabling declarations when using inherited declarations

Author: TristonianJones

cel/cel_test.go

@@ -95,6 +95,90 @@ func Test_ExampleWithBuiltins(t *testing.T) {
 	}
 }
 
+func TestExtendCheckerParity(t *testing.T) {
+	// Base environment carrying standard library functions
+	baseEnv, err := NewEnv(
+		Variable("baseVar", StringType),
+	)
+	if err != nil {
+		t.Fatalf("NewEnv() failed: %v", err)
+	}
+
+	// Extended environment adding child variables (K8s CRD pattern)
+	extEnv, err := baseEnv.Extend(
+		Variable("value", StringType),
+		Variable("oldValue", StringType),
+	)
+	if err != nil {
+		t.Fatalf("baseEnv.Extend() failed: %v", err)
+	}
+
+	// Equivalent flat environment created from scratch
+	flatEnv, err := NewEnv(
+		Variable("baseVar", StringType),
+		Variable("value", StringType),
+		Variable("oldValue", StringType),
+	)
+	if err != nil {
+		t.Fatalf("flat NewEnv() failed: %v", err)
+	}
+
+	testCases := []struct {
+		expr string
+		vars map[string]any
+		want ref.Val
+	}{
+		{
+			expr: `value + " " + oldValue + " " + baseVar`,
+			vars: map[string]any{"value": "new", "oldValue": "old", "baseVar": "base"},
+			want: types.String("new old base"),
+		},
+		{
+			expr: `size(value) > 0 && [1, 2, 3].exists(x, x > 2)`,
+			vars: map[string]any{"value": "test"},
+			want: types.True,
+		},
+	}
+
+	for _, tc := range testCases {
+		extAst, extIss := extEnv.Compile(tc.expr)
+		if extIss.Err() != nil {
+			t.Fatalf("extEnv.Compile(%q) failed: %v", tc.expr, extIss.Err())
+		}
+		flatAst, flatIss := flatEnv.Compile(tc.expr)
+		if flatIss.Err() != nil {
+			t.Fatalf("flatEnv.Compile(%q) failed: %v", tc.expr, flatIss.Err())
+		}
+
+		if extAst.OutputType().TypeName() != flatAst.OutputType().TypeName() {
+			t.Errorf("OutputType mismatch for %q: ext %v, flat %v", tc.expr, extAst.OutputType(), flatAst.OutputType())
+		}
+
+		extPrg, err := extEnv.Program(extAst)
+		if err != nil {
+			t.Fatalf("extEnv.Program() failed: %v", err)
+		}
+		flatPrg, err := flatEnv.Program(flatAst)
+		if err != nil {
+			t.Fatalf("flatEnv.Program() failed: %v", err)
+		}
+
+		extOut, _, err := extPrg.Eval(tc.vars)
+		if err != nil {
+			t.Fatalf("extPrg.Eval() failed: %v", err)
+		}
+		flatOut, _, err := flatPrg.Eval(tc.vars)
+		if err != nil {
+			t.Fatalf("flatPrg.Eval() failed: %v", err)
+		}
+
+		if extOut.Equal(tc.want) != types.True || flatOut.Equal(tc.want) != types.True {
+			t.Errorf("Eval result mismatch for %q: ext %v, flat %v, want %v", tc.expr, extOut, flatOut, tc.want)
+		}
+	}
+}
+
+
 func TestCompile(t *testing.T) {
 	prg, err := Compile(`"hello " + name`, Variable("name", StringType))
 	if err != nil {
@@ -4003,12 +4087,45 @@ func BenchmarkDynamicDispatch(b *testing.B) {
 }
 
 func BenchmarkProgramPlan(b *testing.B) {
-	env, err := NewEnv(
+	b.Run("NewEnv", func(b *testing.B) {
+		b.ReportAllocs()
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_, err := NewEnv(
+				Variable("ai", IntType),
+				Variable("ar", MapType(StringType, StringType)),
+			)
+			if err != nil {
+				b.Fatalf("NewEnv() failed: %v", err)
+			}
+		}
+	})
+
+	baseEnv, err := NewEnv()
+	if err != nil {
+		b.Fatalf("NewEnv() failed: %v", err)
+	}
+
+	b.Run("ExtendEnv", func(b *testing.B) {
+		b.ReportAllocs()
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_, err := baseEnv.Extend(
+				Variable("ai", IntType),
+				Variable("ar", MapType(StringType, StringType)),
+			)
+			if err != nil {
+				b.Fatalf("baseEnv.Extend() failed: %v", err)
+			}
+		}
+	})
+
+	env, err := baseEnv.Extend(
 		Variable("ai", IntType),
 		Variable("ar", MapType(StringType, StringType)),
 	)
 	if err != nil {
-		b.Fatalf("NewEnv() failed: %v", err)
+		b.Fatalf("Extend() failed: %v", err)
 	}
 	astSimple, iss := env.Compile("ai == 20 || ar['foo'] == 'bar'")
 	if iss.Err() != nil {

cel/decls.go

@@ -223,8 +223,8 @@ func FunctionDecls(funcs ...*decls.FunctionDecl) EnvOption {
 		if len(funcs) == 0 {
 			return e, nil
 		}
-		e.ensureMutableFunctions()
 		var err error
+		e.ensureMutableFunctions()
 		for _, fn := range funcs {
 			if existing, found := e.functions[fn.Name()]; found {
 				fn, err = existing.Merge(fn)

cel/env.go

@@ -532,34 +532,17 @@ func (e *Env) CompileSource(src Source) (*Ast, *Issues) {
 // TypeProvider are immutable, or that their underlying implementations are based on the
 // ref.TypeRegistry which provides a Copy method which will be invoked by this method.
 func (e *Env) Extend(opts ...EnvOption) (*Env, error) {
-	chk, chkErr := e.getCheckerOrError()
-	if chkErr != nil {
+	if _, chkErr := e.getCheckerOrError(); chkErr != nil {
 		return nil, chkErr
 	}
 
-	prsrOptsCopy := make([]parser.Option, len(e.prsrOpts))
-	copy(prsrOptsCopy, e.prsrOpts)
-
-	// The type-checker is configured with Declarations. The declarations may either be provided
-	// as options which have not yet been validated, or may come from a previous checker instance
-	// whose types have already been validated.
-	chkOptsCopy := make([]checker.Option, len(e.chkOpts))
-	copy(chkOptsCopy, e.chkOpts)
-
-	// Copy the declarations if needed.
-	if chk != nil {
-		// If the type-checker has already been instantiated, then the e.declarations have been
-		// validated within the chk instance.
-		chkOptsCopy = append(chkOptsCopy, checker.ValidatedDeclarations(chk))
-	}
-	varsCopy := make([]*decls.VariableDecl, len(e.variables))
-	copy(varsCopy, e.variables)
-
-	// Copy macros and program options
-	macsCopy := make([]parser.Macro, len(e.macros))
-	progOptsCopy := make([]ProgramOption, len(e.progOpts))
-	copy(macsCopy, e.macros)
-	copy(progOptsCopy, e.progOpts)
+	prsrOptsCopy := slices.Clone(e.prsrOpts)
+	chkOptsCopy := slices.Clone(e.chkOpts)
+	varsCopy := slices.Clone(e.variables)
+	macsCopy := slices.Clone(e.macros)
+	progOptsCopy := slices.Clone(e.progOpts)
+	validatorsCopy := slices.Clone(e.validators)
+	costOptsCopy := slices.Clone(e.costOptions)
 
 	// Copy the adapter / provider if they appear to be mutable.
 	adapter := e.adapter
@@ -588,12 +571,6 @@ func (e *Env) Extend(opts ...EnvOption) (*Env, error) {
 		adapter = adapterReg.Copy()
 	}
 
-	validatorsCopy := make([]ASTValidator, len(e.validators))
-	copy(validatorsCopy, e.validators)
-
-	costOptsCopy := make([]checker.CostOption, len(e.costOptions))
-	copy(costOptsCopy, e.costOptions)
-
 	ext := &Env{
 		parent:          e,
 		Container:       e.Container,
@@ -687,25 +664,17 @@ func (e *Env) HasFunction(functionName string) bool {
 
 // Functions returns a shallow copy of the Functions, keyed by function name, that have been configured in the environment.
 func (e *Env) Functions() map[string]*decls.FunctionDecl {
-	shallowCopy := make(map[string]*decls.FunctionDecl, len(e.functions))
-	for nm, fn := range e.functions {
-		shallowCopy[nm] = fn
-	}
-	return shallowCopy
+	return maps.Clone(e.functions)
 }
 
 // Variables returns a shallow copy of the variables associated with the environment.
 func (e *Env) Variables() []*decls.VariableDecl {
-	shallowCopy := make([]*decls.VariableDecl, len(e.variables))
-	copy(shallowCopy, e.variables)
-	return shallowCopy
+	return slices.Clone(e.variables)
 }
 
 // Macros returns a shallow copy of macros associated with the environment.
 func (e *Env) Macros() []Macro {
-	shallowCopy := make([]Macro, len(e.macros))
-	copy(shallowCopy, e.macros)
-	return shallowCopy
+	return slices.Clone(e.macros)
 }
 
 // HasValidator returns whether a specific ASTValidator has been configured in the environment.
@@ -718,9 +687,9 @@ func (e *Env) HasValidator(name string) bool {
 	return false
 }
 
-// Validators returns the set of ASTValidators configured on the environment.
+// Validators returns a shallow copy of the set of ASTValidators configured on the environment.
 func (e *Env) Validators() []ASTValidator {
-	return e.validators[:]
+	return slices.Clone(e.validators)
 }
 
 // Parse parses the input expression value `txt` to a Ast and/or a set of Issues.
@@ -1007,6 +976,15 @@ func (e *Env) initChecker() (*checker.Env, error) {
 		chkOpts = append(chkOpts,
 			checker.JSONFieldNames(e.HasFeature(featureJSONFieldNames)))
 
+		if e.parent != nil && e.funcsShared {
+			parentChk, err := e.parent.initChecker()
+			if err != nil {
+				e.setCheckerOrError(nil, err)
+				return
+			}
+			chkOpts = append(chkOpts, checker.ValidatedDeclarations(parentChk))
+		}
+
 		ce, err := checker.NewEnv(e.Container, e.provider, chkOpts...)
 		if err != nil {
 			e.setCheckerOrError(nil, err)
@@ -1019,14 +997,16 @@ func (e *Env) initChecker() (*checker.Env, error) {
 			return
 		}
 		// Add the function declarations which are derived from the FunctionDecl instances.
-		for _, fn := range e.functions {
-			if fn.IsDeclarationDisabled() {
-				continue
-			}
-			err = ce.AddFunctions(fn)
-			if err != nil {
-				e.setCheckerOrError(nil, err)
-				return
+		if e.parent == nil || !e.funcsShared {
+			for _, fn := range e.functions {
+				if fn.IsDeclarationDisabled() {
+					continue
+				}
+				err = ce.AddFunctions(fn)
+				if err != nil {
+					e.setCheckerOrError(nil, err)
+					return
+				}
 			}
 		}
 		// Add function declarations here separately.