diff --git a/desugarer.go b/desugarer.go index 1d91f71..3f1453f 100644 --- a/desugarer.go +++ b/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 - 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 diff --git a/testdata/arrcomp_if.golden b/testdata/arrcomp_if.golden new file mode 100644 index 0000000..3cc1e3d --- /dev/null +++ b/testdata/arrcomp_if.golden @@ -0,0 +1,4 @@ +[ + 4, + 5 +] diff --git a/testdata/arrcomp_if.input b/testdata/arrcomp_if.input new file mode 100644 index 0000000..0555648 --- /dev/null +++ b/testdata/arrcomp_if.input @@ -0,0 +1 @@ +[x for x in [1, 2, 3, 4, 5] if x > 3] diff --git a/testdata/arrcomp_if2.golden b/testdata/arrcomp_if2.golden new file mode 100644 index 0000000..8506d69 --- /dev/null +++ b/testdata/arrcomp_if2.golden @@ -0,0 +1,26 @@ +[ + [ + 4, + 1 + ], + [ + 4, + 2 + ], + [ + 4, + 3 + ], + [ + 5, + 1 + ], + [ + 5, + 2 + ], + [ + 5, + 3 + ] +] diff --git a/testdata/arrcomp_if2.input b/testdata/arrcomp_if2.input new file mode 100644 index 0000000..67a3289 --- /dev/null +++ b/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] diff --git a/testdata/arrcomp_if3.golden b/testdata/arrcomp_if3.golden new file mode 100644 index 0000000..3083933 --- /dev/null +++ b/testdata/arrcomp_if3.golden @@ -0,0 +1,26 @@ +[ + [ + 1, + 3 + ], + [ + 1, + 4 + ], + [ + 1, + 5 + ], + [ + 2, + 3 + ], + [ + 2, + 4 + ], + [ + 2, + 5 + ] +] diff --git a/testdata/arrcomp_if3.input b/testdata/arrcomp_if3.input new file mode 100644 index 0000000..2f403f0 --- /dev/null +++ b/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] diff --git a/testdata/arrcomp_if4.golden b/testdata/arrcomp_if4.golden new file mode 100644 index 0000000..6e925c4 --- /dev/null +++ b/testdata/arrcomp_if4.golden @@ -0,0 +1 @@ +testdata/arrcomp_if4:1:33-34 Unknown variable: y diff --git a/testdata/arrcomp_if4.input b/testdata/arrcomp_if4.input new file mode 100644 index 0000000..2a02baa --- /dev/null +++ b/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]] diff --git a/testdata/arrcomp_if5.golden b/testdata/arrcomp_if5.golden new file mode 100644 index 0000000..1e3ec72 --- /dev/null +++ b/testdata/arrcomp_if5.golden @@ -0,0 +1 @@ +[ ] diff --git a/testdata/arrcomp_if5.input b/testdata/arrcomp_if5.input new file mode 100644 index 0000000..92ea5e6 --- /dev/null +++ b/testdata/arrcomp_if5.input @@ -0,0 +1 @@ +[x for x in [] if error "x"] diff --git a/testdata/arrcomp_if6.golden b/testdata/arrcomp_if6.golden new file mode 100644 index 0000000..57d4eaa --- /dev/null +++ b/testdata/arrcomp_if6.golden @@ -0,0 +1 @@ +RUNTIME ERROR: x diff --git a/testdata/arrcomp_if6.input b/testdata/arrcomp_if6.input new file mode 100644 index 0000000..0b0fdc1 --- /dev/null +++ b/testdata/arrcomp_if6.input @@ -0,0 +1 @@ +[x for x in [1] if error "x"] diff --git a/testdata/arrcomp_if7.golden b/testdata/arrcomp_if7.golden new file mode 100644 index 0000000..ce87f7c --- /dev/null +++ b/testdata/arrcomp_if7.golden @@ -0,0 +1 @@ +RUNTIME ERROR: Unexpected type number, expected boolean diff --git a/testdata/arrcomp_if7.input b/testdata/arrcomp_if7.input new file mode 100644 index 0000000..4fe2992 --- /dev/null +++ b/testdata/arrcomp_if7.input @@ -0,0 +1 @@ +[x for x in [1, 2, 3] if 42]