perf: don't use dynamic dispatch in visitors (#15607)

pola-rs · Apr 12, 2024 · 5e74924 · 5e74924
1 parent 0a594af
commit 5e74924
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 20 deletions.
diff --git a/crates/polars-plan/src/logical_plan/visitor/expr.rs b/crates/polars-plan/src/logical_plan/visitor/expr.rs
@@ -9,9 +9,9 @@ use crate::prelude::*;
 impl TreeWalker for Expr {
     type Arena = ();
 
-    fn apply_children(
+    fn apply_children<F: FnMut(&Self, &Self::Arena) -> PolarsResult<VisitRecursion>>(
         &self,
-        op: &mut dyn FnMut(&Self, &Self::Arena) -> PolarsResult<VisitRecursion>,
+        op: &mut F,
         arena: &Self::Arena,
     ) -> PolarsResult<VisitRecursion> {
         let mut scratch = unitvec![];
@@ -29,9 +29,9 @@ impl TreeWalker for Expr {
         Ok(VisitRecursion::Continue)
     }
 
-    fn map_children(
+    fn map_children<F: FnMut(Self, &mut Self::Arena) -> PolarsResult<Self>>(
         self,
-        f: &mut dyn FnMut(Self, &mut Self::Arena) -> PolarsResult<Self>,
+        f: &mut F,
         _arena: &mut Self::Arena,
     ) -> PolarsResult<Self> {
         use polars_utils::functions::try_arc_map as am;
@@ -256,9 +256,9 @@ impl PartialEq for AExprArena<'_> {
 
 impl TreeWalker for AexprNode {
     type Arena = Arena<AExpr>;
-    fn apply_children(
+    fn apply_children<F: FnMut(&Self, &Self::Arena) -> PolarsResult<VisitRecursion>>(
         &self,
-        op: &mut dyn FnMut(&Self, &Self::Arena) -> PolarsResult<VisitRecursion>,
+        op: &mut F,
         arena: &Self::Arena,
     ) -> PolarsResult<VisitRecursion> {
         let mut scratch = unitvec![];
@@ -276,9 +276,9 @@ impl TreeWalker for AexprNode {
         Ok(VisitRecursion::Continue)
     }
 
-    fn map_children(
+    fn map_children<F: FnMut(Self, &mut Self::Arena) -> PolarsResult<Self>>(
         mut self,
-        op: &mut dyn FnMut(Self, &mut Self::Arena) -> PolarsResult<Self>,
+        op: &mut F,
         arena: &mut Self::Arena,
     ) -> PolarsResult<Self> {
         let mut scratch = unitvec![];

diff --git a/crates/polars-plan/src/logical_plan/visitor/lp.rs b/crates/polars-plan/src/logical_plan/visitor/lp.rs
@@ -46,9 +46,9 @@ pub type IRNodeArena = (Arena<IR>, Arena<AExpr>);
 impl TreeWalker for IRNode {
     type Arena = IRNodeArena;
 
-    fn apply_children(
+    fn apply_children<F: FnMut(&Self, &Self::Arena) -> PolarsResult<VisitRecursion>>(
         &self,
-        op: &mut dyn FnMut(&Self, &Self::Arena) -> PolarsResult<VisitRecursion>,
+        op: &mut F,
         arena: &Self::Arena,
     ) -> PolarsResult<VisitRecursion> {
         let mut scratch = unitvec![];
@@ -66,9 +66,9 @@ impl TreeWalker for IRNode {
         Ok(VisitRecursion::Continue)
     }
 
-    fn map_children(
+    fn map_children<F: FnMut(Self, &mut Self::Arena) -> PolarsResult<Self>>(
         self,
-        op: &mut dyn FnMut(Self, &mut Self::Arena) -> PolarsResult<Self>,
+        op: &mut F,
         arena: &mut Self::Arena,
     ) -> PolarsResult<Self> {
         let mut inputs = vec![];

diff --git a/crates/polars-plan/src/logical_plan/visitor/visitors.rs b/crates/polars-plan/src/logical_plan/visitor/visitors.rs
@@ -6,23 +6,23 @@ use super::*;
 /// Implemented for [`crate::dsl::Expr`] and [`AexprNode`].
 pub trait TreeWalker: Sized {
     type Arena;
-    fn apply_children(
+    fn apply_children<F: FnMut(&Self, &Self::Arena) -> PolarsResult<VisitRecursion>>(
         &self,
-        op: &mut dyn FnMut(&Self, &Self::Arena) -> PolarsResult<VisitRecursion>,
+        op: &mut F,
         arena: &Self::Arena,
     ) -> PolarsResult<VisitRecursion>;
 
-    fn map_children(
+    fn map_children<F: FnMut(Self, &mut Self::Arena) -> PolarsResult<Self>>(
         self,
-        op: &mut dyn FnMut(Self, &mut Self::Arena) -> PolarsResult<Self>,
+        op: &mut F,
         arena: &mut Self::Arena,
     ) -> PolarsResult<Self>;
 
     /// Walks all nodes in depth-first-order.
     #[recursive]
-    fn visit(
+    fn visit<V: Visitor<Node = Self, Arena = Self::Arena>>(
         &self,
-        visitor: &mut dyn Visitor<Node = Self, Arena = Self::Arena>,
+        visitor: &mut V,
         arena: &Self::Arena,
     ) -> PolarsResult<VisitRecursion> {
         match visitor.pre_visit(self, arena)? {
@@ -44,9 +44,9 @@ pub trait TreeWalker: Sized {
     }
 
     #[recursive]
-    fn rewrite(
+    fn rewrite<R: RewritingVisitor<Node = Self, Arena = Self::Arena>>(
         self,
-        rewriter: &mut dyn RewritingVisitor<Node = Self, Arena = Self::Arena>,
+        rewriter: &mut R,
         arena: &mut Self::Arena,
     ) -> PolarsResult<Self> {
         let mutate_this_node = match rewriter.pre_visit(&self, arena)? {