diff --git a/src/ir/tree.rs b/src/ir/tree.rs index 39a99f7a1d818855c6c078377b6326cab47c43f4..960987fff53dcbaa94fccf095ccb039ce48d174b 100644 --- a/src/ir/tree.rs +++ b/src/ir/tree.rs @@ -1,4 +1,5 @@ use super::expression::Expression; +use super::operator::Bool; use super::{Node, Nodes}; use std::cell::RefCell; @@ -13,6 +14,16 @@ pub struct ExpressionIterator<'n> { nodes: &'n Nodes, } +/// Children iterator for "and" node chains. +/// +/// Iterator returns the next child for Bool::And nodes. +#[derive(Debug)] +pub struct AndChainIterator<'n> { + current: &'n usize, + child: RefCell<usize>, + nodes: &'n Nodes, +} + impl<'n> Nodes { #[must_use] pub fn expr_iter(&'n self, current: &'n usize) -> ExpressionIterator<'n> { @@ -22,6 +33,15 @@ impl<'n> Nodes { nodes: self, } } + + #[must_use] + pub fn and_iter(&'n self, current: &'n usize) -> AndChainIterator<'n> { + AndChainIterator { + current, + child: RefCell::new(0), + nodes: self, + } + } } impl<'n> Iterator for ExpressionIterator<'n> { @@ -67,5 +87,29 @@ impl<'n> Iterator for ExpressionIterator<'n> { } } +impl<'n> Iterator for AndChainIterator<'n> { + type Item = &'n usize; + + fn next(&mut self) -> Option<Self::Item> { + if let Some(Node::Expression(Expression::Bool {left, op, right, .. })) = self.nodes.arena.get(*self.current) { + if *op != Bool::And { + return None; + } + let child_step = *self.child.borrow(); + if child_step == 0 { + *self.child.borrow_mut() += 1; + return Some(left); + } else if child_step == 1 { + *self.child.borrow_mut() += 1; + return Some(right); + } + None + } + else { + None + } + } +} + #[cfg(test)] mod tests; diff --git a/src/ir/tree/tests.rs b/src/ir/tree/tests.rs index 7f91ea1316928a3e0b8a833155d327082bf7c80d..f84f2f1e7ae07c6a48db5d7a1819d5f2ff79fb95 100644 --- a/src/ir/tree/tests.rs +++ b/src/ir/tree/tests.rs @@ -2,7 +2,7 @@ use crate::ir::operator::*; use crate::ir::value::*; use crate::ir::*; use pretty_assertions::assert_eq; -use traversal::Bft; +use traversal::{Bft, DftPre}; #[test] fn expression_bft() { @@ -38,3 +38,30 @@ fn expression_bft() { assert_eq!(bft_tree.next(), Some((3, &c3))); assert_eq!(bft_tree.next(), None); } + +#[test] +fn and_chain() { + // (((b1 or b2) and b3) and b4) and b5 + + let mut plan = Plan::new(); + let b1 = plan.nodes.add_const(Value::Boolean(true)); + let b2 = plan.nodes.add_const(Value::Boolean(true)); + let b3 = plan.nodes.add_const(Value::Boolean(true)); + let b4 = plan.nodes.add_const(Value::Boolean(true)); + let b5 = plan.nodes.add_const(Value::Boolean(true)); + + let b1_2 = plan.nodes.add_bool(b1, Bool::Or, b2).unwrap(); + let b1_23 = plan.nodes.add_bool(b1_2, Bool::And, b3).unwrap(); + let b1_234 = plan.nodes.add_bool(b1_23, Bool::And, b4).unwrap(); + let top = plan.nodes.add_bool(b1_234, Bool::And, b5).unwrap(); + + let mut dft_pre = DftPre::new(&top, |node| plan.nodes.and_iter(node)); + assert_eq!(dft_pre.next(), Some((0, &top))); + assert_eq!(dft_pre.next(), Some((1, &b1_234))); + assert_eq!(dft_pre.next(), Some((2, &b1_23))); + assert_eq!(dft_pre.next(), Some((3, &b1_2))); + assert_eq!(dft_pre.next(), Some((3, &b3))); + assert_eq!(dft_pre.next(), Some((2, &b4))); + assert_eq!(dft_pre.next(), Some((1, &b5))); + assert_eq!(dft_pre.next(), None); +}