perf: group multiple matches expressions inside or conditions into a RegexSet (VirusTotal#657)

When writing complex YARA rules, it is incredibly common to evaluate a single variable or structure field access (e.g., vt.net.domain.raw) against dozens of distinct regular expressions within an or condition. Previously, YARA-X evaluated each of these match operations sequentially.

This PR introduces an advanced compile-time and runtime optimization that detects identical match targets within or conditions, groups their regular expressions into a single multi-pattern RegexSet, and evaluates the entire set simultaneously in a single pass.

Author: plusvic

lib/src/compiler/context.rs

@@ -10,9 +10,10 @@ use yara_x_parser::ast::{Ident, WithSpan};
 use crate::compiler::errors::{CompileError, UnknownPattern};
 use crate::compiler::ir::{IR, PatternIdx};
 use crate::compiler::report::ReportBuilder;
-use crate::compiler::{Warnings, ir};
+use crate::compiler::{RegexId, RegexSetId, Warnings, ir};
 use crate::errors::{UnknownField, UnknownIdentifier};
 use crate::modules::BUILTIN_MODULES;
+use crate::string_pool::StringPool;
 use crate::symbols::{StackedSymbolTable, Symbol, SymbolLookup};
 use crate::types::Type;
 use crate::wasm;
@@ -65,6 +66,12 @@ pub(crate) struct CompileContext<'a, 'src> {
     /// Tracks the product of iteration counts of nested loops.
     /// Used to detect loops that may iterate an excessive number of times.
     pub loop_iteration_multiplier: i64,
+
+    /// Grouped RegexSets constructed during IR creation for or-expressions.
+    pub regex_sets: &'a mut rustc_hash::FxHashMap<RegexSetId, Vec<RegexId>>,
+
+    /// Pool for regular expressions.
+    pub regex_pool: &'a mut StringPool<RegexId>,
 }
 
 impl<'src> CompileContext<'_, 'src> {

lib/src/compiler/emit.rs

@@ -27,7 +27,7 @@ use crate::compiler::ir::{
 };
 use crate::compiler::{
     FieldAccess, ForVars, LiteralId, OfExprTuple, OfPatternSet, PatternId,
-    RegexpId, RuleId, RuleInfo, Var,
+    RegexId, RuleId, RuleInfo, Var,
 };
 use crate::scanner::RuntimeObjectHandle;
 use crate::string_pool::{BStringPool, StringPool};
@@ -205,7 +205,7 @@ pub(crate) struct EmitContext<'a> {
     pub wasm_exports: &'a FxHashMap<String, FunctionId>,
 
     /// Pool with regular expressions used in rule conditions.
-    pub regexp_pool: &'a mut StringPool<RegexpId>,
+    pub regex_pool: &'a mut StringPool<RegexId>,
 
     /// Pool with literal strings used in the rules.
     pub lit_pool: &'a mut BStringPool<LiteralId>,
@@ -313,7 +313,7 @@ fn emit_expr(
                     .i64_const(RuntimeString::Literal(literal_id).into_wasm());
             }
             TypeValue::Regexp(Some(regexp)) => {
-                let re_id = ctx.regexp_pool.get_or_intern(regexp.as_str());
+                let re_id = ctx.regex_pool.get_or_intern(regexp.as_str());
 
                 instr.i32_const(re_id.into());
             }
@@ -640,6 +640,14 @@ fn emit_expr(
                 .call(ctx.function_id(wasm::export__str_matches.mangled_name));
         }
 
+        Expr::MatchesMany { lhs, regex_set } => {
+            emit_expr(ctx, ir, *lhs, instr);
+            instr.i32_const((*regex_set).into());
+            instr.call(ctx.function_id(
+                wasm::export__str_matches_regex_set.mangled_name,
+            ));
+        }
+
         Expr::Lookup(lookup) => {
             // Emit code for the primary expression (array or map) that is
             // being indexed.

lib/src/compiler/ir/ast2ir.rs

@@ -3,13 +3,14 @@
 use std::borrow::Borrow;
 use std::cell::RefCell;
 use std::collections::BTreeSet;
+use std::hash::{Hash, Hasher};
 use std::iter;
 use std::ops::RangeInclusive;
 use std::rc::Rc;
 
 use bstr::{BString, ByteSlice};
 use itertools::Itertools;
-
+use rustc_hash::{FxHashMap, FxHashSet, FxHasher};
 use yara_x_parser::Span;
 use yara_x_parser::ast;
 use yara_x_parser::ast::WithSpan;
@@ -26,12 +27,12 @@ use crate::compiler::errors::{
 use crate::compiler::ir::hex2hir::hex_pattern_hir_from_ast;
 use crate::compiler::ir::{
     Error, Expr, ExprId, Iterable, LiteralPattern, MatchAnchor, Pattern,
-    PatternFlags, PatternInRule, Quantifier, Range, RegexpPattern,
+    PatternFlags, PatternInRule, Quantifier, Range, RegexpPattern, dfs,
 };
 use crate::compiler::report::{Level, ReportBuilder};
 use crate::compiler::{
     CompileContext, CompileError, FilesizeBounds, ForVars, PatternIdx,
-    TextPatternAsHex, warnings,
+    RegexId, RegexSetId, TextPatternAsHex, warnings,
 };
 use crate::errors::CustomError;
 use crate::errors::{MethodNotAllowedInWith, PotentiallySlowLoop};
@@ -2179,40 +2180,23 @@ macro_rules! gen_n_ary_operation {
 
             // Make sure that all operands have one of the accepted types.
             for (hir, ast) in iter::zip(operands_hir.iter(), expr.operands()) {
-                check_type(ctx, *hir, ast.span(), accepted_types)?;
                 if let Some(check_fn) = check_fn {
                     check_fn(ctx, *hir, ast.span())?;
                 }
             }
 
-            // Iterate the operands in pairs (first, second), (second, third),
-            // (third, fourth), etc.
-            for ((lhs_hir, rhs_ast), (rhs_hir, lhs_ast)) in
+            for ((lhs_hir, lhs_ast), (rhs_hir, rhs_ast)) in
                 iter::zip(operands_hir.iter(), expr.operands()).tuple_windows()
             {
-                let lhs_ty = ctx.ir.get(*lhs_hir).ty();
-                let rhs_ty = ctx.ir.get(*rhs_hir).ty();
-
-                let types_are_compatible = {
-                    // If the types are the same, they are compatible.
-                    (lhs_ty == rhs_ty) ||
-                        // If the list of compatible types contains the pair
-                        // (lhs_ty, rhs_ty) or (rhs_ty, lhs_ty), they are
-                        // compatible.
-                        compatible_types.contains(&(lhs_ty, rhs_ty))
-                            || compatible_types.contains(&(rhs_ty, lhs_ty))
-
-                };
-
-                if !types_are_compatible {
-                    return Err(MismatchingTypes::build(
-                            ctx.report_builder,
-                            lhs_ty.to_string(),
-                            rhs_ty.to_string(),
-                            ctx.report_builder.span_to_code_loc(expr.first().span().combine(&lhs_ast.span())),
-                            ctx.report_builder.span_to_code_loc(rhs_ast.span()),
-                    ));
-                }
+                check_operands(
+                    ctx,
+                    *lhs_hir,
+                    *rhs_hir,
+                    expr.first().span().combine(&lhs_ast.span()),
+                    rhs_ast.span(),
+                    accepted_types,
+                    compatible_types,
+                )?;
             }
 
             ctx.ir.$variant(operands_hir).map_err(|err| {
@@ -2275,28 +2259,123 @@ gen_n_ary_operation!(
     })
 );
 
-gen_n_ary_operation!(
-    or_expr_from_ast,
-    or,
-    // Boolean operations accept integer, float and string operands.
-    // If operands are not boolean they are casted to boolean.
-    Type::Bool | Type::Integer | Type::Float | Type::String,
-    // All operand types can be mixed in a boolean operation, as they
-    // are casted to boolean anyways.
-    &[
+fn or_expr_from_ast(
+    ctx: &mut CompileContext,
+    expr: &ast::NAryExpr,
+) -> Result<ExprId, CompileError> {
+    let span = expr.span();
+    let accepted_types =
+        &[Type::Bool, Type::Integer, Type::Float, Type::String];
+    let compatible_types = &[
         (Type::Integer, Type::Bool),
         (Type::Integer, Type::Float),
         (Type::Integer, Type::String),
         (Type::String, Type::Bool),
         (Type::String, Type::Float),
-        (Type::Float, Type::Bool)
-    ],
-    Some(|ctx, operand, span| {
-        let ty = ctx.ir.get(operand).ty();
-        warn_if_not_bool(ctx, ty, span);
-        Ok(())
+        (Type::Float, Type::Bool),
+    ];
+
+    // Validate operand types and emit standard warnings. Ensure all operands
+    // in the `or` expression conform to boolean-compatible types, checking
+    // for potential mismatches across adjacent pairs.
+    let or_operands: Vec<ExprId> = expr
+        .operands()
+        .map(|expr| expr_from_ast(ctx, expr))
+        .collect::<Result<Vec<ExprId>, CompileError>>()?;
+
+    for (hir, ast) in iter::zip(or_operands.iter(), expr.operands()) {
+        let ty = ctx.ir.get(*hir).ty();
+        warn_if_not_bool(ctx, ty, ast.span());
+    }
+
+    for ((lhs_hir, lhs_ast), (rhs_hir, rhs_ast)) in
+        iter::zip(or_operands.iter(), expr.operands()).tuple_windows()
+    {
+        check_operands(
+            ctx,
+            *lhs_hir,
+            *rhs_hir,
+            expr.first().span().combine(&lhs_ast.span()),
+            rhs_ast.span(),
+            accepted_types,
+            compatible_types,
+        )?;
+    }
+
+    // Group `matches` expressions by their left-side operand. All the `matches`
+    // expressions sharing the same left-side operand will be aggregated together
+    // into a MatchMany operation that matches the left-side operand against
+    // a set of regular expressions.
+    let mut matches_by_lhs: FxHashMap<u64, Vec<(usize, ExprId, RegexId)>> =
+        FxHashMap::default();
+
+    for (i, &op) in or_operands.iter().enumerate() {
+        if let Expr::Matches { lhs, rhs } = ctx.ir.get(op)
+            && let Expr::Const(TypeValue::Regexp(Some(re))) = ctx.ir.get(*rhs)
+        {
+            let mut hasher = FxHasher::default();
+            for evt in ctx.ir.dfs_iter(*lhs) {
+                if let dfs::Event::Enter((_, expr, _)) = evt {
+                    Hash::hash(expr, &mut hasher);
+                }
+            }
+            let lhs_expr_hash = hasher.finish();
+            let re_id = ctx.regex_pool.get_or_intern(re.as_str());
+            matches_by_lhs
+                .entry(lhs_expr_hash)
+                .or_default()
+                .push((i, *lhs, re_id));
+        }
+    }
+
+    // Collapse grouped regular expressions into `MatchesMany`. For any target
+    // expression associated with two or more regular expressions, construct a
+    // unified `RegexSet`, replace the individual `matches` nodes with a single
+    // `MatchesMany` expression, and record the collapsed indices.
+    let mut final_operands = Vec::new();
+    let mut collapsed_indices = FxHashSet::default();
+
+    for (_, group) in matches_by_lhs {
+        if group.len() >= 2 {
+            let set_id = RegexSetId::from(ctx.regex_sets.len() as i32);
+            let re_ids: Vec<_> =
+                group.iter().map(|&(_, _, re_id)| re_id).collect();
+            ctx.regex_sets.insert(set_id, re_ids);
+
+            let first_lhs = group[0].1;
+            let multimatch = ctx.ir.matches_regex_set(first_lhs, set_id);
+
+            final_operands.push(multimatch);
+
+            for (i, _, _) in group {
+                collapsed_indices.insert(i);
+            }
+        }
+    }
+
+    // Assemble final operands and construct the `or` expression. Preserve all
+    // non-collapsed operands in their original relative order. If all operands
+    // collapse into a single expression, return it directly without a wrapping
+    // `or` node.
+    for (i, &op) in or_operands.iter().enumerate() {
+        if !collapsed_indices.contains(&i) {
+            final_operands.push(op);
+        }
+    }
+
+    if final_operands.len() == 1 {
+        return Ok(final_operands[0]);
+    }
+
+    ctx.ir.or(final_operands).map_err(|err| match err {
+        Error::NumberOutOfRange => NumberOutOfRange::build(
+            ctx.report_builder,
+            i64::MIN,
+            i64::MAX,
+            ctx.report_builder.span_to_code_loc(span),
+        ),
     })
-);
+}
 
 gen_unary_op!(minus_expr_from_ast, minus, Type::Integer | Type::Float, None);
 

lib/src/compiler/ir/dfs.rs

@@ -303,6 +303,10 @@ pub(super) fn dfs_common(
             stack.push(Event::Enter((*lhs, EventContext::None)));
         }
 
+        Expr::MatchesMany { lhs, .. } => {
+            stack.push(Event::Enter((*lhs, EventContext::None)));
+        }
+
         Expr::PatternMatch { anchor, .. }
         | Expr::PatternMatchVar { anchor, .. } => {
             push_anchor(anchor, stack);

lib/src/compiler/ir/hex2hir.rs

@@ -227,6 +227,9 @@ mod tests {
         let mut warnings = Warnings::default();
         let mut rule_patterns = vec![];
 
+        let mut regex_sets = rustc_hash::FxHashMap::default();
+        let mut regex_pool = crate::string_pool::StringPool::new();
+
         let mut ctx = CompileContext {
             ir: &mut ir,
             relaxed_re_syntax: false,
@@ -240,6 +243,8 @@ mod tests {
             vars: VarStack::new(),
             for_of_depth: 0,
             loop_iteration_multiplier: 1,
+            regex_sets: &mut regex_sets,
+            regex_pool: &mut regex_pool,
         };
 
         let mut pattern = HexPattern::new("test_ident");

lib/src/compiler/ir/mod.rs

@@ -51,7 +51,7 @@ use crate::compiler::ir::dfs::{
     DFSIter, DFSWithScopeIter, Event, EventContext, dfs_common,
 };
 
-use crate::compiler::FilesizeBounds;
+use crate::compiler::{FilesizeBounds, RegexSetId};
 use crate::re;
 use crate::symbols::Symbol;
 use crate::types::Value::Const;
@@ -1556,6 +1556,20 @@ impl IR {
         expr_id
     }
 
+    /// Creates a new [`Expr::MatchesMany`].
+    pub fn matches_regex_set(
+        &mut self,
+        lhs: ExprId,
+        regex_set: crate::compiler::RegexSetId,
+    ) -> ExprId {
+        let expr_id = ExprId::from(self.nodes.len());
+        self.parents[lhs.0 as usize] = expr_id;
+        self.parents.push(ExprId::none());
+        self.nodes.push(Expr::MatchesMany { lhs, regex_set });
+        debug_assert_eq!(self.parents.len(), self.nodes.len());
+        expr_id
+    }
+
     /// Creates a new [`Expr::PatternMatch`]
     pub fn pattern_match(
         &mut self,
@@ -2000,6 +2014,7 @@ impl Debug for IR {
                         Expr::IEndsWith { .. } => writeln!(f, "IENDS_WITH -- hash: {expr_hash:#08x}")?,
                         Expr::IEquals { .. } => writeln!(f, "IEQUALS -- hash: {expr_hash:#08x}")?,
                         Expr::Matches { .. } => writeln!(f, "MATCHES -- hash: {expr_hash:#08x}")?,
+                        Expr::MatchesMany { regex_set, .. } => writeln!(f, "MATCHES_MANY RegexSetId({}) -- hash: {expr_hash:#08x}", usize::from(*regex_set))?,
                         Expr::Defined { .. } => writeln!(f, "DEFINED -- hash: {expr_hash:#08x}")?,
                         Expr::FieldAccess { .. } => writeln!(f, "FIELD_ACCESS -- hash: {expr_hash:#08x}")?,
                         Expr::With { .. } => writeln!(f, "WITH -- hash: {expr_hash:#08x}")?,
@@ -2235,6 +2250,10 @@ pub(crate) enum Expr {
     /// `matches` expression.
     Matches { rhs: ExprId, lhs: ExprId },
 
+    /// Match expression similar to `Expr::Matches` but that matches against
+    /// multiple regular expressions at the same time.
+    MatchesMany { lhs: ExprId, regex_set: RegexSetId },
+
     /// A `defined` expression (e.g. `defined foo`)
     Defined { operand: ExprId },
 
@@ -2682,6 +2701,11 @@ impl Expr {
                     *rhs = replacement;
                 }
             }
+            Expr::MatchesMany { lhs, .. } => {
+                if *lhs == child {
+                    *lhs = replacement;
+                }
+            }
             Expr::PatternMatch { anchor, .. }
             | Expr::PatternMatchVar { anchor, .. } => {
                 replace_in_anchor(anchor)
@@ -2786,6 +2810,7 @@ impl Expr {
             | Expr::IEndsWith { .. }
             | Expr::IEquals { .. }
             | Expr::Matches { .. }
+            | Expr::MatchesMany { .. }
             | Expr::PatternMatch { .. }
             | Expr::PatternMatchVar { .. }
             | Expr::OfExprTuple(_)
@@ -2857,6 +2882,7 @@ impl Expr {
             | Expr::IEndsWith { .. }
             | Expr::IEquals { .. }
             | Expr::Matches { .. }
+            | Expr::MatchesMany { .. }
             | Expr::PatternMatch { .. }
             | Expr::PatternMatchVar { .. }
             | Expr::OfExprTuple(_)