use crate::ast::*;
#[derive(Debug)]
pub enum DFSEvent<T> {
Enter(T),
Leave(T),
}
pub enum DFSContext<'src> {
Root,
Body(&'src Expr<'src>),
Quantifier(&'src Expr<'src>),
Operand(&'src Expr<'src>),
WithDeclaration(&'src Expr<'src>),
Range(&'src Expr<'src>),
Anchor(&'src Expr<'src>),
Items(&'src Expr<'src>),
Index(&'src Expr<'src>),
FuncArg(&'src Expr<'src>),
FuncSelf(&'src Expr<'src>),
}
pub struct DFSIter<'src> {
stack: Vec<DFSEvent<(&'src Expr<'src>, DFSContext<'src>)>>,
recently_left_context: Option<DFSContext<'src>>,
}
impl<'src> DFSIter<'src> {
pub fn new(expr: &'src Expr<'src>) -> Self {
Self {
stack: vec![DFSEvent::Enter((expr, DFSContext::Root))],
recently_left_context: None,
}
}
pub fn contexts(
&self,
) -> impl DoubleEndedIterator<Item = &DFSContext<'src>> {
itertools::chain(
self.recently_left_context.iter(),
self.stack.iter().rev().filter_map(|event| match event {
DFSEvent::Enter(_) => None,
DFSEvent::Leave((_, ctx)) => Some(ctx),
}),
)
}
pub fn prune(&mut self) {
while let Some(DFSEvent::Enter(_)) = self.stack.last() {
self.stack.pop();
}
}
}
impl<'src> Iterator for DFSIter<'src> {
type Item = DFSEvent<&'src Expr<'src>>;
fn next(&mut self) -> Option<Self::Item> {
match self.stack.pop()? {
DFSEvent::Enter((expr, context)) => {
self.recently_left_context = None;
self.stack.push(DFSEvent::Leave((expr, context)));
dfs_enter(expr, &mut self.stack);
Some(DFSEvent::Enter(expr))
}
DFSEvent::Leave((expr, context)) => {
self.recently_left_context = Some(context);
Some(DFSEvent::Leave(expr))
}
}
}
}
fn dfs_enter<'a>(
expr: &'a Expr,
stack: &mut Vec<DFSEvent<(&'a Expr<'a>, DFSContext<'a>)>>,
) {
match expr {
Expr::True { .. }
| Expr::False { .. }
| Expr::Filesize { .. }
| Expr::Entrypoint { .. }
| Expr::LiteralString(_)
| Expr::LiteralInteger(_)
| Expr::LiteralFloat(_)
| Expr::Regexp(_)
| Expr::Ident(_) => {}
Expr::PatternCount(p) => {
if let Some(r) = &p.range {
stack.push(DFSEvent::Enter((
&r.upper_bound,
DFSContext::Range(expr),
)));
stack.push(DFSEvent::Enter((
&r.lower_bound,
DFSContext::Range(expr),
)));
}
}
Expr::PatternOffset(p) | Expr::PatternLength(p) => {
if let Some(index) = &p.index {
stack.push(DFSEvent::Enter((index, DFSContext::Index(expr))));
}
}
Expr::PatternMatch(m) => {
if let Some(anchor) = &m.anchor {
match anchor {
MatchAnchor::At(at) => {
stack.push(DFSEvent::Enter((
&at.expr,
DFSContext::Anchor(expr),
)));
}
MatchAnchor::In(in_expr) => {
stack.push(DFSEvent::Enter((
&in_expr.range.upper_bound,
DFSContext::Anchor(expr),
)));
stack.push(DFSEvent::Enter((
&in_expr.range.lower_bound,
DFSContext::Anchor(expr),
)));
}
}
}
}
Expr::Lookup(lookup) => {
stack.push(DFSEvent::Enter((
&lookup.index,
DFSContext::Index(expr),
)));
stack.push(DFSEvent::Enter((
&lookup.primary,
DFSContext::Operand(expr),
)));
}
Expr::FieldAccess(e) => {
for operand in e.operands.iter().rev() {
stack.push(DFSEvent::Enter((
operand,
DFSContext::Operand(expr),
)));
}
}
Expr::FuncCall(func) => {
for arg in func.args.iter().rev() {
stack.push(DFSEvent::Enter((arg, DFSContext::FuncArg(expr))));
}
if let Some(obj) = &func.object {
stack.push(DFSEvent::Enter((obj, DFSContext::FuncSelf(expr))));
}
}
Expr::Defined(e)
| Expr::Not(e)
| Expr::Minus(e)
| Expr::BitwiseNot(e) => {
stack.push(DFSEvent::Enter((
&e.operand,
DFSContext::Operand(expr),
)));
}
Expr::And(e) | Expr::Or(e) => {
for operand in e.operands.iter().rev() {
stack.push(DFSEvent::Enter((
operand,
DFSContext::Operand(expr),
)));
}
}
Expr::Add(e)
| Expr::Sub(e)
| Expr::Mul(e)
| Expr::Div(e)
| Expr::Mod(e) => {
for operand in e.operands.iter().rev() {
stack.push(DFSEvent::Enter((
operand,
DFSContext::Operand(expr),
)));
}
}
Expr::Shl(e)
| Expr::Shr(e)
| Expr::BitwiseAnd(e)
| Expr::BitwiseOr(e)
| Expr::BitwiseXor(e)
| Expr::Eq(e)
| Expr::Ne(e)
| Expr::Lt(e)
| Expr::Gt(e)
| Expr::Le(e)
| Expr::Ge(e)
| Expr::Contains(e)
| Expr::IContains(e)
| Expr::StartsWith(e)
| Expr::IStartsWith(e)
| Expr::EndsWith(e)
| Expr::IEndsWith(e)
| Expr::IEquals(e)
| Expr::Matches(e) => {
stack.push(DFSEvent::Enter((&e.rhs, DFSContext::Operand(expr))));
stack.push(DFSEvent::Enter((&e.lhs, DFSContext::Operand(expr))));
}
Expr::Of(of) => {
if let Some(anchor) = &of.anchor {
match anchor {
MatchAnchor::At(at) => {
stack.push(DFSEvent::Enter((
&at.expr,
DFSContext::Anchor(expr),
)));
}
MatchAnchor::In(in_expr) => {
stack.push(DFSEvent::Enter((
&in_expr.range.upper_bound,
DFSContext::Anchor(expr),
)));
stack.push(DFSEvent::Enter((
&in_expr.range.lower_bound,
DFSContext::Anchor(expr),
)));
}
}
}
if let OfItems::BoolExprTuple(tuple) = &of.items {
for item in tuple.iter().rev() {
stack.push(DFSEvent::Enter((
item,
DFSContext::Items(expr),
)));
}
}
match &of.quantifier {
Quantifier::Percentage(quantifier)
| Quantifier::Expr(quantifier) => {
stack.push(DFSEvent::Enter((
quantifier,
DFSContext::Quantifier(expr),
)));
}
_ => {}
}
}
Expr::ForOf(for_of) => {
stack
.push(DFSEvent::Enter((&for_of.body, DFSContext::Body(expr))));
match &for_of.quantifier {
Quantifier::Percentage(quantifier)
| Quantifier::Expr(quantifier) => {
stack.push(DFSEvent::Enter((
quantifier,
DFSContext::Quantifier(expr),
)));
}
_ => {}
}
}
Expr::ForIn(for_in) => {
stack
.push(DFSEvent::Enter((&for_in.body, DFSContext::Body(expr))));
match &for_in.iterable {
Iterable::Range(range) => {
stack.push(DFSEvent::Enter((
&range.upper_bound,
DFSContext::Items(expr),
)));
stack.push(DFSEvent::Enter((
&range.lower_bound,
DFSContext::Items(expr),
)));
}
Iterable::ExprTuple(tuple) => {
for item in tuple.iter().rev() {
stack.push(DFSEvent::Enter((
item,
DFSContext::Items(expr),
)));
}
}
Iterable::Expr(iterable_expr) => {
stack.push(DFSEvent::Enter((
iterable_expr,
DFSContext::Items(expr),
)));
}
}
match &for_in.quantifier {
Quantifier::Percentage(quantifier)
| Quantifier::Expr(quantifier) => {
stack.push(DFSEvent::Enter((
quantifier,
DFSContext::Quantifier(expr),
)));
}
_ => {}
}
}
Expr::With(with) => {
stack.push(DFSEvent::Enter((&with.body, DFSContext::Body(expr))));
for declaration in with.declarations.iter().rev() {
stack.push(DFSEvent::Enter((
&declaration.expression,
DFSContext::WithDeclaration(expr),
)));
}
}
}
}
#[cfg(test)]
mod tests {
use crate::ast::dfs::{DFSContext, DFSEvent, DFSIter};
use crate::ast::{Expr, AST};
#[test]
fn dfs() {
let source = br#"
rule test {
condition:
(true and false) or (1 + 2 > 5)
}
"#;
let ast = AST::from(source.as_slice());
let mut dfs = DFSIter::new(&ast.rules().next().unwrap().condition);
assert!(matches!(dfs.next(), Some(DFSEvent::Enter(Expr::Or(_)))));
assert!(matches!(dfs.next(), Some(DFSEvent::Enter(Expr::And(_)))));
dfs.prune();
assert!(matches!(dfs.next(), Some(DFSEvent::Leave(Expr::And(_)))));
assert!(matches!(dfs.next(), Some(DFSEvent::Enter(Expr::Gt(_)))));
assert!(matches!(dfs.next(), Some(DFSEvent::Enter(Expr::Add(_)))));
assert!(matches!(
dfs.next(),
Some(DFSEvent::Enter(Expr::LiteralInteger(_)))
));
assert!(matches!(
dfs.next(),
Some(DFSEvent::Leave(Expr::LiteralInteger(_)))
));
dfs.prune();
assert!(matches!(dfs.next(), Some(DFSEvent::Leave(Expr::Add(_)))));
assert!(matches!(
dfs.next(),
Some(DFSEvent::Enter(Expr::LiteralInteger(_)))
));
assert!(matches!(
dfs.next(),
Some(DFSEvent::Leave(Expr::LiteralInteger(_)))
));
assert!(matches!(dfs.next(), Some(DFSEvent::Leave(Expr::Gt(_)))));
assert!(matches!(dfs.next(), Some(DFSEvent::Leave(Expr::Or(_)))));
assert!(dfs.next().is_none());
}
#[test]
fn dfs_contexts() {
let source = r#"
rule test {
strings:
$a = "foo"
condition:
for 1 of ($a) : (
with i = 1 : (
i == 1
)
)
}
"#;
let ast = AST::from(source);
let mut dfs = DFSIter::new(&ast.rules().next().unwrap().condition);
assert!(dfs.contexts().next().is_none());
assert!(matches!(dfs.next(), Some(DFSEvent::Enter(Expr::ForOf(_)))));
let mut contexts = dfs.contexts();
assert!(matches!(contexts.next(), Some(DFSContext::Root)));
assert!(contexts.next().is_none());
drop(contexts);
assert!(matches!(
dfs.next(),
Some(DFSEvent::Enter(Expr::LiteralInteger(_)))
));
let mut contexts = dfs.contexts();
assert!(matches!(
contexts.next(),
Some(DFSContext::Quantifier(Expr::ForOf(_)))
));
assert!(matches!(contexts.next(), Some(DFSContext::Root)));
assert!(contexts.next().is_none());
drop(contexts);
assert!(matches!(
dfs.next(),
Some(DFSEvent::Leave(Expr::LiteralInteger(_)))
));
let mut contexts = dfs.contexts();
assert!(matches!(
contexts.next(),
Some(DFSContext::Quantifier(Expr::ForOf(_)))
));
assert!(matches!(contexts.next(), Some(DFSContext::Root)));
assert!(contexts.next().is_none());
drop(contexts);
assert!(matches!(dfs.next(), Some(DFSEvent::Enter(Expr::With(_)))));
let mut contexts = dfs.contexts();
assert!(matches!(
contexts.next(),
Some(DFSContext::Body(Expr::ForOf(_)))
));
assert!(matches!(contexts.next(), Some(DFSContext::Root)));
assert!(contexts.next().is_none());
drop(contexts);
assert!(matches!(
dfs.next(),
Some(DFSEvent::Enter(Expr::LiteralInteger(_)))
));
assert!(matches!(
dfs.next(),
Some(DFSEvent::Leave(Expr::LiteralInteger(_)))
));
assert!(matches!(dfs.next(), Some(DFSEvent::Enter(Expr::Eq(_)))));
let mut contexts = dfs.contexts();
assert!(matches!(
contexts.next(),
Some(DFSContext::Body(Expr::With(_)))
));
assert!(matches!(
contexts.next(),
Some(DFSContext::Body(Expr::ForOf(_)))
));
assert!(matches!(contexts.next(), Some(DFSContext::Root)));
assert!(contexts.next().is_none());
drop(contexts);
assert!(matches!(dfs.next(), Some(DFSEvent::Enter(Expr::Ident(_)))));
let mut contexts = dfs.contexts();
assert!(matches!(
contexts.next(),
Some(DFSContext::Operand(Expr::Eq(_)))
));
assert!(matches!(
contexts.next(),
Some(DFSContext::Body(Expr::With(_)))
));
assert!(matches!(
contexts.next(),
Some(DFSContext::Body(Expr::ForOf(_)))
));
assert!(matches!(contexts.next(), Some(DFSContext::Root)));
assert!(contexts.next().is_none());
drop(contexts);
assert!(matches!(dfs.next(), Some(DFSEvent::Leave(Expr::Ident(_)))));
assert!(matches!(
dfs.next(),
Some(DFSEvent::Enter(Expr::LiteralInteger(_)))
));
assert!(matches!(
dfs.next(),
Some(DFSEvent::Leave(Expr::LiteralInteger(_)))
));
assert!(matches!(dfs.next(), Some(DFSEvent::Leave(Expr::Eq(_)))));
assert!(matches!(dfs.next(), Some(DFSEvent::Leave(Expr::With(_)))));
assert!(matches!(dfs.next(), Some(DFSEvent::Leave(Expr::ForOf(_)))));
let mut contexts = dfs.contexts();
assert!(matches!(contexts.next(), Some(DFSContext::Root)));
assert!(contexts.next().is_none());
}
}