hydro_lang/rewrites/
insert_counter.rs
1use std::time::Duration;
2
3use stageleft::Quoted;
4
5use crate::ir::*;
6
7fn insert_counter_node(node: &mut HydroNode, next_stmt_id: &mut usize, duration: syn::Expr) {
8 match node {
9 HydroNode::Placeholder
10 | HydroNode::Unpersist { .. }
11 | HydroNode::Counter { .. } => {
12 std::panic!("Unexpected {:?} found in insert_counter_node", node.print_root());
13 }
14 HydroNode::Source { metadata, .. }
15 | HydroNode::CycleSource { metadata, .. }
16 | HydroNode::Persist { metadata, .. }
17 | HydroNode::Delta { metadata, .. }
18 | HydroNode::Chain { metadata, .. } | HydroNode::CrossSingleton { metadata, .. }
20 | HydroNode::CrossProduct { metadata, .. } | HydroNode::Join { metadata, .. }
22 | HydroNode::Difference { metadata, .. }
23 | HydroNode::AntiJoin { metadata, .. }
24 | HydroNode::FlatMap { metadata, .. }
25 | HydroNode::Filter { metadata, .. }
26 | HydroNode::FilterMap { metadata, .. }
27 | HydroNode::Unique { metadata, .. }
28 | HydroNode::Fold { metadata, .. } | HydroNode::Reduce { metadata, .. } | HydroNode::FoldKeyed { metadata, .. }
31 | HydroNode::ReduceKeyed { metadata, .. }
32 | HydroNode::Network { metadata, .. }
33 => {
34 let metadata = metadata.clone();
35 let node_content = std::mem::replace(node, HydroNode::Placeholder);
36
37 let counter = HydroNode::Counter {
38 tag: next_stmt_id.to_string(),
39 duration: duration.into(),
40 input: Box::new(node_content),
41 metadata: metadata.clone(),
42 };
43
44 *next_stmt_id += 1;
46
47 *node = counter;
48 }
49 HydroNode::Tee { .. } | HydroNode::Map { .. } | HydroNode::DeferTick { .. } | HydroNode::Enumerate { .. }
53 | HydroNode::Inspect { .. }
54 | HydroNode::Sort { .. }
55 => {}
56 }
57}
58
59pub fn insert_counter(ir: &mut [HydroLeaf], duration: impl Quoted<'static, Duration>) {
60 let duration = duration.splice_typed();
61 traverse_dfir(
62 ir,
63 |_, _| {},
64 |node, next_stmt_id| {
65 insert_counter_node(node, next_stmt_id, duration.clone());
66 },
67 );
68}