Fixed inability to use dynamic values in LIMIT clause (#706)

pkg/compiler/compiler_limit_test.go

@@ -93,4 +93,81 @@ func TestForLimit(t *testing.T) {
 		So(err, ShouldBeNil)
 		So(string(out), ShouldEqual, `[1,4]`)
 	})
+
+	Convey("Should be able to use variable", t, func() {
+		c := compiler.New()
+
+		p, err := c.Compile(`
+			LET li = 2
+			FOR i IN [ 1,2,3,4,5,6,7,8 ]
+				LIMIT li
+				RETURN i
+		`)
+
+		So(err, ShouldBeNil)
+
+		out, err := p.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[1,2]`)
+	})
+
+	Convey("Should be able to use function call", t, func() {
+		c := compiler.New()
+		c.RegisterFunction("TEST", func(ctx context.Context, args ...core.Value) (core.Value, error) {
+			return values.NewInt(2), nil
+		})
+
+		p, err := c.Compile(`
+			FOR i IN [ 1,2,3,4,5,6,7,8 ]
+				LIMIT TEST()
+				RETURN i
+		`)
+
+		So(err, ShouldBeNil)
+
+		out, err := p.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[1,2]`)
+	})
+
+	Convey("Should be able to use member expression (object)", t, func() {
+		c := compiler.New()
+
+		p, err := c.Compile(`
+			LET o = {
+				limit: 2
+			}
+			FOR i IN [ 1,2,3,4,5,6,7,8 ]
+				LIMIT o.limit
+				RETURN i
+		`)
+
+		So(err, ShouldBeNil)
+
+		out, err := p.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[1,2]`)
+	})
+
+	Convey("Should be able to use member expression (array)", t, func() {
+		c := compiler.New()
+
+		p, err := c.Compile(`
+			LET o = [1,2]
+
+			FOR i IN [ 1,2,3,4,5,6,7,8 ]
+				LIMIT o[1]
+				RETURN i
+		`)
+
+		So(err, ShouldBeNil)
+
+		out, err := p.Run(context.Background())
+
+		So(err, ShouldBeNil)
+		So(string(out), ShouldEqual, `[1,2]`)
+	})
 }

pkg/compiler/visitor.go

@@ -422,20 +422,25 @@ func (v *visitor) visitLimitClause(c fql.ILimitClauseContext, scope *scope) (cor
 
 func (v *visitor) visitLimitClauseValue(c fql.ILimitClauseValueContext, scope *scope) (core.Expression, error) {
 	ctx := c.(*fql.LimitClauseValueContext)
-	literalCtx := ctx.IntegerLiteral()
 
-	if literalCtx != nil {
+	if literalCtx := ctx.IntegerLiteral(); literalCtx != nil {
 		i, err := strconv.Atoi(literalCtx.GetText())
 		if err != nil {
 			return nil, err
 		}
 
 		return literals.NewIntLiteral(i), nil
+	} else if paramCtx := ctx.Param(); paramCtx != nil {
+		return v.visitParam(paramCtx.(fql.IParamContext), scope)
+	} else if variableCtx := ctx.Variable(); variableCtx != nil {
+		return v.visitVariable(variableCtx, scope)
+	} else if funcCtx := ctx.FunctionCallExpression(); funcCtx != nil {
+		return v.visitFunctionCallExpression(funcCtx, scope)
+	} else if memCtx := ctx.MemberExpression(); memCtx != nil {
+		return v.visitMemberExpression(memCtx, scope)
 	}
 
-	paramCtx := ctx.Param()
-
-	return v.visitParam(paramCtx.(fql.IParamContext), scope)
+	return nil, v.unexpectedToken(ctx)
 }
 
 func (v *visitor) visitFilterClause(c fql.IFilterClauseContext, scope *scope) (core.Expression, error) {

pkg/parser/antlr/FqlParser.g4

@@ -96,6 +96,9 @@ limitClause
 limitClauseValue
     : IntegerLiteral
     | param
+    | variable
+    | functionCallExpression
+    | memberExpression
     ;
 
 sortClause