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, .. } // Can technically be derived by summing parent cardinalities
19        | HydroNode::CrossSingleton { metadata, .. }
20        | HydroNode::CrossProduct { metadata, .. } // Can technically be derived by multiplying parent cardinalities
21        | 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, .. } // Output 1 value per tick
29        | HydroNode::Reduce { metadata, .. } // Output 1 value per tick
30        | 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            // when we emit this IR, the counter will bump the stmt id, so simulate that here
45            *next_stmt_id += 1;
46
47            *node = counter;
48        }
49        HydroNode::Tee { .. } // Do nothing, we will count the parent of the Tee
50        | HydroNode::Map { .. } // Equal to parent cardinality
51        | HydroNode::DeferTick { .. } // Equal to parent cardinality
52        | 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}