feat: implement with statement

lib/src/compiler/emit.rs

@@ -22,7 +22,7 @@ use walrus::{FunctionId, InstrSeqBuilder, ValType};
 
 use crate::compiler::ir::{
     Expr, ForIn, ForOf, Iterable, MatchAnchor, Of, OfItems, PatternIdx,
-    Quantifier,
+    Quantifier, With,
 };
 use crate::compiler::{
     LiteralId, PatternId, RegexpId, RuleId, RuleInfo, Var, VarStackFrame,
@@ -635,6 +635,10 @@ fn emit_expr(
             }
         },
 
+        Expr::With(with) => {
+            emit_with(ctx, instr, with);
+        }
+
         Expr::FuncCall(fn_call) => {
             // Emit the arguments first.
             for expr in fn_call.args.iter_mut() {
@@ -2134,6 +2138,32 @@ fn emit_for<I, B, C, A>(
     });
 }
 
+/// Emits the code for a `with` statement.
+///
+/// Each `with` statement has a corresponding <identifier> = <expression> pair.
+/// Each pair is stored in the `identifiers` and `expressions` fields of the
+/// `with` statement.
+/// For each pair, the code emitted by this function sets the variable
+/// corresponding to the identifier to the value of the emmited expression.
+/// Those variables are later used in the condition of the `with` statement.
+fn emit_with(
+    ctx: &mut EmitContext,
+    instr: &mut InstrSeqBuilder,
+    with: &mut With,
+) {
+    // Emit the code that sets the variables in the `with` statement.
+    for (idx, &item) in with.identifiers.iter().enumerate() {
+        let expr = &mut with.expressions[idx];
+        set_var(ctx, instr, item, |ctx, instr| {
+            emit_expr(ctx, instr, expr);
+        });
+    }
+
+    // Emit the code that evaluates the condition of the `with` statement.
+    // This condition is a boolean expression that uses the variables set
+    emit_bool_expr(ctx, instr, &mut with.condition)
+}
+
 /// Produces a switch statement by calling a `branch_generator` function
 /// multiple times.
 ///

lib/src/compiler/ir/ast2ir.rs

@@ -24,7 +24,7 @@ use crate::compiler::ir::hex2hir::hex_pattern_hir_from_ast;
 use crate::compiler::ir::{
     Expr, ForIn, ForOf, FuncCall, Iterable, LiteralPattern, Lookup,
     MatchAnchor, Of, OfItems, Pattern, PatternFlagSet, PatternFlags,
-    PatternIdx, PatternInRule, Quantifier, Range, RegexpPattern,
+    PatternIdx, PatternInRule, Quantifier, Range, RegexpPattern, With,
 };
 use crate::compiler::report::ReportBuilder;
 use crate::compiler::{warnings, CompileContext, CompileError};
@@ -520,6 +520,7 @@ pub(in crate::compiler) fn expr_from_ast(
         ast::Expr::Of(of) => of_expr_from_ast(ctx, of),
         ast::Expr::ForOf(for_of) => for_of_expr_from_ast(ctx, for_of),
         ast::Expr::ForIn(for_in) => for_in_expr_from_ast(ctx, for_in),
+        ast::Expr::With(with) => with_expr_from_ast(ctx, with),
         ast::Expr::FuncCall(fn_call) => func_call_from_ast(ctx, fn_call),
 
         ast::Expr::FieldAccess(expr) => {
@@ -1118,6 +1119,48 @@ fn for_in_expr_from_ast(
     })))
 }
 
+fn with_expr_from_ast(
+    ctx: &mut CompileContext,
+    with: &ast::With,
+) -> Result<Expr, CompileError> {
+    // Create stack frame with capacity for the with statement variables
+    let mut stack_frame = ctx.vars.new_frame(with.items.len() as i32);
+    let mut symbols = SymbolTable::new();
+    let mut identifiers = Vec::new();
+    let mut expressions = Vec::new();
+
+    // Iterate over all items in the with statement and create a new variable
+    // for each one. Both identifiers and corresponding expressions are stored
+    // in separate vectors.
+    for item in with.items.iter() {
+        let type_value = expr_from_ast(ctx, &item.expression)?
+            .type_value()
+            .clone_without_value();
+        let var = stack_frame.new_var(type_value.ty());
+
+        identifiers.push(var);
+        expressions.push(expr_from_ast(ctx, &item.expression)?);
+
+        // Insert the variable into the symbol table.
+        symbols.insert(
+            item.identifier.name,
+            Symbol::new(type_value, SymbolKind::Var(var)),
+        );
+    }
+
+    // Put the with variables into scope.
+    ctx.symbol_table.push(Rc::new(symbols));
+
+    let condition = bool_expr_from_ast(ctx, &with.condition)?;
+
+    // Leaving with statement condition's scope. Remove with statement variables.
+    ctx.symbol_table.pop();
+
+    ctx.vars.unwind(&stack_frame);
+
+    Ok(Expr::With(Box::new(With { identifiers, expressions, condition })))
+}
+
 fn iterable_from_ast(
     ctx: &mut CompileContext,
     iter: &ast::Iterable,

lib/src/compiler/ir/dfs.rs

@@ -208,10 +208,18 @@ impl<'a> Iterator for DepthFirstSearch<'a> {
                     }
                     push_quantifier(&for_in.quantifier, &mut self.stack);
                 }
