Repository: pola-rs/polars
File: crates/polars-plan/src/plans/iterator.rs (Rust)

Code before:
        self.stack
            .pop()
            .inspect(|current_expr| current_expr.nodes(&mut self.stack))
    }
}

pub struct ExprMapper<F> {
    f: F,
}

impl<F: FnMut(Expr) -> PolarsResult<Expr>> RewritingVisitor for ExprMapper<F> {
    type Node = Expr;
    type Arena = ();

    fn mutate(&mut self, node: Self::Node, _arena: &mut Self::Arena) -> PolarsResult<Self::Node> {
        (self.f)(node)
    }
}

impl Expr {
    pub fn nodes<'a>(&'a self, container: &mut UnitVec<&'a Expr>) {
        let push = |c: &mut UnitVec<&'a Expr>, e: &'a Expr| c.push(e);
        push_expr!(self, container, push, push, iter);
    }

    pub fn nodes_owned(self, container: &mut UnitVec<Expr>) {
        let push_arc = |c: &mut UnitVec<Expr>, e: Arc<Expr>| c.push(Arc::unwrap_or_clone(e));
        let push_owned = |c: &mut UnitVec<Expr>, e: Expr| c.push(e);
        push_expr!(self, container, push_arc, push_owned, into_iter);
    }

    pub fn map_expr<F: FnMut(Self) -> Self>(self, mut f: F) -> Self {
        self.rewrite(&mut ExprMapper { f: |e| Ok(f(e)) }, &mut ())
            .unwrap()
    }

    pub fn try_map_expr<F: FnMut(Self) -> PolarsResult<Self>>(self, f: F) -> PolarsResult<Self> {
        self.rewrite(&mut ExprMapper { f }, &mut ())
    }
}

impl<'a> IntoIterator for &'a Expr {
    type Item = &'a Expr;
    type IntoIter = ExprIter<'a>;

    fn into_iter(self) -> Self::IntoIter {
        let stack = unitvec!(self);
        ExprIter { stack }
    }
}


Review comment (suggestion):
Note, this was not used anywhere else yet and it would throw an assertion error when used because UnitVec is not allowed with types that are larger than `usize`.

Code after:
        self.stack
            .pop()
            .inspect(|current_expr| current_expr.nodes(&mut self.stack))
    }
}

pub struct ExprMapper<F> {
    f: F,
}

impl<F: FnMut(Expr) -> PolarsResult<Expr>> RewritingVisitor for ExprMapper<F> {
    type Node = Expr;
    type Arena = ();

    fn mutate(&mut self, node: Self::Node, _arena: &mut Self::Arena) -> PolarsResult<Self::Node> {
        (self.f)(node)
    }
}

impl Expr {
    pub fn nodes<'a>(&'a self, container: &mut UnitVec<&'a Expr>) {
        let push = |c: &mut UnitVec<&'a Expr>, e: &'a Expr| c.push(e);
        push_expr!(self, container, push, push, iter);
    }

    pub fn nodes_owned(self, container: &mut Vec<Expr>) {
        let push_arc = |c: &mut Vec<Expr>, e: Arc<Expr>| c.push(Arc::unwrap_or_clone(e));
        let push_owned = |c: &mut Vec<Expr>, e: Expr| c.push(e);
        push_expr!(self, container, push_arc, push_owned, into_iter);
    }

    pub fn map_expr<F: FnMut(Self) -> Self>(self, mut f: F) -> Self {
        self.rewrite(&mut ExprMapper { f: |e| Ok(f(e)) }, &mut ())
            .unwrap()
    }

    pub fn try_map_expr<F: FnMut(Self) -> PolarsResult<Self>>(self, f: F) -> PolarsResult<Self> {
        self.rewrite(&mut ExprMapper { f }, &mut ())
    }
}

impl<'a> IntoIterator for &'a Expr {
    type Item = &'a Expr;
    type IntoIter = ExprIter<'a>;

    fn into_iter(self) -> Self::IntoIter {
        let stack = unitvec!(self);
        ExprIter { stack }
    }
}