eval: performance improvements (#392)

* eval: performance improvements

These changes improve evaluation performance by memoizing exported
values and merged object keys and avoiding copies when building object
schemas.

This gives a tremendous improvement in execution time for
evaluation-dominated scenarios, but results in little to no change for
open- or load-dominated scenarios.

Local benchmark results:

goos: darwin
goarch: arm64
pkg: github.com/pulumi/esc/eval
cpu: Apple M1 Max
BenchmarkEval-10           	    453	  2566938 ns/op	2204129 B/op	  18883 allocs/op
BenchmarkEval-10           	    468	  2586703 ns/op	2204146 B/op	  18884 allocs/op
BenchmarkEval-10           	    462	  2588916 ns/op	2204202 B/op	  18883 allocs/op
BenchmarkEval-10           	    464	  2586365 ns/op	2204228 B/op	  18884 allocs/op
BenchmarkEval-10           	    463	  2577383 ns/op	2204362 B/op	  18884 allocs/op
BenchmarkEval-10           	    477	  2537640 ns/op	2204389 B/op	  18884 allocs/op
BenchmarkEval-10           	    468	  2582930 ns/op	2204594 B/op	  18884 allocs/op
BenchmarkEval-10           	    463	  2582915 ns/op	2204246 B/op	  18884 allocs/op
BenchmarkEval-10           	    469	  2608014 ns/op	2204382 B/op	  18884 allocs/op
BenchmarkEval-10           	    465	  2554270 ns/op	2204313 B/op	  18884 allocs/op
BenchmarkEvalOpen-10       	      9	119163125 ns/op	2208651 B/op	  18926 allocs/op
BenchmarkEvalOpen-10       	      9	118168319 ns/op	2209928 B/op	  18928 allocs/op
BenchmarkEvalOpen-10       	      9	118805454 ns/op	2208294 B/op	  18924 allocs/op
BenchmarkEvalOpen-10       	      9	118506347 ns/op	2208712 B/op	  18922 allocs/op
BenchmarkEvalOpen-10       	      9	118898060 ns/op	2210256 B/op	  18926 allocs/op
BenchmarkEvalOpen-10       	      9	118450250 ns/op	2208210 B/op	  18924 allocs/op
BenchmarkEvalOpen-10       	      9	117723833 ns/op	2207528 B/op	  18922 allocs/op
BenchmarkEvalOpen-10       	      9	117134787 ns/op	2209227 B/op	  18925 allocs/op
BenchmarkEvalOpen-10       	      9	116210269 ns/op	2208843 B/op	  18926 allocs/op
BenchmarkEvalOpen-10       	      9	116987444 ns/op	2208736 B/op	  18925 allocs/op
BenchmarkEvalEnvLoad-10    	      4	298021334 ns/op	2216058 B/op	  18951 allocs/op
BenchmarkEvalEnvLoad-10    	      4	302557979 ns/op	2213974 B/op	  18944 allocs/op
BenchmarkEvalEnvLoad-10    	      4	293050229 ns/op	2212098 B/op	  18945 allocs/op
BenchmarkEvalEnvLoad-10    	      4	304410510 ns/op	2211322 B/op	  18946 allocs/op
BenchmarkEvalEnvLoad-10    	      4	301698562 ns/op	2212554 B/op	  18947 allocs/op
BenchmarkEvalEnvLoad-10    	      4	299588854 ns/op	2214102 B/op	  18946 allocs/op
BenchmarkEvalEnvLoad-10    	      4	295087740 ns/op	2211650 B/op	  18944 allocs/op
BenchmarkEvalEnvLoad-10    	      4	295875531 ns/op	2212638 B/op	  18950 allocs/op
BenchmarkEvalEnvLoad-10    	      4	294871781 ns/op	2212038 B/op	  18945 allocs/op
BenchmarkEvalEnvLoad-10    	      4	294592875 ns/op	2211682 B/op	  18945 allocs/op
BenchmarkEvalAll-10        	      3	405058722 ns/op	2215330 B/op	  18976 allocs/op
BenchmarkEvalAll-10        	      3	407002764 ns/op	2215688 B/op	  18978 allocs/op
BenchmarkEvalAll-10        	      3	409757153 ns/op	2214973 B/op	  18976 allocs/op
BenchmarkEvalAll-10        	      3	404553611 ns/op	2215261 B/op	  18977 allocs/op
BenchmarkEvalAll-10        	      3	402620945 ns/op	2216994 B/op	  18980 allocs/op
BenchmarkEvalAll-10        	      3	405302139 ns/op	2213112 B/op	  18973 allocs/op
BenchmarkEvalAll-10        	      3	404533556 ns/op	2215848 B/op	  18978 allocs/op
BenchmarkEvalAll-10        	      3	403431236 ns/op	2215896 B/op	  18979 allocs/op
BenchmarkEvalAll-10        	      3	402586597 ns/op	2217245 B/op	  18983 allocs/op
BenchmarkEvalAll-10        	      3	404775236 ns/op	2217122 B/op	  18980 allocs/op

* CL

CHANGELOG_PENDING.md

@@ -1,5 +1,11 @@
 ### Improvements
 
+- Improve evaluation performance and memory footprint.
+  [#392](https://github.com/pulumi/esc/pull/392)
+
 ### Bug Fixes
 
 ### Breaking changes
+
+- `schema`: `ObjectBuilder.Properties` and `Record` now take a `MapBuilder` in order to avoid copies.
+  [#392](https://github.com/pulumi/esc/pull/392)

analysis/common_test.go

@@ -27,18 +27,18 @@ import (
 )
 
 var testProviderSchema = schema.Object().
-	Properties(map[string]schema.Builder{
+	Properties(schema.BuilderMap{
 		"address": schema.String().
 			Description("The URL of the Vault server. Must contain a scheme and hostname, but no path."),
 		"jwt": schema.Object().
-			Properties(map[string]schema.Builder{
+			Properties(schema.BuilderMap{
 				"mount": schema.String().Description("The name of the authentication engine mount."),
 				"role":  schema.String().Description("The name of the role to use for login."),
 			}).
 			Required("role").
 			Description("Options for JWT login. JWT login uses an OIDC token issued by the Pulumi Cloud to generate an ephemeral token."),
 		"token": schema.Object().
-			Properties(map[string]schema.Builder{
+			Properties(schema.BuilderMap{
 				"displayName": schema.String().Description("The display name of the ephemeral token. Defaults to 'pulumi'."),
 				"token":       schema.String().Description("The parent token."),
 				"maxTtl": schema.String().

ast/environment.go

@@ -127,7 +127,7 @@ func (d *MapDecl[T]) parse(name string, node syntax.Node) syntax.Diagnostics {
 		kvp := obj.Index(i)
 
 		var v T
-		vname := fmt.Sprintf("%s.%s", name, kvp.Key.Value())
+		vname := name + "." + kvp.Key.Value()
 		vdiags := parseNode(vname, &v, kvp.Value)
 		diags.Extend(vdiags...)
 

ast/expr.go

@@ -230,7 +230,7 @@ type SymbolExpr struct {
 func Symbol(accessors ...PropertyAccessor) *SymbolExpr {
 	property := &PropertyAccess{Accessors: accessors}
 	return &SymbolExpr{
-		exprNode: expr(syntax.String(fmt.Sprintf("${%v}", property))),
+		exprNode: expr(syntax.String("${" + property.String() + "}")),
 		Property: property,
 	}
 }

cmd/esc/cli/client/client_test.go

@@ -478,7 +478,7 @@ func TestCheckYAMLEnvironment(t *testing.T) {
 		expected := &esc.Environment{
 			Exprs:      map[string]esc.Expr{"foo": {Literal: "bar"}},
 			Properties: map[string]esc.Value{"foo": esc.NewValue("bar")},
-			Schema:     schema.Record(map[string]schema.Builder{"foo": schema.String().Const("bar")}).Schema(),
+			Schema:     schema.Record(schema.BuilderMap{"foo": schema.String().Const("bar")}).Schema(),
 		}
 
 		client := newTestClient(t, http.MethodPost, "/api/esc/environments/test-org/yaml/check", func(w http.ResponseWriter, r *http.Request) {
@@ -604,7 +604,7 @@ func TestGetOpenEnvironment(t *testing.T) {
 		expected := &esc.Environment{
 			Exprs:      map[string]esc.Expr{"foo": {Literal: "bar"}},
 			Properties: map[string]esc.Value{"foo": esc.NewValue("bar")},
-			Schema:     schema.Record(map[string]schema.Builder{"foo": schema.String().Const("bar")}).Schema(),
+			Schema:     schema.Record(schema.BuilderMap{"foo": schema.String().Const("bar")}).Schema(),
 		}
 
 		client := newTestClient(t, http.MethodGet, "/api/esc/environments/test-org/test-project/test-env/open/session", func(w http.ResponseWriter, r *http.Request) {
@@ -638,7 +638,7 @@ func TestGetAnonymousOpenEnvironment(t *testing.T) {
 		expected := &esc.Environment{
 			Exprs:      map[string]esc.Expr{"foo": {Literal: "bar"}},
 			Properties: map[string]esc.Value{"foo": esc.NewValue("bar")},
-			Schema:     schema.Record(map[string]schema.Builder{"foo": schema.String().Const("bar")}).Schema(),
+			Schema:     schema.Record(schema.BuilderMap{"foo": schema.String().Const("bar")}).Schema(),
 		}
 
 		client := newTestClient(t, http.MethodGet, "/api/esc/environments/test-org/yaml/open/session", func(w http.ResponseWriter, r *http.Request) {

eval/eval.go

@@ -360,7 +360,7 @@ func (e *evalContext) evaluate() (*value, syntax.Diagnostics) {
 	// root.
 	properties := make(map[string]*expr, len(e.env.Values.GetEntries()))
 	e.root = &expr{
-		path: fmt.Sprintf("<%v>", e.name),
+		path: "<" + e.name + ">",
 		repr: &objectExpr{
 			node:       ast.Object(),
 			properties: properties,
@@ -402,7 +402,7 @@ func (e *evalContext) evaluateImports() {
 		e.evaluateImport(myImports, entry)
 	}
 
-	properties := make(map[string]schema.Builder, len(myImports))
+	properties := make(schema.SchemaMap, len(myImports))
 	for k, v := range myImports {
 		properties[k] = v.schema
 	}
@@ -582,7 +582,7 @@ func (e *evalContext) evaluateObject(x *expr, repr *objectExpr) *value {
 	keys := maps.Keys(repr.properties)
 	sort.Strings(keys)
 
-	object, properties := make(map[string]*value, len(keys)), make(map[string]schema.Builder, len(keys))
+	object, properties := make(map[string]*value, len(keys)), make(schema.SchemaMap, len(keys))
 	for _, k := range keys {
 		pv := e.evaluateExpr(repr.properties[k])
 		object[k], properties[k] = pv, pv.schema
@@ -921,7 +921,7 @@ func (e *evalContext) evaluateBuiltinOpen(x *expr, repr *openExpr) *value {
 
 	output, err := provider.Open(e.ctx, inputs.export("").Value.(map[string]esc.Value), e.execContext)
 	if err != nil {
-		e.errorf(repr.syntax(), err.Error())
+		e.errorf(repr.syntax(), "%s", err.Error())
 		v.unknown = true
 		return v
 	}

eval/eval_test.go

@@ -43,7 +43,7 @@ func accept() bool {
 type errorProvider struct{}
 
 func (errorProvider) Schema() (*schema.Schema, *schema.Schema) {
-	return schema.Record(map[string]schema.Builder{"why": schema.String()}).Schema(), schema.Always()
+	return schema.Record(schema.BuilderMap{"why": schema.String()}).Schema(), schema.Always()
 }
 
 func (errorProvider) Open(ctx context.Context, inputs map[string]esc.Value, context esc.EnvExecContext) (esc.Value, error) {
@@ -54,12 +54,12 @@ type testSchemaProvider struct{}
 
 func (testSchemaProvider) Schema() (*schema.Schema, *schema.Schema) {
 	s := schema.Object().
-		Defs(map[string]schema.Builder{
-			"defRecord": schema.Record(map[string]schema.Builder{
+		Defs(schema.BuilderMap{
+			"defRecord": schema.Record(schema.BuilderMap{
 				"baz": schema.String().Const("qux"),
 			}),
 		}).
-		Properties(map[string]schema.Builder{
+		Properties(schema.BuilderMap{
 			"null":    schema.Null(),
 			"boolean": schema.Boolean(),
 			"false":   schema.Boolean().Const(false),
@@ -71,7 +71,7 @@ func (testSchemaProvider) Schema() (*schema.Schema, *schema.Schema) {
 			"array":   schema.Array().Items(schema.Always()),
 			"tuple":   schema.Tuple(schema.String().Const("hello"), schema.String().Const("world")),
 			"map":     schema.Object().AdditionalProperties(schema.Always()),
-			"record": schema.Record(map[string]schema.Builder{
+			"record": schema.Record(schema.BuilderMap{
 				"foo": schema.String(),
 			}),
 			"anyOf": schema.AnyOf(schema.String(), schema.Number()),
@@ -87,7 +87,7 @@ func (testSchemaProvider) Schema() (*schema.Schema, *schema.Schema) {
 			"double":       schema.Tuple(schema.String(), schema.Number()),
 			"triple":       schema.Tuple(schema.String(), schema.Number(), schema.Boolean()),
 			"dependentReq": schema.Object().
-				Properties(map[string]schema.Builder{
+				Properties(schema.BuilderMap{
 					"foo": schema.String(),
 					"bar": schema.Number(),
 				}).DependentRequired(map[string][]string{"foo": {"bar"}}),

eval/expr.go

@@ -185,8 +185,8 @@ func (x *expr) export(environment string) esc.Expr {
 			ex.Builtin = &esc.BuiltinExpr{
 				Name:      name,
 				NameRange: convertRange(repr.node.Name().Syntax().Syntax().Range(), environment),
-				ArgSchema: schema.Record(map[string]schema.Builder{
-					"provider": schema.String(),
+				ArgSchema: schema.Record(schema.SchemaMap{
+					"provider": schema.String().Schema(),
 					"inputs":   repr.inputSchema,
 				}).Schema(),
 				Arg: esc.Expr{

eval/value.go

@@ -39,6 +39,9 @@ type value struct {
 	base   *value         // the base value, if any
 	schema *schema.Schema // the value's schema
 
+	mergedKeys []string   // the value's merged keys. computed lazily--use keys().
+	exported   *esc.Value // non-nil if this value has already been exported
+
 	// true if the value is unknown (e.g. because it did not evaluate successfully or is the result of an unevaluated
 	// fn::open)
 	unknown bool
@@ -149,20 +152,36 @@ func (v *value) combine(others ...*value) {
 // keys returns the value's keys if the value is an object. This method should be used instead of accessing the
 // underlying map[string]*value directly, as it takes JSON merge patch semantics into account.
 func (v *value) keys() []string {
-	keySet := make(map[string]struct{})
-	for v != nil {
+	if v == nil {
+		return nil
+	}
+	if v.mergedKeys == nil {
 		m, ok := v.repr.(map[string]*value)
 		if !ok {
-			break
+			return nil
 		}
-		for k := range m {
-			keySet[k] = struct{}{}
+
+		baseKeys := v.base.keys()
+		if len(baseKeys) == 0 {
+			v.mergedKeys = maps.Keys(m)
+		} else {
+			l := len(baseKeys)
+			if l < len(m) {
+				l = len(m)
+			}
+			keySet := make(map[string]struct{}, l)
+
+			for _, k := range baseKeys {
+				keySet[k] = struct{}{}
+			}
+			for k := range m {
+				keySet[k] = struct{}{}
+			}
+			v.mergedKeys = maps.Keys(keySet)
 		}
-		v = v.base
+		sort.Strings(v.mergedKeys)
 	}
-	keys := maps.Keys(keySet)
-	sort.Strings(keys)
-	return keys
+	return v.mergedKeys
 }
 
 // property returns the named property (if any) as per JSON merge patch semantics. If the receiver is unknown,
@@ -269,6 +288,10 @@ func (v *value) toString() (str string, unknown bool, secret bool) {
 
 // export converts the value into its serializable representation.
 func (v *value) export(environment string) esc.Value {
+	if v.exported != nil {
+		return *v.exported
+	}
+
 	var pv any
 	switch repr := v.repr.(type) {
 	case []*value:
@@ -295,7 +318,7 @@ func (v *value) export(environment string) esc.Value {
 		base = &b
 	}
 
-	return esc.Value{
+	v.exported = &esc.Value{
 		Value:   pv,
 		Secret:  v.secret,
 		Unknown: v.unknown,
@@ -304,6 +327,7 @@ func (v *value) export(environment string) esc.Value {
 			Base: base,
 		},
 	}
+	return *v.exported
 }
 
 // unexport creates a value from a Value. This is used when interacting with providers, as the Provider API works on
@@ -327,7 +351,7 @@ func unexport(v esc.Value, x *expr) *value {
 		}
 		vv.repr, vv.schema = a, schema.Tuple(items...).Schema()
 	case map[string]esc.Value:
-		m, properties := make(map[string]*value, len(pv)), make(map[string]schema.Builder, len(pv))
+		m, properties := make(map[string]*value, len(pv)), make(schema.SchemaMap, len(pv))
 		for k, v := range pv {
 			uv := unexport(v, x)
 			m[k], properties[k] = uv, uv.schema
@@ -348,7 +372,12 @@ func mergedSchema(base, top *schema.Schema) *schema.Schema {
 		return top
 	}
 
-	record := make(map[string]schema.Builder)
+	l := len(base.Properties)
+	if l < len(top.Properties) {
+		l = len(top.Properties)
+	}
+
+	record := make(schema.SchemaMap, l)
 	for k, base := range base.Properties {
 		record[k] = base
 	}

schema/objects.go

@@ -30,11 +30,13 @@ func Object() *ObjectBuilder {
 	return &ObjectBuilder{}
 }
 
-func Record(m map[string]Builder) *ObjectBuilder {
-	names := maps.Keys(m)
+func Record(m MapBuilder) *ObjectBuilder {
+	props := m.Build()
+
+	names := maps.Keys(props)
 	sort.Strings(names)
 
-	return Object().Properties(m).Required(names...)
+	return Object().Properties(SchemaMap(props)).Required(names...)
 }
 
 func (b *ObjectBuilder) Defs(defs map[string]Builder) *ObjectBuilder {
@@ -53,11 +55,8 @@ func (b *ObjectBuilder) OneOf(oneOf ...Builder) *ObjectBuilder {
 	return buildOneOf(b, oneOf)
 }
 
-func (b *ObjectBuilder) Properties(m map[string]Builder) *ObjectBuilder {
-	b.s.Properties = make(map[string]*Schema, len(m))
-	for k, v := range m {
-		b.s.Properties[k] = v.Schema()
-	}
+func (b *ObjectBuilder) Properties(m MapBuilder) *ObjectBuilder {
+	b.s.Properties = m.Build()
 	return b
 }
 

schema/objects_test.go

@@ -23,7 +23,7 @@ import (
 
 func TestObjects(t *testing.T) {
 	t.Run("record", func(t *testing.T) {
-		record := map[string]Builder{
+		record := BuilderMap{
 			"bool":   Boolean(),
 			"string": String(),
 		}
@@ -39,7 +39,7 @@ func TestObjects(t *testing.T) {
 	})
 
 	t.Run("defs", func(t *testing.T) {
-		defs := map[string]Builder{
+		defs := BuilderMap{
 			"bool": Boolean(),
 		}
 		s := Object().
@@ -89,7 +89,7 @@ func TestObjects(t *testing.T) {
 	})
 
 	t.Run("properties", func(t *testing.T) {
-		properties := map[string]Builder{
+		properties := BuilderMap{
 			"bool":   Boolean(),
 			"string": String(),
 		}

schema/schema.go

@@ -30,6 +30,26 @@ type Builder interface {
 	Schema() *Schema
 }
 
+type MapBuilder interface {
+	Build() map[string]*Schema
+}
+
+type BuilderMap map[string]Builder
+
+func (m BuilderMap) Build() map[string]*Schema {
+	s := make(map[string]*Schema, len(m))
+	for k, v := range m {
+		s[k] = v.Schema()
+	}
+	return s
+}
+
+type SchemaMap map[string]*Schema
+
+func (m SchemaMap) Build() map[string]*Schema {
+	return m
+}
+
 func Never() *Schema {
 	return &Schema{Never: true}
 }

schema/schema_test.go

@@ -111,7 +111,7 @@ func TestProperty(t *testing.T) {
 
 	t.Run("valid", func(t *testing.T) {
 		s := Object().
-			Properties(map[string]Builder{"a": String(), "b": Number()}).
+			Properties(BuilderMap{"a": String(), "b": Number()}).
 			AdditionalProperties(Boolean()).
 			Schema()
 
@@ -122,12 +122,12 @@ func TestProperty(t *testing.T) {
 
 	t.Run("anyOf", func(t *testing.T) {
 		a := Object().
-			Properties(map[string]Builder{"a": String(), "b": Number()}).
+			Properties(BuilderMap{"a": String(), "b": Number()}).
 			AdditionalProperties(Never()).
 			Schema()
 
 		b := Object().
-			Properties(map[string]Builder{"a": Number(), "b": Boolean()}).
+			Properties(BuilderMap{"a": Number(), "b": Boolean()}).
 			AdditionalProperties(Number()).
 			Schema()
 
@@ -141,12 +141,12 @@ func TestProperty(t *testing.T) {
 
 	t.Run("oneOf", func(t *testing.T) {
 		a := Object().
-			Properties(map[string]Builder{"a": String(), "b": Number()}).
+			Properties(BuilderMap{"a": String(), "b": Number()}).
 			AdditionalProperties(Never()).
 			Schema()
 
 		b := Object().
-			Properties(map[string]Builder{"a": Number(), "b": Boolean()}).
+			Properties(BuilderMap{"a": Number(), "b": Boolean()}).
 			AdditionalProperties(Number()).
 			Schema()