+
                 Expr::Lookup(lookup) => {
                     self.stack.push(Event::Enter(&lookup.index));
                     self.stack.push(Event::Enter(&lookup.primary));
                 }
+
+                Expr::With(with) => {
+                    self.stack.push(Event::Enter(&with.condition));
+                    for expr in with.expressions.iter().rev() {
+                        self.stack.push(Event::Enter(expr))
+                    }
+                }
             }
         }
 

lib/src/compiler/ir/mod.rs

@@ -552,6 +552,9 @@ pub(in crate::compiler) enum Expr {
     /// A `for <quantifier> <vars> in ...` expression. (e.g. `for all i in (1..100) : ( ... )`)
     ForIn(Box<ForIn>),
 
+    /// A `with <identifiers> : ...` expression. (e.g. `with $a, $b : ( ... )`)
+    With(Box<With>),
+
     /// Array or dictionary lookup expression (e.g. `array[1]`, `dict["key"]`)
     Lookup(Box<Lookup>),
 }
@@ -610,6 +613,14 @@ pub(in crate::compiler) struct ForIn {
     pub stack_frame: VarStackFrame,
 }
 
+/// A `with` expression (e.g `with $a, $b : (..)`)
+#[derive(Debug)]
+pub(in crate::compiler) struct With {
+    pub identifiers: Vec<Var>,
+    pub expressions: Vec<Expr>,
+    pub condition: Expr,
+}
+
 /// A quantifier used in `for` and `of` expressions.
 #[derive(Debug)]
 pub(in crate::compiler) enum Quantifier {
@@ -878,7 +889,8 @@ impl Expr {
             | Expr::PatternMatchVar { .. }
             | Expr::Of(_)
             | Expr::ForOf(_)
-            | Expr::ForIn(_) => Type::Bool,
+            | Expr::ForIn(_)
+            | Expr::With(_) => Type::Bool,
 
             Expr::Minus { operand, .. } => match operand.ty() {
                 Type::Integer => Type::Integer,
@@ -947,7 +959,8 @@ impl Expr {
             | Expr::PatternMatchVar { .. }
             | Expr::Of(_)
             | Expr::ForOf(_)
-            | Expr::ForIn(_) => TypeValue::Bool(Value::Unknown),
+            | Expr::ForIn(_)
+            | Expr::With(_) => TypeValue::Bool(Value::Unknown),
 
             Expr::Minus { operand, .. } => match operand.ty() {
                 Type::Integer => TypeValue::Integer(Value::Unknown),
@@ -1193,6 +1206,7 @@ impl Debug for Expr {
                         Expr::Of(_) => writeln!(f, "OF")?,
                         Expr::ForOf(_) => writeln!(f, "FOR_OF")?,
                         Expr::ForIn(_) => writeln!(f, "FOR_IN")?,
+                        Expr::With(_) => writeln!(f, "WITH")?,
                         Expr::Lookup(_) => writeln!(f, "LOOKUP")?,
                         Expr::PatternMatch { pattern, anchor } => writeln!(
                             f,

lib/src/compiler/tests/testdata/errors/134.in

@@ -0,0 +1,4 @@
+rule test {
+	condition:
+	   with foo = "foo" : ( bar )
+}

lib/src/compiler/tests/testdata/errors/134.out

@@ -0,0 +1,6 @@
+error[E009]: unknown identifier `bar`
+ --> line:3:26
+  |
+3 |     with foo = "foo" : ( bar )
+  |                          ^^^ this identifier has not been declared
+  |

lib/src/compiler/tests/testdata/errors/135.in

@@ -0,0 +1,5 @@
+rule test {
+	condition:
+	   with foo = "foo" : ( foo == "foo" )
+       and foo == "foo"
+}

lib/src/compiler/tests/testdata/errors/135.out

@@ -0,0 +1,6 @@
+error[E009]: unknown identifier `foo`
+ --> line:4:12
+  |
+4 |        and foo == "foo"
+  |            ^^^ this identifier has not been declared
+  |

lib/src/compiler/tests/testdata/errors/72.out

@@ -2,5 +2,5 @@ error[E001]: syntax error
  --> line:1:24
   |
 1 | rule test { condition: }
-  |                        ^ expecting expression, `for`, `all`, `none` or `any`, found `}`
+  |                        ^ expecting expression, `for`, `all`, `none`, `any` or `with`, found `}`
   |

lib/src/tests/mod.rs

@@ -514,6 +514,69 @@ fn for_in() {
     );
 }
 
+#[test]
+fn with() {
+    condition_true!(r#"with foo = 1 + 1 : (foo == 2)"#);
+    condition_false!(r#"with foo = 1 + 1 : (foo == 3)"#);
+    condition_true!(r#"with foo = 1 + 1, bar = 2 + 2 : (foo + bar == 6)"#);
+    condition_false!(r#"with foo = 1 + 1, bar = 2 + 2 : (foo + bar == 7)"#);
+
+    #[cfg(feature = "test_proto2-module")]
+    condition_true!(r#"with foo = test_proto2.array_int64[0]: (foo == 1)"#);
+
+    #[cfg(feature = "test_proto2-module")]
+    condition_false!(
+        r#"with foo = test_proto2.array_int64[test_proto2.int64_zero]: (foo == 10)"#
+    );
+
+    #[cfg(feature = "test_proto2-module")]
+    condition_true!(
+        r#"with foo = test_proto2.map_string_struct["foo"].nested_int64_one: (foo == 1)"#
+    );
+
+    #[cfg(feature = "test_proto2-module")]
+    condition_true!(
+        r#"with 
+            bar = test_proto2.array_string[1],
+            baz = test_proto2.array_string[2] : 
+                (
+                    bar == "bar" and baz == "baz"
+                )
+        "#
+    );
+
+    #[cfg(feature = "test_proto2-module")]
+    condition_true!(
+        r#"for any i in (0..1): (
+            with foo = test_proto2.array_int64[i]: (foo == 1)
+        )"#
+    );
+
+    #[cfg(feature = "test_proto2-module")]
+    condition_true!(
+        r#"for all i in (0..0): (
+            with 
+                foo = test_proto2.array_int64[i],
+                bar = test_proto2.array_int64[i + 1] : 
+                    (
+                        foo == 1 and bar == 10
+                    )
+        )"#
+    );
+
+    #[cfg(feature = "test_proto2-module")]
+    condition_false!(
+        r#"for all i in (0..2): (
+            with 
+                foo = test_proto2.array_int64[i],
+                bar = test_proto2.array_int64[i + 1] : 
+                    (
+                        foo == 1 and bar == foo * 10
+                    )
+        )"#
+    );
+}
+
 #[test]
 fn text_patterns() {
     pattern_true!(r#""issi""#, b"mississippi");

parser/src/ast/ascii_tree.rs

@@ -424,9 +424,29 @@ pub(crate) fn expr_ascii_tree(expr: &Expr) -> Tree {
 
             Node(node_title, children)
         }
+        Expr::With(w) => Node(
+            "with <identifiers> : ( <boolean expression> )".to_string(),
+            vec![
+                Node(
+                    "<identifiers>".to_string(),
+                    w.items.iter().flat_map(with_items_ascii_tree).collect(),
+                ),
+                Node(
+                    "<boolean expression>".to_string(),
+                    vec![expr_ascii_tree(&w.condition)],
+                ),
+            ],
+        ),
     }
 }
 
+fn with_items_ascii_tree(item: &WithItems) -> Vec<Tree> {
+    vec![
+        Leaf(vec![format!("{}", item.identifier.name)]),
+        expr_ascii_tree(&item.expression),
+    ]
+}
+
 pub(crate) fn quantifier_ascii_tree(quantifier: &Quantifier) -> Tree {
     match quantifier {
         Quantifier::None { .. } => Leaf(vec!["none".to_string()]),

parser/src/ast/cst2ast.rs

@@ -905,6 +905,7 @@ impl<'src> Builder<'src> {
             }
             Event::Begin(FOR_EXPR) => self.for_expr()?,
             Event::Begin(OF_EXPR) => self.of_expr()?,
+            Event::Begin(WITH_EXPR) => self.with_expr()?,
             Event::Begin(EXPR) => self.pratt_parser(Self::expr, 0)?,
             event => panic!("unexpected {:?}", event),
         };
@@ -1018,6 +1019,44 @@ impl<'src> Builder<'src> {
         Ok(Expr::Of(Box::new(Of { span, quantifier, items, anchor })))
     }
 
+    fn with_expr(&mut self) -> Result<Expr<'src>, Abort> {
+        self.begin(WITH_EXPR)?;
+
+        let mut span = self.expect(WITH_KW)?;
+
+        self.begin(WITH_IDENTIFIERS)?;
+
+        let item = |i: &mut Self| -> Result<WithItems<'src>, Abort> {
+            let identifier = i.identifier()?;
+            let mut span = identifier.span();
+            span = span.combine(&i.expect(EQUAL)?);
+            let expression = i.expr()?;
+            span = span.combine(&expression.span());
+
+            Ok(WithItems { span, identifier, expression })
+        };
+
+        let mut items = vec![item(self)?];
+
+        while let Event::Token { kind: COMMA, .. } = self.peek() {
+            self.expect(COMMA)?;
+            items.push(item(self)?);
+        }
+
+        self.end(WITH_IDENTIFIERS)?;
+
+        self.expect(COLON)?;
+        self.expect(L_PAREN)?;
+
+        let condition = self.boolean_expr()?;
+
+        span = span.combine(&self.expect(R_PAREN)?);
+
+        self.end(WITH_EXPR)?;
+
+        Ok(Expr::With(Box::new(With { span, items, condition })))
+    }
+
     fn quantifier(&mut self) -> Result<Quantifier<'src>, Abort> {
         self.begin(QUANTIFIER)?;