From d76e4fb21311d030ba49026b1362a4c4a01b4ef2 Mon Sep 17 00:00:00 2001
From: Joe Betz <jpbetz@google.com>
Date: Tue, 27 Feb 2024 16:30:03 -0500
Subject: [PATCH] Track sizes of comprehension results

---
 checker/cost.go      |  3 +++
 checker/cost_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/checker/cost.go b/checker/cost.go
index 3470d0a3f..04244694d 100644
--- a/checker/cost.go
+++ b/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)
 
diff --git a/checker/cost_test.go b/checker/cost_test.go
index c6c956439..eb34c0520 100644
--- a/checker/cost_test.go
+++ b/checker/cost_test.go
@@ -503,6 +503,31 @@ 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},
+		},
 	}
 
 	for _, tc := range cases {