Finished merge

davidchuyaya · Mar 14, 2024 · 00d28a0 · 00d28a0
1 parent 88dd9b7
commit 00d28a0
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 4 deletions.
diff --git a/hydroflow_plus/src/ir.rs b/hydroflow_plus/src/ir.rs
@@ -79,7 +79,7 @@ pub enum HfPlusLeaf {
 impl HfPlusLeaf {
     pub fn transform_children(
         self,
-        transform: impl Fn(HfPlusNode, &mut SeenTees) -> HfPlusNode,
+        mut transform: impl FnMut(HfPlusNode, &mut SeenTees) -> HfPlusNode,
         seen_tees: &mut SeenTees,
     ) -> HfPlusLeaf {
         match self {
@@ -246,7 +246,7 @@ pub type SeenTees = HashMap<*const RefCell<HfPlusNode>, Rc<RefCell<HfPlusNode>>>
 impl HfPlusNode {
     pub fn transform_children(
         self,
-        transform: impl Fn(HfPlusNode, &mut SeenTees) -> HfPlusNode,
+        mut transform: impl FnMut(HfPlusNode, &mut SeenTees) -> HfPlusNode,
         seen_tees: &mut SeenTees,
     ) -> HfPlusNode {
         match self {

diff --git a/hydroflow_plus/src/lib.rs b/hydroflow_plus/src/lib.rs
@@ -36,6 +36,7 @@ pub use builder::FlowBuilder;
 pub mod ir;
 
 pub mod persist_pullup;
+pub mod profiler;
 
 #[derive(Clone)]
 pub struct RuntimeContext<'a> {

diff --git a/hydroflow_plus/src/profiler.rs b/hydroflow_plus/src/profiler.rs
@@ -52,11 +52,12 @@ mod tests {
             .map(q!(|v| v + 1))
             .for_each(q!(|n| println!("{}", n)));
 
-        let built = flow.build();
+        let built = flow.extract();
 
         insta::assert_debug_snapshot!(&built.ir);
 
-        let pushed_down = built.optimize_with(|ir| super::profiling(ir, flow.runtime_context()));
+        let counters = RuntimeData::new("Fake");
+        let pushed_down = built.optimize_with(|ir| super::profiling(ir, flow.runtime_context(), counters));
 
         insta::assert_debug_snapshot!(&pushed_down.ir);
     }

diff --git a/.../src/snapshots/hydroflow_plus__profiler__tests__predicate_pushdown_through_map-2.snap.new b/.../src/snapshots/hydroflow_plus__profiler__tests__predicate_pushdown_through_map-2.snap.new
@@ -0,0 +1,30 @@
+---
+source: hydroflow_plus/src/profiler.rs
+assertion_line: 62
+expression: "&pushed_down.ir"
+---
+[
+    ForEach {
+        f: { use crate :: __staged :: profiler :: tests :: * ; | n | println ! ("{}" , n) },
+        input: Map {
+            f: { use crate :: __staged :: profiler :: * ; let counters_copy = Fake ; let my_id = 0u32 ; let context = & context ; let my_id = 32668u32 ; let counters = Fake ; let my_id = 32668u32 ; { counters_copy . borrow_mut () [my_id as usize] = 0 ; move | v | { hydroflow_plus :: profiler :: increment_counter (context . current_tick () , my_id , & mut counters . borrow_mut () [my_id as usize]) ; v } } },
+            input: Map {
+                f: { use crate :: __staged :: profiler :: tests :: * ; | v | v + 1 },
+                input: Map {
+                    f: { use crate :: __staged :: profiler :: * ; let counters_copy = Fake ; let my_id = 1u32 ; let context = & context ; let my_id = 32668u32 ; let counters = Fake ; let my_id = 32668u32 ; { counters_copy . borrow_mut () [my_id as usize] = 0 ; move | v | { hydroflow_plus :: profiler :: increment_counter (context . current_tick () , my_id , & mut counters . borrow_mut () [my_id as usize]) ; v } } },
+                    input: Persist(
+                        Map {
+                            f: { use crate :: __staged :: profiler :: * ; let counters_copy = Fake ; let my_id = 2u32 ; let context = & context ; let my_id = 32668u32 ; let counters = Fake ; let my_id = 32668u32 ; { counters_copy . borrow_mut () [my_id as usize] = 0 ; move | v | { hydroflow_plus :: profiler :: increment_counter (context . current_tick () , my_id , & mut counters . borrow_mut () [my_id as usize]) ; v } } },
+                            input: Source {
+                                source: Iter(
+                                    { use crate :: __staged :: profiler :: tests :: * ; 0 .. 10 },
+                                ),
+                                location_id: 0,
+                            },
+                        },
+                    ),
+                },
+            },
+        },
+    },
+]
diff --git a/...us/src/snapshots/hydroflow_plus__profiler__tests__predicate_pushdown_through_map.snap.new b/...us/src/snapshots/hydroflow_plus__profiler__tests__predicate_pushdown_through_map.snap.new
@@ -0,0 +1,21 @@
+---
+source: hydroflow_plus/src/profiler.rs
+assertion_line: 57
+expression: "&built.ir"
+---
+[
+    ForEach {
+        f: { use crate :: __staged :: profiler :: tests :: * ; | n | println ! ("{}" , n) },
+        input: Map {
+            f: { use crate :: __staged :: profiler :: tests :: * ; | v | v + 1 },
+            input: Persist(
+                Source {
+                    source: Iter(
+                        { use crate :: __staged :: profiler :: tests :: * ; 0 .. 10 },
+                    ),
+                    location_id: 0,
+                },
+            ),
+        },
+    },
+]