Fix runtime cost ref.Val stack (#517)

Fix runtime cost ref.Val stack

interpreter/interpretable.go

@@ -677,7 +677,19 @@ func (m *evalMap) Eval(ctx Activation) ref.Val {
 }
 
 func (m *evalMap) InitVals() []Interpretable {
-	return append(m.keys, m.vals...)
+	if len(m.keys) != len(m.vals) {
+		return nil
+	}
+	result := make([]Interpretable, len(m.keys)+len(m.vals))
+	idx := 0
+	for i, k := range m.keys {
+		v := m.vals[i]
+		result[idx] = k
+		idx++
+		result[idx] = v
+		idx++
+	}
+	return result
 }
 
 func (m *evalMap) Type() ref.Type {

interpreter/runtimecost.go

@@ -41,26 +41,42 @@ func CostObserver(tracker *CostTracker) EvalObserver {
 		case ConstantQualifier:
 			// TODO: Push identifiers on to the stack before observing constant qualifiers that apply to them
 			// and enable the below pop. Once enabled this can case can be collapsed into the Qualifier case.
-			//tracker.stack.pop(1)
 			tracker.cost++
 		case InterpretableConst:
 			// zero cost
 		case InterpretableAttribute:
-			// Ternary has no direct cost. All cost is from the conditional and the true/false branch expressions.
-			_, isConditional := t.Attr().(*conditionalAttribute)
-			if !isConditional {
+			switch a := t.Attr().(type) {
+			case *conditionalAttribute:
+				// Ternary has no direct cost. All cost is from the conditional and the true/false branch expressions.
+				tracker.stack.drop(a.falsy.ID(), a.truthy.ID(), a.expr.ID())
+			default:
+				tracker.stack.drop(t.Attr().ID())
 				tracker.cost += common.SelectAndIdentCost
 			}
-		case *evalExhaustiveConditional, *evalOr, *evalAnd, *evalExhaustiveOr, *evalExhaustiveAnd:
+		case *evalExhaustiveConditional:
 			// Ternary has no direct cost. All cost is from the conditional and the true/false branch expressions.
+			tracker.stack.drop(t.attr.falsy.ID(), t.attr.truthy.ID(), t.attr.expr.ID())
+
+		// While the field names are identical, the boolean operation eval structs do not share an interface and so
+		// must be handled individually.
+		case *evalOr:
+			tracker.stack.drop(t.rhs.ID(), t.lhs.ID())
+		case *evalAnd:
+			tracker.stack.drop(t.rhs.ID(), t.lhs.ID())
+		case *evalExhaustiveOr:
+			tracker.stack.drop(t.rhs.ID(), t.lhs.ID())
+		case *evalExhaustiveAnd:
+			tracker.stack.drop(t.rhs.ID(), t.lhs.ID())
+		case *evalFold:
+			tracker.stack.drop(t.iterRange.ID())
 		case Qualifier:
-			tracker.stack.pop(1)
 			tracker.cost++
 		case InterpretableCall:
-			if argVals, ok := tracker.stack.pop(len(t.Args())); ok {
+			if argVals, ok := tracker.stack.dropArgs(t.Args()); ok {
 				tracker.cost += tracker.costCall(t, argVals, val)
 			}
 		case InterpretableConstructor:
+			tracker.stack.dropArgs(t.InitVals())
 			switch t.Type() {
 			case types.ListType:
 				tracker.cost += common.ListCreateBaseCost
@@ -70,7 +86,7 @@ func CostObserver(tracker *CostTracker) EvalObserver {
 				tracker.cost += common.StructCreateBaseCost
 			}
 		}
-		tracker.stack.push(val)
+		tracker.stack.push(val, id)
 
 		if tracker.Limit != nil && tracker.cost > *tracker.Limit {
 			panic(EvalCancelledError{Cause: CostLimitExceeded, Message: "operation cancelled: actual cost limit exceeded"})
@@ -170,19 +186,56 @@ func (c CostTracker) actualSize(value ref.Val) uint64 {
 	return 1
 }
 
+type stackVal struct {
+	Val ref.Val
+	ID  int64
+}
+
 // refValStack keeps track of values of the stack for cost calculation purposes
-type refValStack []ref.Val
+type refValStack []stackVal
 
-func (s *refValStack) push(value ref.Val) {
+func (s *refValStack) push(val ref.Val, id int64) {
+	value := stackVal{Val: val, ID: id}
 	*s = append(*s, value)
 }
 
-func (s *refValStack) pop(count int) ([]ref.Val, bool) {
-	if len(*s) < count {
+// TODO: Allowing drop and dropArgs to remove stack items above the IDs they are provided is a workaround. drop and dropArgs
+// should find and remove only the stack items matching the provided IDs once all attributes are properly pushed and popped from stack.
+
+// drop searches the stack for each ID and removes the ID and all stack items above it.
+// If none of the IDs are found, the stack is not modified.
+// WARNING: It is possible for multiple expressions with the same ID to exist (due to how macros are implemented) so it's
+// possible that a dropped ID will remain on the stack.  They should be removed when IDs on the stack are popped.
+func (s *refValStack) drop(ids ...int64) {
+	for _, id := range ids {
+		for idx := len(*s) - 1; idx >= 0; idx-- {
+			if (*s)[idx].ID == id {
+				*s = (*s)[:idx]
+				break
+			}
+		}
+	}
+}
+
+// dropArgs searches the stack for all the args by their IDs, accumulates their associated ref.Vals and drops any
+// stack items above any of the arg IDs. If any of the IDs are not found the stack, false is returned.
+// Args are assumed to be found in the stack in reverse order, i.e. the last arg is expected to be found highest in
+// the stack.
+// WARNING: It is possible for multiple expressions with the same ID to exist (due to how macros are implemented) so it's
+// possible that a dropped ID will remain on the stack.  They should be removed when IDs on the stack are popped.
+func (s *refValStack) dropArgs(args []Interpretable) ([]ref.Val, bool) {
+	result := make([]ref.Val, len(args))
+argloop:
+	for nIdx := len(args) - 1; nIdx >= 0; nIdx-- {
+		for idx := len(*s) - 1; idx >= 0; idx-- {
+			if (*s)[idx].ID == args[nIdx].ID() {
+				el := (*s)[idx]
+				*s = (*s)[:idx]
+				result[nIdx] = el.Val
+				continue argloop
+			}
+		}
 		return nil, false
 	}
-	idx := len(*s) - count
-	el := (*s)[idx:]
-	*s = (*s)[:idx]
-	return el, true
+	return result, true
 }

interpreter/runtimecost_test.go

@@ -153,6 +153,10 @@ func computeCost(t *testing.T, expr string, decls []*exprpb.Decl, ctx Activation
 		}
 	}()
 	prg.Eval(ctx)
+	// TODO: enable this once all attributes are properly pushed and popped from stack.
+	//if len(costTracker.stack) != 1 {
+	//	t.Fatalf(`Expected resulting stack size to be 1 but got %d: %#+v`, len(costTracker.stack), costTracker.stack)
+	//}
 	return costTracker.cost, est, err
 }
 
@@ -264,7 +268,7 @@ func TestRuntimeCost(t *testing.T) {
 		},
 		{
 			name:  "select: array index",
-			expr:  `input[1]`,
+			expr:  `input[0]`,
 			decls: []*exprpb.Decl{decls.NewVar("input", decls.NewListType(decls.String))},
 			want:  2,
 			in:    map[string]interface{}{"input": []string{"v"}},
@@ -293,7 +297,7 @@ func TestRuntimeCost(t *testing.T) {
 		},
 		{
 			name:  "expr select: array index",
-			expr:  `input[3-2]`,
+			expr:  `input[3-3]`,
 			decls: []*exprpb.Decl{decls.NewVar("input", decls.NewListType(decls.String))},
 			want:  3,
 			in:    map[string]interface{}{"input": []string{"v"}},
@@ -373,6 +377,11 @@ func TestRuntimeCost(t *testing.T) {
 		},
 		{
 			name: "or",
+			expr: `false || false`,
+			want: 0,
+		},
+		{
+			name: "or short-circuit",
 			expr: `true || false`,
 			want: 0,
 		},
@@ -381,6 +390,11 @@ func TestRuntimeCost(t *testing.T) {
 			expr: `true && false`,
 			want: 0,
 		},
+		{
+			name: "and short-circuit",
+			expr: `false && true`,
+			want: 0,
+		},
 		{
 			name: "lt",
 			expr: `1 < 2`,
@@ -651,6 +665,13 @@ func TestRuntimeCost(t *testing.T) {
 			in:    map[string]interface{}{},
 			want:  3,
 		},
+		{
+			name:  "list map literal",
+			expr:  `[{'k1': 1}, {'k2': 2}].all(x, true)`,
+			decls: []*exprpb.Decl{},
+			in:    map[string]interface{}{},
+			want:  77,
+		},
 	}
 
 	for _, tc := range cases {