List append optimizations for comprehension loops (#491)

cel/cel_test.go

@@ -857,7 +857,7 @@ func TestContextEval(t *testing.T) {
 		t.Errorf("ContextEval() got %v, wanted 75", out)
 	}
 
-	evalCtx, cancel := context.WithTimeout(ctx, 100*time.Microsecond)
+	evalCtx, cancel := context.WithTimeout(ctx, 50*time.Microsecond)
 	defer cancel()
 
 	out, _, err = ContextEval(evalCtx, prg, map[string]interface{}{"items": items})
@@ -869,6 +869,40 @@ func TestContextEval(t *testing.T) {
 	}
 }
 
+func BenchmarkContextEval(b *testing.B) {
+	env, err := NewEnv(
+		Declarations(
+			decls.NewVar("items", decls.NewListType(decls.Int)),
+		),
+	)
+	if err != nil {
+		b.Fatalf("NewEnv() failed: %v", err)
+	}
+	ast, iss := env.Compile("items.map(i, i * 2).filter(i, i >= 50).size()")
+	if iss.Err() != nil {
+		b.Fatalf("env.Compile(expr) failed: %v", iss.Err())
+	}
+	prg, err := env.Program(ast, EvalOptions(OptOptimize))
+	if err != nil {
+		b.Fatalf("env.Program() failed: %v", err)
+	}
+
+	ctx := context.TODO()
+	items := make([]int64, 100)
+	for i := int64(0); i < 100; i++ {
+		items[i] = i
+	}
+	for i := 0; i < b.N; i++ {
+		out, _, err := ContextEval(ctx, prg, map[string]interface{}{"items": items})
+		if err != nil {
+			b.Fatalf("ContextEval() failed: %v", err)
+		}
+		if out != types.Int(75) {
+			b.Errorf("ContextEval() got %v, wanted 75", out)
+		}
+	}
+}
+
 func TestEvalRecover(t *testing.T) {
 	e, err := NewEnv(
 		Declarations(

common/types/list.go

@@ -95,6 +95,18 @@ func NewJSONList(adapter ref.TypeAdapter, l *structpb.ListValue) traits.Lister {
 	}
 }
 
+// NewMutableList creates a new mutable list whose internal state can be modified.
+//
+// The mutable list only handles `Add` calls correctly as it is intended only for use within
+// comprehension loops which generate an immutable result upon completion.
+func NewMutableList(adapter ref.TypeAdapter) traits.Lister {
+	return &mutableList{
+		TypeAdapter:   adapter,
+		baseList:      nil,
+		mutableValues: []ref.Val{},
+	}
+}
+
 // baseList points to a list containing elements of any type.
 // The `value` is an array of native values, and refValue is its reflection object.
 // The `ref.TypeAdapter` enables native type to CEL type conversions.
@@ -279,6 +291,32 @@ func (l *baseList) Value() interface{} {
 	return l.value
 }
 
+// mutableList aggregates values into its internal storage. For use with internal CEL variables only.
+type mutableList struct {
+	ref.TypeAdapter
+	*baseList
+	mutableValues []ref.Val
+}
+
+// Add copies elements from the other list into the internal storage of the mutable list.
+func (l *mutableList) Add(other ref.Val) ref.Val {
+	otherList, ok := other.(traits.Lister)
+	if !ok {
+		return MaybeNoSuchOverloadErr(otherList)
+	}
+	for i := IntZero; i < otherList.Size().(Int); i++ {
+		l.mutableValues = append(l.mutableValues, otherList.Get(i))
+	}
+	return l
+}
+
+// ToImmutableList returns an immutable list based on the internal storage of the mutable list.
+func (l *mutableList) ToImmutableList() traits.Lister {
+	// The reference to internal state is guaranteed to be safe as this call is only performed
+	// when mutations have been completed.
+	return NewRefValList(l.TypeAdapter, l.mutableValues)
+}
+
 // concatList combines two list implementations together into a view.
 // The `ref.TypeAdapter` enables native type to CEL type conversions.
 type concatList struct {

common/types/traits/lister.go

@@ -25,3 +25,8 @@ type Lister interface {
 	Iterable
 	Sizer
 }
+
+// MutableLister interface which emits an immutable result after an intermediate computation.
+type MutableLister interface {
+	ToImmutableList() Lister
+}

interpreter/interpretable.go

@@ -696,6 +696,7 @@ type evalFold struct {
 	cond      Interpretable
 	step      Interpretable
 	result    Interpretable
+	adapter   ref.TypeAdapter
 }
 
 // ID implements the Interpretable interface method.
@@ -714,6 +715,14 @@ func (fold *evalFold) Eval(ctx Activation) ref.Val {
 	accuCtx.parent = ctx
 	accuCtx.name = fold.accuVar
 	accuCtx.val = fold.accu.Eval(ctx)
+	// If the accumulator starts as an empty list, then the comprehension will build a list
+	// so create a mutable list to optimize the cost of the inner loop.
+	l, ok := accuCtx.val.(traits.Lister)
+	buildingList := false
+	if ok && l.Size() == types.IntZero {
+		buildingList = true
+		accuCtx.val = types.NewMutableList(fold.adapter)
+	}
 	iterCtx := varActivationPool.Get().(*varActivation)
 	iterCtx.parent = accuCtx
 	iterCtx.name = fold.iterVar
@@ -734,6 +743,10 @@ func (fold *evalFold) Eval(ctx Activation) ref.Val {
 	}
 	// Compute the result.
 	res := fold.result.Eval(accuCtx)
+	// Convert a mutable list to an immutable one, if the comprehension has generated a list as a result.
+	if buildingList {
+		res = res.(traits.MutableLister).ToImmutableList()
+	}
 	varActivationPool.Put(iterCtx)
 	varActivationPool.Put(accuCtx)
 	return res

interpreter/interpreter_test.go

@@ -557,7 +557,7 @@ var (
 		},
 		{
 			name: "macro_filter",
-			expr: `[1, 2, 3].filter(x, x > 2) == [3]`,
+			expr: `[-10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3].filter(x, x > 0) == [1, 2, 3]`,
 		},
 		{
 			name:           "macro_has_map_key",
@@ -1104,10 +1104,10 @@ func BenchmarkInterpreter(b *testing.B) {
 			b.Fatal(err)
 		}
 		// Benchmark the eval.
-		b.Run(tst.name, func(bb *testing.B) {
+		b.Run(tst.name, func(b *testing.B) {
 			b.ResetTimer()
 			b.ReportAllocs()
-			for i := 0; i < bb.N; i++ {
+			for i := 0; i < b.N; i++ {
 				prg.Eval(vars)
 			}
 		})

interpreter/planner.go

@@ -617,6 +617,7 @@ func (p *planner) planComprehension(expr *exprpb.Expr) (Interpretable, error) {
 		cond:      cond,
 		step:      step,
 		result:    result,
+		adapter:   p.adapter,
 	}, nil
 }