Support for ifs in array comprehensions

desugarer.go

@@ -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
+	var body ast.Node
-	function := simpleLambda(inside, forSpec.VarName)
+	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

testdata/arrcomp_if.golden

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

testdata/arrcomp_if.input

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

testdata/arrcomp_if2.golden

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

testdata/arrcomp_if2.input

@@ -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]

testdata/arrcomp_if3.golden

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

testdata/arrcomp_if3.input

@@ -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]

testdata/arrcomp_if4.golden

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

testdata/arrcomp_if4.input

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

testdata/arrcomp_if5.golden

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

testdata/arrcomp_if5.input

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

testdata/arrcomp_if6.golden

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

testdata/arrcomp_if6.input

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

testdata/arrcomp_if7.golden

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

testdata/arrcomp_if7.input

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