Fix CEL cost handling of zero length replacement strings

Kubernetes-commit: f3f88b8e7b51f15c91b4f004ff4ae04ed6b25ac0
kubernetes · Aug 22, 2023 · 6394fb8 · 6394fb8
1 parent d8f0b41
commit 6394fb8
Show file tree

Hide file tree

Showing 2 changed files with 152 additions and 5 deletions.
diff --git a/pkg/cel/library/cost.go b/pkg/cel/library/cost.go
@@ -126,13 +126,47 @@ func (l *CostEstimator) EstimateCallCost(function, overloadId string, target *ch
 			sz := l.sizeEstimate(*target)
 			toReplaceSz := l.sizeEstimate(args[0])
 			replaceWithSz := l.sizeEstimate(args[1])
-			// smallest possible result: smallest input size composed of the largest possible substrings being replaced by smallest possible replacement
-			minSz := uint64(math.Ceil(float64(sz.Min)/float64(toReplaceSz.Max))) * replaceWithSz.Min
-			// largest possible result: largest input size composed of the smallest possible substrings being replaced by largest possible replacement
-			maxSz := uint64(math.Ceil(float64(sz.Max)/float64(toReplaceSz.Min))) * replaceWithSz.Max
+
+			var replaceCount, retainedSz checker.SizeEstimate
+			// find the longest replacement:
+			if toReplaceSz.Min == 0 {
+				// if the string being replaced is empty, replace surrounds all characters in the input string with the replacement.
+				if replaceCount.Max < math.MaxUint64 {
+					replaceCount.Max = sz.Max + 1
+				}
+				// Include the length of the longest possible original string length.
+				retainedSz.Max = sz.Max
+			} else if replaceWithSz.Max <= toReplaceSz.Min {
+				// If the replacement does not make the result longer, use the original string length.
+				replaceCount.Max = 0
+				retainedSz.Max = sz.Max
+			} else {
+				// Replace the smallest possible substrings with the largest possible replacement
+				// as many times as possible.
+				replaceCount.Max = uint64(math.Ceil(float64(sz.Max) / float64(toReplaceSz.Min)))
+			}
+
+			// find the shortest replacement:
+			if toReplaceSz.Max == 0 {
+				// if the string being replaced is empty, replace surrounds all characters in the input string with the replacement.
+				if replaceCount.Min < math.MaxUint64 {
+					replaceCount.Min = sz.Min + 1
+				}
+				// Include the length of the shortest possible original string length.
+				retainedSz.Min = sz.Min
+			} else if toReplaceSz.Max <= replaceWithSz.Min {
+				// If the replacement does not make the result shorter, use the original string length.
+				replaceCount.Min = 0
+				retainedSz.Min = sz.Min
+			} else {
+				// Replace the largest possible substrings being with the smallest possible replacement
+				// as many times as possible.
+				replaceCount.Min = uint64(math.Ceil(float64(sz.Min) / float64(toReplaceSz.Max)))
+			}
+			size := replaceCount.Multiply(replaceWithSz).Add(retainedSz)
 
 			// cost is the traversal plus the construction of the result
-			return &checker.CallEstimate{CostEstimate: sz.MultiplyByCostFactor(2 * common.StringTraversalCostFactor), ResultSize: &checker.SizeEstimate{Min: minSz, Max: maxSz}}
+			return &checker.CallEstimate{CostEstimate: sz.MultiplyByCostFactor(2 * common.StringTraversalCostFactor), ResultSize: &size}
 		}
 	case "split":
 		if target != nil {

diff --git a/pkg/cel/library/cost_test.go b/pkg/cel/library/cost_test.go
@@ -263,6 +263,18 @@ func TestStringLibrary(t *testing.T) {
 			expectEsimatedCost: checker.CostEstimate{Min: 3, Max: 3},
 			expectRuntimeCost:  3,
 		},
+		{
+			name:               "replace between all chars",
+			expr:               "'abc 123 def 123'.replace('', 'x')",
+			expectEsimatedCost: checker.CostEstimate{Min: 3, Max: 3},
+			expectRuntimeCost:  3,
+		},
+		{
+			name:               "replace with empty",
+			expr:               "'abc 123 def 123'.replace('123', '')",
+			expectEsimatedCost: checker.CostEstimate{Min: 3, Max: 3},
+			expectRuntimeCost:  3,
+		},
 		{
 			name:               "replace with limit",
 			expr:               "'abc 123 def 123'.replace('123', '456', 1)",
@@ -437,6 +449,107 @@ func testCost(t *testing.T, expr string, expectEsimatedCost checker.CostEstimate
 	}
 }
 
+func TestSize(t *testing.T) {
+	exactSize := func(size int) checker.SizeEstimate {
+		return checker.SizeEstimate{Min: uint64(size), Max: uint64(size)}
+	}
+	exactSizes := func(sizes ...int) []checker.SizeEstimate {
+		results := make([]checker.SizeEstimate, len(sizes))
+		for i, size := range sizes {
+			results[i] = exactSize(size)
+		}
+		return results
+	}
+	cases := []struct {
+		name       string
+		function   string
+		overload   string
+		targetSize checker.SizeEstimate
+		argSizes   []checker.SizeEstimate
+		expectSize checker.SizeEstimate
+	}{
+		{
+			name:       "replace empty with char",
+			function:   "replace",
+			targetSize: exactSize(3),     // e.g. abc
+			argSizes:   exactSizes(0, 1), // e.g. replace "" with "_"
+			expectSize: exactSize(7),     // e.g. _a_b_c_
+		},
+		{
+			name:       "maybe replace char with empty",
+			function:   "replace",
+			targetSize: exactSize(3),
+			argSizes:   exactSizes(1, 0),
+			expectSize: checker.SizeEstimate{Min: 0, Max: 3},
+		},
+		{
+			name:       "maybe replace repeated",
+			function:   "replace",
+			targetSize: exactSize(4),
+			argSizes:   exactSizes(2, 4),
+			expectSize: checker.SizeEstimate{Min: 4, Max: 8},
+		},
+		{
+			name:       "maybe replace empty",
+			function:   "replace",
+			targetSize: exactSize(4),
+			argSizes:   []checker.SizeEstimate{{Min: 0, Max: 1}, {Min: 0, Max: 2}},
+			expectSize: checker.SizeEstimate{Min: 0, Max: 14}, // len(__a__a__a__a__) == 14
+		},
+		{
+			name:       "replace non-empty size range, maybe larger",
+			function:   "replace",
+			targetSize: exactSize(4),
+			argSizes:   []checker.SizeEstimate{{Min: 1, Max: 1}, {Min: 1, Max: 2}},
+			expectSize: checker.SizeEstimate{Min: 4, Max: 8},
+		},
+		{
+			name:       "replace non-empty size range, maybe smaller",
+			function:   "replace",
+			targetSize: exactSize(4),
+			argSizes:   []checker.SizeEstimate{{Min: 1, Max: 2}, {Min: 1, Max: 1}},
+			expectSize: checker.SizeEstimate{Min: 2, Max: 4},
+		},
+	}
+	est := &CostEstimator{SizeEstimator: &testCostEstimator{}}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			var targetNode checker.AstNode = testSizeNode{size: tc.targetSize}
+			argNodes := make([]checker.AstNode, len(tc.argSizes))
+			for i, arg := range tc.argSizes {
+				argNodes[i] = testSizeNode{size: arg}
+			}
+			result := est.EstimateCallCost(tc.function, tc.overload, &targetNode, argNodes)
+			if result.ResultSize == nil {
+				t.Fatalf("Expected ResultSize but got none")
+			}
+			if *result.ResultSize != tc.expectSize {
+				t.Fatalf("Expected %+v but got %+v", tc.expectSize, *result.ResultSize)
+			}
+		})
+	}
+}
+
+type testSizeNode struct {
+	size checker.SizeEstimate
+}
+
+func (t testSizeNode) Path() []string {
+	return nil // not needed
+}
+
+func (t testSizeNode) Type() *expr.Type {
+	return nil // not needed
+}
+
+func (t testSizeNode) Expr() *expr.Expr {
+	return nil // not needed
+}
+
+func (t testSizeNode) ComputedSize() *checker.SizeEstimate {
+	return &t.size
+}
+
 type testCostEstimator struct {
 }