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
8pub 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, .. } | HydroNode::DeferTick { input, metadata, .. } | 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}