From b02e1c7bc6baeea04bf43c6ec1a586c37ba47c99 Mon Sep 17 00:00:00 2001
From: Danis Kirasirov <dbgbbu@gmail.com>
Date: Mon, 29 Jan 2024 15:05:23 +0300
Subject: [PATCH] =?UTF-8?q?=D0=A0=D0=B5=D1=84=D0=B0=D0=BA=D1=82=D0=BE?=
 =?UTF-8?q?=D1=80=D0=B8=D0=BD=D0=B3=20=D0=BE=D0=B1=D1=80=D0=B0=D0=B1=D0=BE?=
 =?UTF-8?q?=D1=82=D0=BA=D0=B8=20=D1=84=D1=83=D0=BD=D0=BA=D1=86=D0=B8=D0=B8?=
 =?UTF-8?q?=20len?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/expr/mongo.go      | 96 +++++++++++++++++++++++++-----------------
 pkg/expr/mongo_test.go |  2 +-
 2 files changed, 59 insertions(+), 39 deletions(-)

diff --git a/pkg/expr/mongo.go b/pkg/expr/mongo.go
index 3cdbf045..de14bb57 100644
--- a/pkg/expr/mongo.go
+++ b/pkg/expr/mongo.go
@@ -231,19 +231,12 @@ func (c *compiler) identifier(node ast.Node) string {
 }
 
 func (c *compiler) BinaryNode(node *ast.BinaryNode) interface{} {
+	if result, ok := c.handleLenNode(node); ok {
+		return result
+	}
+
 	switch node.Operator {
 	case "==":
-		op := c.eval(node.Right)
-		lenNode, ok := node.Left.(*ast.BuiltinNode)
-		if ok && lenNode.Name == "len" && len(lenNode.Arguments) == 1 {
-			// Оптимизация для случая len(arr) == 0, т.к. $size не использует индекс
-			if op == 0 {
-				return bson.M{c.identifier(lenNode.Arguments[0]): bson.M{"$eq": bson.A{}}}
-			}
-
-			return bson.M{c.identifier(lenNode.Arguments[0]): bson.M{"$size": op}}
-		}
-
 		return bson.M{c.identifier(node.Left): c.eval(node.Right)}
 
 	case "!=":
@@ -262,36 +255,10 @@ func (c *compiler) BinaryNode(node *ast.BinaryNode) interface{} {
 		return bson.M{c.identifier(node.Left): bson.M{"$nin": c.eval(node.Right)}}
 
 	case "<":
-		op := c.eval(node.Right)
-		lenNode, ok := node.Left.(*ast.BuiltinNode)
-		if ok && lenNode.Name == "len" && len(lenNode.Arguments) == 1 {
-			length, ok := op.(int)
-			if !ok {
-				panic("len must be compared with an integer")
-			}
-			return bson.M{c.identifier(lenNode.Arguments[0]) + "." + strconv.Itoa(length): bson.M{"$exists": false}}
-		}
-
 		return bson.M{c.identifier(node.Left): bson.M{"$lt": c.eval(node.Right)}}
 
 	case ">":
-		op := c.eval(node.Right)
-		lenNode, ok := node.Left.(*ast.BuiltinNode)
-		if ok && lenNode.Name == "len" && len(lenNode.Arguments) == 1 {
-			length, ok := op.(int)
-			if !ok {
-				panic("len must be compared with an integer")
-			}
-
-			// Оптимизация для случая len(arr) > 0
-			if length == 0 {
-				return bson.M{c.identifier(lenNode.Arguments[0]): bson.M{"$exists": true, "$type": "array", "$ne": bson.A{}}}
-			}
-
-			return bson.M{c.identifier(lenNode.Arguments[0]) + "." + strconv.Itoa(length): bson.M{"$exists": true}}
-		}
-
-		return bson.M{c.identifier(node.Left): bson.M{"$gt": op}}
+		return bson.M{c.identifier(node.Left): bson.M{"$gt": c.eval(node.Right)}}
 
 	case "<=":
 		return bson.M{c.identifier(node.Left): bson.M{"$lte": c.eval(node.Right)}}
@@ -688,3 +655,56 @@ func (c *compiler) PairNode(node *ast.PairNode) interface{} {
 	//c.compile(node.Key)
 	//c.compile(node.Value)
 }
+
+func (c *compiler) handleLenNode(node *ast.BinaryNode) (result bson.M, ok bool) {
+	lenNode, ok := node.Left.(*ast.BuiltinNode)
+	if !ok || lenNode.Name != "len" {
+		return nil, false
+	}
+
+	if len(lenNode.Arguments) != 1 {
+		panic("len() expects exactly 1 argument")
+	}
+
+	length, ok := c.eval(node.Right).(int)
+	if !ok || length < 0 {
+		panic("len() can only be compared with non-negative number")
+	}
+
+	switch node.Operator {
+	case "==": // +
+		if length == 0 {
+			return bson.M{c.identifier(lenNode.Arguments[0]): bson.M{"$eq": bson.A{}}}, true
+		}
+		return bson.M{c.identifier(lenNode.Arguments[0]): bson.M{"$size": length}}, true
+
+	case "!=": // +
+		if length == 0 {
+			return bson.M{c.identifier(lenNode.Arguments[0]): bson.M{"$exists": true, "$type": "array", "$ne": bson.A{}}}, true
+		}
+		return bson.M{c.identifier(lenNode.Arguments[0]): bson.M{"$not": bson.M{"$size": length}}}, true
+
+	case ">": // +
+		if length == 0 {
+			return bson.M{c.identifier(lenNode.Arguments[0]): bson.M{"$exists": true, "$type": "array", "$ne": bson.A{}}}, true
+		}
+		return bson.M{c.identifier(lenNode.Arguments[0]) + "." + strconv.Itoa(length): bson.M{"$exists": true}}, true
+
+	case ">=":
+		if length == 0 {
+			return bson.M{c.identifier(lenNode.Arguments[0]): bson.M{"$exists": true, "$type": "array"}}, true
+		}
+		return bson.M{c.identifier(lenNode.Arguments[0]) + "." + strconv.Itoa(length-1): bson.M{"$exists": true}}, true
+
+	case "<":
+		if length == 0 {
+			panic("invalid comparison: len() cannot be less than 0")
+		}
+		return bson.M{c.identifier(lenNode.Arguments[0]) + "." + strconv.Itoa(length-1): bson.M{"$exists": false}}, true
+
+	case "<=":
+		return bson.M{c.identifier(lenNode.Arguments[0]) + "." + strconv.Itoa(length): bson.M{"$exists": false}}, true
+	}
+
+	panic("invalid comparison operator with len()")
+}
diff --git a/pkg/expr/mongo_test.go b/pkg/expr/mongo_test.go
index c8f8d4ba..ae6bee16 100644
--- a/pkg/expr/mongo_test.go
+++ b/pkg/expr/mongo_test.go
@@ -35,7 +35,7 @@ func TestConvertToMongo(t *testing.T) {
 		{"len equal empty", "len(s) == 0", nil, bson.M{"s": bson.M{"$eq": bson.A{}}}, false},
 		{"len gt", "len(s) > 1", nil, bson.M{"s.1": bson.M{"$exists": true}}, false},
 		{"len gt 0", "len(s) > 0", nil, bson.M{"s": bson.M{"$exists": true, "$type": "array", "$ne": bson.A{}}}, false},
-		{"len lt", "len(s) < 1", nil, bson.M{"s.1": bson.M{"$exists": false}}, false},
+		{"len lt", "len(s) < 1", nil, bson.M{"s.0": bson.M{"$exists": false}}, false},
 		{"field#1", "s.test > 3", nil, bson.M{"s.test": bson.M{"$gt": 3}}, false},
 		{"field#2", "s['test'] > 3", nil, bson.M{"s.test": bson.M{"$gt": 3}}, false},
 		{"field#3", "s[test] > 3", nil, bson.M{"s.test": bson.M{"$gt": 3}}, false},
-- 
GitLab