hydro_lang/rewrites/
analyze_counter.rs

1use std::collections::HashMap;
2
3use regex::Regex;
4
5pub use crate::internal_constants::COUNTER_PREFIX;
6use crate::ir::*;
7
8/// Returns (op_id, count)
9pub fn parse_counter_usage(measurement: String) -> (usize, usize) {
10    let regex = Regex::new(r"\((\d+)\): (\d+)").unwrap();
11    let matches = regex.captures_iter(&measurement).last().unwrap_or_else(|| {
12        ::core::panic!(
13            "Failed to parse counter usage from measurement: {:?}",
14            measurement
15        )
16    });
17    let op_id = matches[1].parse::<usize>().unwrap();
18    let count = matches[2].parse::<usize>().unwrap();
19    (op_id, count)
20}
21
22fn inject_count_node(
23    node: &mut HydroNode,
24    next_stmt_id: &mut usize,
25    op_to_count: &HashMap<usize, usize>,
26) {
27    if let Some(count) = op_to_count.get(next_stmt_id) {
28        let metadata = node.metadata_mut();
29        metadata.cardinality = Some(*count);
30    } else {
31        match node {
32            HydroNode::Tee { inner ,metadata, .. } => {
33                metadata.cardinality = inner.0.borrow().metadata().cardinality;
34            }
35            | HydroNode::Map { input, metadata, .. } // Equal to parent cardinality
36            | HydroNode::DeferTick { input, metadata, .. } // Equal to parent cardinality
37            | HydroNode::Enumerate { input, metadata, .. }
38            | HydroNode::Inspect { input, metadata, .. }
39            | HydroNode::Sort { input, metadata, .. }
40            | HydroNode::Counter { input, metadata, .. }
41            => {
42                metadata.cardinality = input.metadata().cardinality;
43            }
44            _ => {}
45        }
46    }
47}
48
49pub fn inject_count(ir: &mut [HydroLeaf], op_to_count: &HashMap<usize, usize>) {
50    traverse_dfir(
51        ir,
52        |_, _| {},
53        |node, next_stmt_id| {
54            inject_count_node(node, next_stmt_id, op_to_count);
55        },
56    );
57}