Merge branch 'feature/eval-func-enhance' into 'develop'

[CN-1617] feat: the EVAL function supports bitwise operations

See merge request galaxy/platform/groot-stream!26

6 files changed, 82 insertions, 10 deletions

diff --git a/groot-core/src/main/antlr4/com/geedgenetworks/core/expressions/parser/EvalExpression.g4 b/groot-core/src/main/antlr4/com/geedgenetworks/core/expressions/parser/EvalExpression.g4
index 2df35f1..ee7cd5a 100644
--- a/groot-core/src/main/antlr4/com/geedgenetworks/core/expressions/parser/EvalExpression.g4
+++ b/groot-core/src/main/antlr4/com/geedgenetworks/core/expressions/parser/EvalExpression.g4
@@ -67,6 +67,7 @@ comparisonOperator
 
 arithmeticOperator
     : PLUS | MINUS | ASTERISK | SLASH | PERCENT
+    | LEFT_SHIFT | RIGHT_SHIFT | BITWISE_AND | BITWISE_OR
     ;
 
 predicateOperator
@@ -87,9 +88,13 @@ MINUS: '-';
 ASTERISK: '*';
 SLASH: '/';
 PERCENT: '%';
+LEFT_SHIFT: '<<';
+RIGHT_SHIFT: '>>';
+BITWISE_OR: '|';
+BITWISE_AND: '&';
 
-AND: '&' | '&&';
-OR: '|' | '||';
+AND: '&&';
+OR: '||';
 NOT: '!';
 
 TRUE: 'TRUE';
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/expressions/ArithmeticOperator.java b/groot-core/src/main/java/com/geedgenetworks/core/expressions/ArithmeticOperator.java
index de9d99e..36846b3 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/expressions/ArithmeticOperator.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/expressions/ArithmeticOperator.java
@@ -6,5 +6,9 @@ public enum ArithmeticOperator {
     MINUS,
     ASTERISK,
     SLASH,
-    PERCENT
+    PERCENT,
+    LEFT_SHIFT,
+    RIGHT_SHIFT,
+    BITWISE_AND,
+    BITWISE_OR
 }
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/expressions/EvalExpressionResolver.java b/groot-core/src/main/java/com/geedgenetworks/core/expressions/EvalExpressionResolver.java
index 3b18315..95015ae 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/expressions/EvalExpressionResolver.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/expressions/EvalExpressionResolver.java
@@ -39,7 +39,7 @@ public class EvalExpressionResolver extends EvalExpressionBaseVisitor<Object> {
                         message +=
                                 "Expect: "
                                         + e.getExpectedTokens()
-                                                .toString(recognizer.getVocabulary());
+                                        .toString(recognizer.getVocabulary());
                         message += "\n";
                         message += StringUtils.split(expr, "\n")[currentToken.getLine() - 1];
                         message += "\n";
@@ -187,7 +187,7 @@ public class EvalExpressionResolver extends EvalExpressionBaseVisitor<Object> {
                 (short)
                         (factor
                                 * Short.parseShort(
-                                        ctx.SMALLINT_LITERAL().getText().replace("S", "")));
+                                ctx.SMALLINT_LITERAL().getText().replace("S", "")));
         return new ValueExpression.NumberConstantExpression(value);
     }
 
@@ -243,9 +243,19 @@ public class EvalExpressionResolver extends EvalExpressionBaseVisitor<Object> {
         } else if (ctx.ASTERISK() != null) {
             return ArithmeticOperator.ASTERISK;
         } else if (ctx.SLASH() != null) {
-            return ArithmeticOperator.ASTERISK;
-        } else {
+            return ArithmeticOperator.SLASH;
+        } else if (ctx.PERCENT() != null) {
             return ArithmeticOperator.PERCENT;
+        } else if (ctx.LEFT_SHIFT() != null) {
+            return ArithmeticOperator.LEFT_SHIFT;
+        } else if (ctx.RIGHT_SHIFT() != null) {
+            return ArithmeticOperator.RIGHT_SHIFT;
+        } else if (ctx.BITWISE_AND() != null) {
+            return ArithmeticOperator.BITWISE_AND;
+        } else if (ctx.BITWISE_OR() != null) {
+            return ArithmeticOperator.BITWISE_OR;
+        } else {
+            throw new IllegalArgumentException("Unsupported operate.");
         }
     }
 
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/expressions/codegen/EvalCodeGenerator.java b/groot-core/src/main/java/com/geedgenetworks/core/expressions/codegen/EvalCodeGenerator.java
index 797be04..bb6b826 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/expressions/codegen/EvalCodeGenerator.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/expressions/codegen/EvalCodeGenerator.java
@@ -97,6 +97,18 @@ public class EvalCodeGenerator implements ExpressionVisitor<String> {
             case SLASH:
                 code += String.format("%s / %s", leftTerm, rightTerm);
                 break;
+            case LEFT_SHIFT:
+                code += String.format("((java.lang.Integer)%s).intValue() << ((java.lang.Integer)%s).intValue()", leftTerm, rightTerm);
+                break;
+            case RIGHT_SHIFT:
+                code += String.format("((java.lang.Integer)%s).intValue() >> ((java.lang.Integer)%s).intValue()", leftTerm, rightTerm);
+                break;
+            case BITWISE_AND:
+                code += String.format("((java.lang.Integer)%s).intValue() & ((java.lang.Integer)%s).intValue()", leftTerm, rightTerm);
+                break;
+            case BITWISE_OR:
+                code += String.format("((java.lang.Integer)%s).intValue() | ((java.lang.Integer)%s).intValue()", leftTerm, rightTerm);
+                break;
         }
         ctx.addEvalCodeStatements(code + ";");
         return name;
diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/test/eval/EvalExpressionResolverTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/test/eval/EvalExpressionResolverTest.java
index 54d67e9..ccfa03e 100644
--- a/groot-core/src/test/java/com/geedgenetworks/core/udf/test/eval/EvalExpressionResolverTest.java
+++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/test/eval/EvalExpressionResolverTest.java
@@ -60,10 +60,10 @@ class EvalExpressionResolverTest {
                         new ValueExpression.ArithmeticExpression(ArithmeticOperator.PLUS,
                                 new ValueExpression.StringConstantExpression("S2C:"),
                                 new ValueExpression.FieldRefExpression("common_s2c_pkt_num")))
-                );
+        );
         assertEquals(
                 resolver.resolve("common_c2s_pkt_num > common_s2c_pkt_num " +
-                        "& common_c2s_byte_num > common_s2c_byte_num ? common_client_port : common_server_port"),
+                        "&& common_c2s_byte_num > common_s2c_byte_num ? common_client_port : common_server_port"),
                 new ValueExpression.TernaryExpression(
                         new BooleanExpression.PredicateExpression(PredicateOperator.AND,
                                 new BooleanExpression.ComparisonExpression(ComparisonOperator.GT,
@@ -75,6 +75,14 @@ class EvalExpressionResolverTest {
                         new ValueExpression.FieldRefExpression("common_client_port"),
                         new ValueExpression.FieldRefExpression("common_server_port"))
         );
+        assertEquals(
+                resolver.resolve("common_c2s_pkt_num & 2"),
+                new ValueExpression.ArithmeticExpression(
+                        ArithmeticOperator.BITWISE_AND,
+                        new ValueExpression.FieldRefExpression("common_c2s_pkt_num"),
+                        new ValueExpression.NumberConstantExpression(2)
+                )
+        );
     }
 
 
diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/test/eval/codegen/CalcCodeConvertorTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/test/eval/codegen/CalcCodeConvertorTest.java
index ded872e..8572db2 100644
--- a/groot-core/src/test/java/com/geedgenetworks/core/udf/test/eval/codegen/CalcCodeConvertorTest.java
+++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/test/eval/codegen/CalcCodeConvertorTest.java
@@ -8,6 +8,7 @@ import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 
 import java.math.BigDecimal;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -76,13 +77,45 @@ class CalcCodeConvertorTest {
                 commonClientPort > commonServerPort ? "C2S:" + commonC2sPktNum : "S2C:" + commonS2cPktNum);
         assertEquals(
                 CalcCodeConvertor.convert("common_c2s_pkt_num > common_s2c_pkt_num " +
-                        "& common_c2s_byte_num > common_s2c_byte_num ? common_client_port : common_server_port").eval(origin),
+                        "&& common_c2s_byte_num > common_s2c_byte_num ? common_client_port : common_server_port").eval(origin),
                 commonC2sPktNum > commonS2cPktNum
                         && commonC2sByteNum > commonS2cByteNum ? commonClientPort : commonServerPort
         );
     }
 
     @Test
+    public void testArithmeticExpression() throws Exception {
+        final Map<String, Object> map = Collections.emptyMap();
+        assertEquals(
+                1, CalcCodeConvertor.convert("5 & 3").eval(map));
+        assertEquals(
+                7, CalcCodeConvertor.convert("5 | 3").eval(map));
+        assertEquals(
+                20, CalcCodeConvertor.convert("5 << 2").eval(map));
+        assertEquals(
+                1, CalcCodeConvertor.convert("5 >> 2").eval(map));
+        assertEquals(
+                8, CalcCodeConvertor.convert("5 + 3").eval(map));
+        assertEquals(
+                2, CalcCodeConvertor.convert("5 - 3").eval(map));
+        assertEquals(
+                15, CalcCodeConvertor.convert("5 * 3").eval(map));
+        assertEquals(
+                2, CalcCodeConvertor.convert("5 / 2").eval(map));
+        assertEquals(
+                2, CalcCodeConvertor.convert("5 % 3").eval(map));
+
+
+        final Map<String, Object> map1 = new HashMap<>();
+        map1.put("flags", 8);
+
+        assertEquals("B",
+                CalcCodeConvertor.convert(" flags & 8 == 9 ? 'A' : 'B'").eval(map1));
+        assertEquals("A",
+                CalcCodeConvertor.convert(" 5 | flags == 13 ? 'A' : 'B'").eval(map1));
+    }
+
+    @Test
     public void testError() {
         assertThrows(SyntaxErrorException.class,
                 () -> CalcCodeConvertor.convert("common_c2s_pkt_num > common_s2c_pkt_num"));