Added possibility to use FOR loop in ternary expression

pkg/compiler/compiler.go

@@ -90,3 +90,13 @@ func (c *FqlCompiler) Compile(query string) (program *runtime.Program, err error
 
 	return program, err
 }
+
+func (c *FqlCompiler) CompileP(query string) *runtime.Program {
+	program, err := c.Compile(query)
+
+	if err != nil {
+		panic(err)
+	}
+
+	return program
+}

pkg/compiler/compiler_test.go

@@ -1602,6 +1602,92 @@ func TestInOperator(t *testing.T) {
 	})
 }
 
+func TestForTernaryExpression(t *testing.T) {
+	Convey("RETURN foo ? TRUE : (FOR i IN 1..5 RETURN i*2)", t, func() {
+		c := compiler.New()
+
+		out1, err := c.CompileP(`
+			LET foo = FALSE
+			RETURN foo ? TRUE : (FOR i IN 1..5 RETURN i*2)
+		`).Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out1), ShouldEqual, `[2,4,6,8,10]`)
+
+		out2, err := c.CompileP(`
+			LET foo = TRUE
+			RETURN foo ? TRUE : (FOR i IN 1..5 RETURN i*2)
+		`).Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out2), ShouldEqual, `true`)
+	})
+
+	Convey("RETURN foo ? (FOR i IN 1..5 RETURN i) : (FOR i IN 1..5 RETURN i*2)", t, func() {
+		c := compiler.New()
+
+		out1, err := c.CompileP(`
+			LET foo = FALSE
+			RETURN foo ? (FOR i IN 1..5 RETURN i) : (FOR i IN 1..5 RETURN i*2)
+		`).Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out1), ShouldEqual, `[2,4,6,8,10]`)
+
+		out2, err := c.CompileP(`
+			LET foo = TRUE
+			RETURN foo ? (FOR i IN 1..5 RETURN i) : (FOR i IN 1..5 RETURN i*2)
+		`).Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out2), ShouldEqual, `[1,2,3,4,5]`)
+	})
+
+	Convey("LET res =  foo ? TRUE : (FOR i IN 1..5 RETURN i*2)", t, func() {
+		c := compiler.New()
+
+		out1, err := c.CompileP(`
+			LET foo = FALSE
+			LET res = foo ? TRUE : (FOR i IN 1..5 RETURN i*2) 
+			RETURN res
+		`).Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out1), ShouldEqual, `[2,4,6,8,10]`)
+
+		out2, err := c.CompileP(`
+			LET foo = TRUE
+			LET res = foo ? TRUE : (FOR i IN 1..5 RETURN i*2)
+			RETURN res
+		`).Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out2), ShouldEqual, `true`)
+	})
+
+	Convey("LET res = foo ? (FOR i IN 1..5 RETURN i) : (FOR i IN 1..5 RETURN i*2)", t, func() {
+		c := compiler.New()
+
+		out1, err := c.CompileP(`
+			LET foo = FALSE
+			LET res = foo ? (FOR i IN 1..5 RETURN i) : (FOR i IN 1..5 RETURN i*2)
+			RETURN res
+		`).Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out1), ShouldEqual, `[2,4,6,8,10]`)
+
+		out2, err := c.CompileP(`
+			LET foo = TRUE
+			LET res = foo ? (FOR i IN 1..5 RETURN i) : (FOR i IN 1..5 RETURN i*2)
+			RETURN res
+		`).Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out2), ShouldEqual, `[1,2,3,4,5]`)
+	})
+}
+
 //func TestHtml(t *testing.T) {
 //	Convey("Should load a document", t, func() {
 //		c := compiler.New()

pkg/compiler/visitor.go

@@ -117,7 +117,13 @@ func (v *visitor) doVisitReturnExpression(ctx *fql.ReturnExpressionContext, scop
 		}
 
 		exp = out
