Track sizes of comprehension results

checker/cost.go

@@ -506,6 +506,9 @@ func (c *coster) costComprehension(e ast.Expr) CostEstimate {
 	c.iterRanges.pop(comp.IterVar())
 	sum = sum.Add(c.cost(comp.Result()))
 	rangeCnt := c.sizeEstimate(c.newAstNode(comp.IterRange()))
+
+	c.computedSizes[e.ID()] = rangeCnt
+
 	rangeCost := rangeCnt.MultiplyByCost(stepCost.Add(loopCost))
 	sum = sum.Add(rangeCost)
 

checker/cost_test.go

@@ -503,6 +503,46 @@ func TestCost(t *testing.T) {
 			},
 			wanted: CostEstimate{Min: 3, Max: 3},
 		},
+		{
+			name:   ".filter list literal",
+			expr:   `[1,2,3,4,5].filter(x, x % 2 == 0)`,
+			wanted: CostEstimate{Min: 41, Max: 101},
+		},
+		{
+			name:   ".map list literal",
+			expr:   `[1,2,3,4,5].map(x, x)`,
+			wanted: CostEstimate{Min: 86, Max: 86},
+		},
+		{
+			name:   ".map.filter list literal",
+			expr:   `[1,2,3,4,5].map(x, x).filter(x, x % 2 == 0)`,
+			wanted: CostEstimate{Min: 117, Max: 177},
+		},
+		{
+			name:   ".map.exists list literal",
+			expr:   `[1,2,3,4,5].map(x, x).exists(x, x % 2 == 0) == true`,
+			wanted: CostEstimate{Min: 108, Max: 123},
+		},
+		{
+			name:   ".map.map list literal",
+			expr:   `[1,2,3,4,5].map(x, x).map(x, x)`,
+			wanted: CostEstimate{Min: 162, Max: 162},
+		},
+		{
+			name:   ".map list literal selection",
+			expr:   `[1,2,3,4,5].map(x, x)[4]`,
+			wanted: CostEstimate{Min: 87, Max: 87},
+		},
+		{
+			name:   "nested array selection",
+			expr:   `[[1,2],[1,2],[1,2],[1,2],[1,2]][4]`,
+			wanted: CostEstimate{Min: 61, Max: 61},
+		},
+		{
+			name:   "nested array selection",
+			expr:   `{'a': [1,2], 'b': [1,2], 'c': [1,2], 'd': [1,2], 'e': [1,2]}.b`,
+			wanted: CostEstimate{Min: 81, Max: 81},
+		},
 	}
 
 	for _, tc := range cases {