feat: add ExprRef sumtype with better expression iteration

argus-core/src/expr.rs

@@ -12,6 +12,8 @@ pub use traits::*;
 
 use crate::{ArgusResult, Error};
 
+use self::iter::AstIter;
+
 /// All expressions that are numeric
 #[derive(Clone, Debug, PartialEq)]
 pub enum NumExpr {
@@ -29,17 +31,29 @@ pub enum NumExpr {
 }
 
 impl Expr for NumExpr {
+    fn is_numeric(&self) -> bool {
+        true
+    }
+
+    fn is_boolean(&self) -> bool {
+        false
+    }
+
+    fn args(&self) -> Vec<ExprRef<'_>> {
+        match self {
+            NumExpr::Neg { arg } => vec![arg.as_ref().into()],
+            NumExpr::Add { args } | NumExpr::Mul { args } => args.iter().map(|arg| arg.into()).collect(),
+            NumExpr::Div { dividend, divisor } => vec![dividend.as_ref().into(), divisor.as_ref().into()],
+            _ => vec![],
+        }
+    }
+
     fn as_any(&self) -> &dyn Any {
         self
     }
 
-    fn args(&self) -> Vec<&dyn Expr> {
-        match self {
-            NumExpr::Neg { arg } => vec![arg.as_ref()],
-            NumExpr::Add { args } | NumExpr::Mul { args } => args.iter().map(|arg| arg as &dyn Expr).collect(),
-            NumExpr::Div { dividend, divisor } => vec![dividend.as_ref(), divisor.as_ref()],
-            _ => vec![],
-        }
+    fn iter(&self) -> iter::AstIter<'_> {
+        AstIter::new(self.into())
     }
 }
 
@@ -64,11 +78,19 @@ pub enum BoolExpr {
 }
 
 impl Expr for BoolExpr {
-    fn args(&self) -> Vec<&dyn Expr> {
+    fn is_numeric(&self) -> bool {
+        false
+    }
+
+    fn is_boolean(&self) -> bool {
+        true
+    }
+
+    fn args(&self) -> Vec<ExprRef<'_>> {
         match self {
-            BoolExpr::Cmp { op: _, lhs, rhs } => vec![lhs.as_ref(), rhs.as_ref()],
-            BoolExpr::Not { arg } => vec![arg.as_ref()],
-            BoolExpr::And { args } | BoolExpr::Or { args } => args.iter().map(|arg| arg as &dyn Expr).collect(),
+            BoolExpr::Cmp { op: _, lhs, rhs } => vec![lhs.as_ref().into(), rhs.as_ref().into()],
+            BoolExpr::Not { arg } => vec![arg.as_ref().into()],
+            BoolExpr::And { args } | BoolExpr::Or { args } => args.iter().map(|arg| arg.into()).collect(),
             _ => vec![],
         }
     }
@@ -76,6 +98,16 @@ impl Expr for BoolExpr {
     fn as_any(&self) -> &dyn Any {
         self
     }
+
+    fn iter(&self) -> AstIter<'_> {
+        AstIter::new(self.into())
+    }
+}
+
+#[derive(Clone, Copy, Debug, derive_more::From)]
+pub enum ExprRef<'a> {
+    Bool(&'a BoolExpr),
+    Num(&'a NumExpr),
 }
 
 /// Expression builder

argus-core/src/expr/iter.rs

@@ -0,0 +1,42 @@
+use std::collections::VecDeque;
+
+use super::{Expr, ExprRef};
+
+/// Iterator that starts from some root [`Expr`] and travels down to it's leaf
+/// expressions.
+///
+/// This essentially implements breadth-first search over the expression tree rooted at
+/// the given [`Expr`].
+pub struct AstIter<'a> {
+    queue: VecDeque<ExprRef<'a>>,
+}
+
+impl<'a> AstIter<'a> {
+    /// Create an iterator that traverses an [`Expr`] from root to leaf.
+    pub fn new(root: ExprRef<'a>) -> Self {
+        let mut queue = VecDeque::new();
+        queue.push_back(root);
+        Self { queue }
+    }
+}
+
+impl<'a> Iterator for AstIter<'a> {
+    type Item = ExprRef<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let expr_ref = self.queue.pop_front()?;
+
+        let expr: &dyn Expr = match expr_ref {
+            ExprRef::Bool(expr) => expr,
+            ExprRef::Num(expr) => expr,
+        };
+
+        // We need to get all the arguments of the current expression (not including
+        //    any intervals), and push them into the queue.
+        for arg in expr.args().into_iter() {
+            self.queue.push_back(arg);
+        }
+        // 4. Give the user their expr
+        Some(expr_ref)
+    }
+}

argus-core/src/expr/traits.rs

@@ -1,19 +1,17 @@
 use std::any::Any;
 
-use super::iter::AstIter;
+use super::{iter::AstIter, ExprRef};
 
 /// A trait representing expressions
 pub trait Expr {
-    fn args(&self) -> Vec<&dyn Expr>;
+    fn is_numeric(&self) -> bool;
+    fn is_boolean(&self) -> bool;
+
+    fn args(&self) -> Vec<ExprRef<'_>>;
 
     fn as_any(&self) -> &dyn Any;
 
-    fn iter(&self) -> AstIter<'_>
-    where
-        Self: Sized,
-    {
-        AstIter::new(self)
-    }
+    fn iter(&self) -> AstIter<'_>;
 }
 
 impl dyn Expr {

argus-core/src/lib.rs

@@ -12,3 +12,4 @@ pub enum Error {
 
 pub type ArgusError = Error;
 pub type ArgusResult<T> = Result<T, Error>;
+