-	} else {
+
+		return expressions.NewReturnExpression(v.getSourceMap(ctx), exp)
+	}
+
+	forIn := ctx.ForExpression()
+
+	if forIn != nil {
 		out, err := v.doVisitForExpression(ctx.ForExpression().(*fql.ForExpressionContext), scope.Fork())
 
 		if err != nil {
@@ -125,9 +131,23 @@ func (v *visitor) doVisitReturnExpression(ctx *fql.ReturnExpressionContext, scop
 		}
 
 		exp = out
+
+		return expressions.NewReturnExpression(v.getSourceMap(ctx), exp)
 	}
 
-	return expressions.NewReturnExpression(v.getSourceMap(ctx), exp)
+	forInTernary := ctx.ForTernaryExpression()
+
+	if forInTernary != nil {
+		out, err := v.doVisitForTernaryExpression(forInTernary.(*fql.ForTernaryExpressionContext), scope)
+
+		if err != nil {
+			return nil, err
+		}
+
+		return expressions.NewReturnExpression(v.getSourceMap(ctx), out)
+	}
+
+	return nil, ErrNotImplemented
 }
 
 func (v *visitor) doVisitForExpression(ctx *fql.ForExpressionContext, scope *scope) (core.Expression, error) {
@@ -609,7 +629,9 @@ func (v *visitor) doVisitVariableDeclaration(ctx *fql.VariableDeclarationContext
 
 	if exp != nil {
 		init, err = v.doVisitExpression(ctx.Expression().(*fql.ExpressionContext), scope)
-	} else {
+	}
+
+	if init == nil && err == nil {
 		forIn := ctx.ForExpression()
 
 		if forIn != nil {
@@ -617,6 +639,14 @@ func (v *visitor) doVisitVariableDeclaration(ctx *fql.VariableDeclarationContext
 		}
 	}
 
+	if init == nil && err == nil {
+		forTer := ctx.ForTernaryExpression()
+
+		if forTer != nil {
+			init, err = v.doVisitForTernaryExpression(forTer.(*fql.ForTernaryExpressionContext), scope)
+		}
+	}
+
 	if err != nil {
 		return nil, err
 	}
@@ -661,16 +691,22 @@ func (v *visitor) doVisitChildren(node antlr.RuleNode, scope *scope) ([]core.Exp
 		return make([]core.Expression, 0, 0), nil
 	}
 
-	result := make([]core.Expression, len(children))
+	result := make([]core.Expression, 0, len(children))
+
+	for _, child := range children {
+		_, ok := child.(antlr.TerminalNode)
+
+		if ok {
+			continue
+		}
 
-	for idx, child := range children {
 		out, err := v.visit(child, scope)
 
 		if err != nil {
 			return nil, err
 		}
 
-		result[idx] = out
+		result = append(result, out)
 	}
 
 	return result, nil
@@ -794,24 +830,10 @@ func (v *visitor) doVisitExpression(ctx *fql.ExpressionContext, scope *scope) (c
 			return nil, err
 		}
 
-		var test core.Expression
-		var consequent core.Expression
-		var alternate core.Expression
-
-		if len(exps) == 3 {
-			test = exps[0]
-			consequent = exps[1]
-			alternate = exps[2]
-		} else {
-			test = exps[0]
-			alternate = exps[1]
-		}
-
-		return expressions.NewConditionExpression(
+		return v.createTernaryOperator(
 			v.getSourceMap(ctx),
-			test,
-			consequent,
-			alternate,
+			exps,
+			scope,
 		)
 	}
 
@@ -860,14 +882,6 @@ func (v *visitor) doVisitExpression(ctx *fql.ExpressionContext, scope *scope) (c
 		return v.doVisitRangeOperator(rangeOp.(*fql.RangeOperatorContext), scope)
 	}
 
-	seq := ctx.ExpressionSequence()
-
-	if seq != nil {
-		// seq := seq.(*fql.ExpressionSequenceContext)
-
-		return nil, core.Error(ErrNotImplemented, "expression sequence")
-	}
-
 	// TODO: Complete it
 	return nil, ErrNotImplemented
 }
@@ -981,6 +995,42 @@ func (v *visitor) visit(node antlr.Tree, scope *scope) (core.Expression, error)
 	return out, err
 }
 
+func (v *visitor) doVisitForTernaryExpression(ctx *fql.ForTernaryExpressionContext, scope *scope) (*expressions.ConditionExpression, error) {
+	exps, err := v.doVisitChildren(ctx, scope)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return v.createTernaryOperator(
+		v.getSourceMap(ctx),
+		exps,
+		scope,
+	)
+}
+
+func (v *visitor) createTernaryOperator(src core.SourceMap, exps []core.Expression, scope *scope) (*expressions.ConditionExpression, error) {
+	var test core.Expression
+	var consequent core.Expression
+	var alternate core.Expression
+
+	if len(exps) == 3 {
+		test = exps[0]
+		consequent = exps[1]
+		alternate = exps[2]
+	} else {
+		test = exps[0]
+		alternate = exps[1]
+	}
+
+	return expressions.NewConditionExpression(
+		src,
+		test,
+		consequent,
+		alternate,
+	)
+}
+
 func (v *visitor) unexpectedToken(node antlr.Tree) error {
 	name := "undefined"
 	ctx, ok := node.(antlr.RuleContext)