Skip to content

Commit

Permalink
Track sizes of comprehension results
Browse files Browse the repository at this point in the history
  • Loading branch information
jpbetz committed Feb 27, 2024
1 parent b1ed395 commit 030abf1
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 0 deletions.
3 changes: 3 additions & 0 deletions checker/cost.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
40 changes: 40 additions & 0 deletions checker/cost_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down

0 comments on commit 030abf1

Please sign in to comment.