Support for ifs in array comprehensions

This commit is contained in:
Stanisław Barzowski 2017-09-08 11:23:21 -04:00 committed by Dave Cunningham
parent 1e4797071f
commit 7ebcb06c98
15 changed files with 86 additions and 2 deletions

View File

@ -217,9 +217,26 @@ func simpleLambda(body ast.Node, paramName ast.Identifier) ast.Node {
}
}
func buildAnd(left ast.Node, right ast.Node) ast.Node {
return &ast.Binary{Op: ast.BopAnd, Left: left, Right: right}
}
func desugarForSpec(inside ast.Node, forSpec *ast.ForSpec) (ast.Node, error) {
// TODO(sbarzowski) support ifs
function := simpleLambda(inside, forSpec.VarName)
var body ast.Node
if len(forSpec.Conditions) > 0 {
cond := forSpec.Conditions[0].Expr
for i := 1; i < len(forSpec.Conditions); i++ {
cond = buildAnd(cond, forSpec.Conditions[i].Expr)
}
body = &ast.Conditional{
Cond: cond,
BranchTrue: inside,
BranchFalse: &ast.Array{},
}
} else {
body = inside
}
function := simpleLambda(body, forSpec.VarName)
current := buildStdCall("flatMap", function, forSpec.Expr)
if forSpec.Outer == nil {
return current, nil

4
testdata/arrcomp_if.golden vendored Normal file
View File

@ -0,0 +1,4 @@
[
4,
5
]

1
testdata/arrcomp_if.input vendored Normal file
View File

@ -0,0 +1 @@
[x for x in [1, 2, 3, 4, 5] if x > 3]

26
testdata/arrcomp_if2.golden vendored Normal file
View File

@ -0,0 +1,26 @@
[
[
4,
1
],
[
4,
2
],
[
4,
3
],
[
5,
1
],
[
5,
2
],
[
5,
3
]
]

1
testdata/arrcomp_if2.input vendored Normal file
View File

@ -0,0 +1 @@
[[x, y] for x in [1, 2, 3, 4, 5] for y in [1, 2, 3, 4, 5] if x > 3 if y < 4]

26
testdata/arrcomp_if3.golden vendored Normal file
View File

@ -0,0 +1,26 @@
[
[
1,
3
],
[
1,
4
],
[
1,
5
],
[
2,
3
],
[
2,
4
],
[
2,
5
]
]

1
testdata/arrcomp_if3.input vendored Normal file
View File

@ -0,0 +1 @@
[[x, y] for x in [1,2,3,4,5] if x < 3 for y in [1,2,3,4,5] if y > 2]

1
testdata/arrcomp_if4.golden vendored Normal file
View File

@ -0,0 +1 @@
testdata/arrcomp_if4:1:33-34 Unknown variable: y

1
testdata/arrcomp_if4.input vendored Normal file
View File

@ -0,0 +1 @@
[[x, y] for x in [1,2,3,4,5] if y == 3 for y in [1,2,3,4,5]]

1
testdata/arrcomp_if5.golden vendored Normal file
View File

@ -0,0 +1 @@
[ ]

1
testdata/arrcomp_if5.input vendored Normal file
View File

@ -0,0 +1 @@
[x for x in [] if error "x"]

1
testdata/arrcomp_if6.golden vendored Normal file
View File

@ -0,0 +1 @@
RUNTIME ERROR: x

1
testdata/arrcomp_if6.input vendored Normal file
View File

@ -0,0 +1 @@
[x for x in [1] if error "x"]

1
testdata/arrcomp_if7.golden vendored Normal file
View File

@ -0,0 +1 @@
RUNTIME ERROR: Unexpected type number, expected boolean

1
testdata/arrcomp_if7.input vendored Normal file
View File

@ -0,0 +1 @@
[x for x in [1, 2, 3] if 42]