optimize performace when setting multiple variables at once (closes #4)

Author: xrstf

pkg/builtin/lists/functions.go

@@ -68,12 +68,17 @@ func rangeVectorFunction(ctx types.Context, data []any, namingVec ast.Expression
 	)
 
 	for i, item := range data {
-		// do not use separate contexts for each loop iteration, as the loop might build up a counter
-		loopCtx = loopCtx.WithVariable(loopVarName, item)
+		vars := map[string]any{
+			loopVarName: item,
+		}
+
 		if loopIndexName != "" {
-			loopCtx = loopCtx.WithVariable(loopIndexName, i)
+			vars[loopIndexName] = i
 		}
 
+		// do not use separate contexts for each loop iteration, as the loop might build up a counter
+		loopCtx = loopCtx.WithVariables(vars)
+
 		loopCtx, result, err = eval.EvalExpression(loopCtx, expr)
 		if err != nil {
 			return nil, err
@@ -98,12 +103,17 @@ func rangeObjectFunction(ctx types.Context, data map[string]any, namingVec ast.E
 	)
 
 	for key, value := range data {
-		// do not use separate contexts for each loop iteration, as the loop might build up a counter
-		loopCtx = loopCtx.WithVariable(loopVarName, value)
+		vars := map[string]any{
+			loopVarName: value,
+		}
+
 		if loopIndexName != "" {
-			loopCtx = loopCtx.WithVariable(loopIndexName, key)
+			vars[loopIndexName] = key
 		}
 
+		// do not use separate contexts for each loop iteration, as the loop might build up a counter
+		loopCtx = loopCtx.WithVariables(vars)
+
 		loopCtx, result, err = eval.EvalExpression(loopCtx, expr)
 		if err != nil {
 			return nil, err
@@ -140,11 +150,16 @@ func mapVectorExpressionFunction(ctx types.Context, data []any, namingVec ast.Ex
 	}
 
 	mapHandler := func(ctx types.Context, index any, value any) (types.Context, any, error) {
-		ctx = ctx.WithVariable(valueVarName, value)
+		vars := map[string]any{
+			valueVarName: value,
+		}
+
 		if indexVarName != "" {
-			ctx = ctx.WithVariable(indexVarName, index)
+			vars[indexVarName] = index
 		}
 
+		ctx = ctx.WithVariables(vars)
+
 		return eval.EvalExpression(ctx, expr)
 	}
 
@@ -197,11 +212,16 @@ func mapObjectExpressionFunction(ctx types.Context, data map[string]any, namingV
 	}
 
 	mapHandler := func(ctx types.Context, key any, value any) (types.Context, any, error) {
-		ctx = ctx.WithVariable(valueVarName, value)
+		vars := map[string]any{
+			valueVarName: value,
+		}
+
 		if keyVarName != "" {
-			ctx = ctx.WithVariable(keyVarName, key)
+			vars[keyVarName] = key
 		}
 
+		ctx = ctx.WithVariables(vars)
+
 		return eval.EvalExpression(ctx, expr)
 	}
 
@@ -254,11 +274,16 @@ func filterVectorExpressionFunction(ctx types.Context, data []any, namingVec ast
 	}
 
 	mapHandler := func(ctx types.Context, index any, value any) (types.Context, any, error) {
-		ctx = ctx.WithVariable(valueVarName, value)
+		vars := map[string]any{
+			valueVarName: value,
+		}
+
 		if indexVarName != "" {
-			ctx = ctx.WithVariable(indexVarName, index)
+			vars[indexVarName] = index
 		}
 
+		ctx = ctx.WithVariables(vars)
+
 		return eval.EvalExpression(ctx, expr)
 	}
 
@@ -317,11 +342,16 @@ func filterObjectExpressionFunction(ctx types.Context, data map[string]any, nami
 	}
 
 	mapHandler := func(ctx types.Context, key any, value any) (types.Context, any, error) {
-		ctx = ctx.WithVariable(valueVarName, value)
+		vars := map[string]any{
+			valueVarName: value,
+		}
+
 		if keyVarName != "" {
-			ctx = ctx.WithVariable(keyVarName, key)
+			vars[keyVarName] = key
 		}
 
+		ctx = ctx.WithVariables(vars)
+
 		return eval.EvalExpression(ctx, expr)
 	}
 

pkg/builtin/rudifunc/functions.go

@@ -75,17 +75,17 @@ func (f rudispaceFunc) Evaluate(ctx types.Context, args []ast.Expression) (any,
 		return nil, fmt.Errorf("expected %d argument(s), got %d", len(f.params), len(args))
 	}
 
-	funcCtx := ctx
+	funcArgs := map[string]any{}
 	for i, paramName := range f.params {
 		_, arg, err := eval.EvalExpression(ctx, args[i])
 		if err != nil {
 			return nil, err
 		}
 
-		funcCtx = funcCtx.WithVariable(paramName, arg)
+		funcArgs[paramName] = arg
 	}
 
-	_, result, err := eval.EvalExpression(funcCtx, f.body)
+	_, result, err := eval.EvalExpression(ctx.WithVariables(funcArgs), f.body)
 
 	return result, err
 }

pkg/eval/types/context.go

@@ -113,6 +113,21 @@ func (c Context) WithVariable(name string, val any) Context {
 	}
 }
 
+func (c Context) WithVariables(vars map[string]any) Context {
+	if len(vars) == 0 {
+		return c
+	}
+
+	return Context{
+		ctx:        c.ctx,
+		document:   c.document,
+		fixedFuncs: c.fixedFuncs,
+		userFuncs:  c.userFuncs,
+		variables:  c.variables.WithMany(vars),
+		coalescer:  c.coalescer,
+	}
+}
+
 func (c Context) WithCoalescer(coalescer coalescing.Coalescer) Context {
 	return Context{
 		ctx:        c.ctx,
@@ -246,6 +261,20 @@ func (v Variables) With(name string, val any) Variables {
 	return v.DeepCopy().Set(name, val)
 }
 
+// WithMany is like With(), but for adding multiple new variables at once. This
+// should be preferred to With() to prevent unnecessary DeepCopies.
+func (v Variables) WithMany(vars map[string]any) Variables {
+	if len(vars) == 0 {
+		return v
+	}
+
+	out := v.DeepCopy()
+	for k, v := range vars {
+		out.Set(k, v)
+	}
+	return out
+}
+
 func (v Variables) DeepCopy() Variables {
 	result := NewVariables()
 	for key, val := range v {