use std::collections::HashMap;
use nmbrs_workload::model::{BindingsDef, ParsedOp, WorkloadPhase};
use nmbrs_workload::polydat_matter::{HasPolydatMatter, PolydatMatter};
use crate::scope_tree::{ScopeKind, ScopeNodeIdx, ScopeTree};
pub struct ClassifyInputs<'a> {
pub bindings: &'a BindingsDef,
pub params: &'a HashMap<String, String>,
pub phases: &'a HashMap<String, WorkloadPhase>,
}
pub fn classify_and_mark(tree: &mut ScopeTree, inputs: &ClassifyInputs<'_>) {
let mut owning_phase: std::collections::HashMap<ScopeNodeIdx, String> =
std::collections::HashMap::new();
for (idx, node) in tree.iter_dfs() {
if !matches!(node.kind, ScopeKind::OpTemplate { .. }) {
continue;
}
let mut cursor = node.parent;
while let Some(p) = cursor {
if let ScopeKind::Phase { name } = &tree.nodes[p].kind {
owning_phase.insert(idx, name.clone());
break;
}
cursor = tree.nodes[p].parent;
}
}
tree.mark_scope_elision(|kind, idx| {
let matter = scope_kind_polydat_matter(kind, idx, inputs, &owning_phase);
matches!(matter, PolydatMatter::Definitions)
});
}
fn scope_kind_polydat_matter(
kind: &ScopeKind,
idx: ScopeNodeIdx,
inputs: &ClassifyInputs<'_>,
owning_phase: &std::collections::HashMap<ScopeNodeIdx, String>,
) -> PolydatMatter {
match kind {
ScopeKind::Session => PolydatMatter::None,
ScopeKind::Workload => {
if !inputs.bindings.is_empty() || !inputs.params.is_empty() {
PolydatMatter::Definitions
} else {
PolydatMatter::None
}
}
ScopeKind::Scenario { .. } => PolydatMatter::None,
ScopeKind::Phase { name } => inputs
.phases
.get(name)
.map(WorkloadPhase::polydat_matter)
.unwrap_or(PolydatMatter::None),
ScopeKind::OpTemplate { name } => {
let phase_name = match owning_phase.get(&idx) {
Some(n) => n,
None => return PolydatMatter::None,
};
inputs
.phases
.get(phase_name)
.and_then(|p| p.ops.iter().find(|op| op.name == *name))
.map(ParsedOp::polydat_matter)
.unwrap_or(PolydatMatter::None)
}
ScopeKind::Comprehension { .. } | ScopeKind::DoWhile { .. } | ScopeKind::DoUntil { .. } => {
PolydatMatter::Definitions
}
ScopeKind::IncludedScenario { .. } => PolydatMatter::None,
ScopeKind::Bindings { .. } => PolydatMatter::Definitions,
}
}
pub fn elision_summary(
tree: &ScopeTree,
) -> Vec<(ScopeNodeIdx, usize, Option<bool>, String, String)> {
tree.iter_dfs()
.map(|(idx, node)| {
(
idx,
node.depth,
node.materialised,
node.logical_name.clone(),
node.kind.label().to_string(),
)
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use nmbrs_workload::model::{BindingsDef, ScenarioNode, WorkloadPhase};
fn empty_phase() -> WorkloadPhase {
WorkloadPhase {
key_metrics: Vec::new(),
dimensions: Default::default(),
cycles: None,
concurrency: None,
rate: None,
daemon: false,
adapter: None,
errors: None,
tries: None,
tries_backoff: None,
interval: None,
repeat: None,
error_rate_max: None,
timeout: None,
stop_when: Vec::new(),
throttle: None,
continue_if: None,
tags: None,
ops: vec![],
for_each: None,
loop_scope: None,
iter_scope: None,
checkpoint: None,
status_metrics: vec![],
metrics: Default::default(),
bindings: BindingsDef::default(),
poll: None,
optimize: None,
}
}
fn mark_with(
tree: &mut ScopeTree,
bindings: &BindingsDef,
params: &HashMap<String, String>,
phases: &HashMap<String, WorkloadPhase>,
) {
let inputs = ClassifyInputs {
bindings,
params,
phases,
};
classify_and_mark(tree, &inputs);
}
#[test]
fn empty_workload_elides_everything_below_root() {
let mut phases = HashMap::new();
phases.insert("p".into(), empty_phase());
let mut tree = ScopeTree::build("default", &[ScenarioNode::Phase("p".into())]);
mark_with(&mut tree, &BindingsDef::default(), &HashMap::new(), &phases);
assert_eq!(tree.nodes[0].materialised, Some(true));
let scenario_idx = tree.nodes[tree.nodes[0].children[0]].children[0];
let phase_idx = tree.nodes[scenario_idx].children[0];
assert_eq!(tree.nodes[scenario_idx].materialised, Some(false));
assert_eq!(tree.nodes[phase_idx].materialised, Some(false));
assert_eq!(tree.nodes[0].logical_name, "");
assert_eq!(
tree.nodes[scenario_idx].logical_name,
"workload.scenario.default"
);
assert_eq!(
tree.nodes[phase_idx].logical_name,
"workload.scenario.default.phase.p"
);
}
#[test]
fn phase_with_bindings_materialises() {
let mut phases = HashMap::new();
let mut p1 = empty_phase();
p1.bindings = BindingsDef::PolydatSource("k := 5".into());
phases.insert("p1".into(), p1);
phases.insert("p2".into(), empty_phase());
let mut tree = ScopeTree::build(
"default",
&[
ScenarioNode::Phase("p1".into()),
ScenarioNode::Phase("p2".into()),
],
);
mark_with(&mut tree, &BindingsDef::default(), &HashMap::new(), &phases);
let scenario_idx = tree.nodes[tree.nodes[0].children[0]].children[0];
let p1_idx = tree.nodes[scenario_idx].children[0];
let p2_idx = tree.nodes[scenario_idx].children[1];
assert_eq!(tree.nodes[p1_idx].materialised, Some(true));
assert_eq!(tree.nodes[p2_idx].materialised, Some(false));
}
#[test]
fn workload_with_top_level_bindings_materialises_root() {
let mut phases = HashMap::new();
phases.insert("p".into(), empty_phase());
let bindings = BindingsDef::PolydatSource("dataset := \"sift\"".into());
let mut tree = ScopeTree::build("default", &[ScenarioNode::Phase("p".into())]);
mark_with(&mut tree, &bindings, &HashMap::new(), &phases);
assert_eq!(tree.nodes[0].materialised, Some(true));
}
#[test]
fn workload_with_params_classifies_root_as_definitions() {
let mut phases = HashMap::new();
phases.insert("p".into(), empty_phase());
let mut params = HashMap::new();
params.insert("dataset".into(), "sift".into());
let mut tree = ScopeTree::build("default", &[ScenarioNode::Phase("p".into())]);
mark_with(&mut tree, &BindingsDef::default(), ¶ms, &phases);
assert_eq!(tree.nodes[0].materialised, Some(true));
}
#[test]
fn comprehension_node_always_materialises() {
let mut phases = HashMap::new();
phases.insert("p".into(), empty_phase());
let comp = polydat::iteration::comprehension::Comprehension::cartesian(vec![]);
let mut tree = ScopeTree::build(
"default",
&[ScenarioNode::Comprehension {
anchor: None,
comprehension: comp,
children: vec![ScenarioNode::Phase("p".into())],
continue_if: None,
}],
);
mark_with(&mut tree, &BindingsDef::default(), &HashMap::new(), &phases);
let scenario_idx = tree.nodes[tree.nodes[0].children[0]].children[0];
let comp_idx = tree.nodes[scenario_idx].children[0];
assert_eq!(tree.nodes[comp_idx].materialised, Some(true));
}
#[test]
fn op_template_with_metrics_materialises() {
use nmbrs_workload::model::{MetricSpec, ParsedOp};
let mut phases = HashMap::new();
let mut p = empty_phase();
let mut op = ParsedOp::simple("a", "noop");
op.metrics.insert(
"m".into(),
MetricSpec {
cell: Default::default(),
value: "factor * 2.0".into(), family: None,
kind: None,
unit: None,
format: None,
},
);
p.ops.push(op);
phases.insert("p".into(), p);
let mut tree = ScopeTree::build("default", &[ScenarioNode::Phase("p".into())]);
tree.extend_with_op_templates(&phases);
mark_with(&mut tree, &BindingsDef::default(), &HashMap::new(), &phases);
let op_idx = tree
.iter_dfs()
.find(|(_, n)| {
matches!(&n.kind,
crate::scope_tree::ScopeKind::OpTemplate { name } if name == "a")
})
.map(|(i, _)| i)
.expect("op-template node");
assert_eq!(tree.nodes[op_idx].materialised, Some(true));
}
#[test]
fn op_template_bare_name_metric_materialises() {
use nmbrs_workload::model::{MetricSpec, ParsedOp};
let mut phases = HashMap::new();
let mut p = empty_phase();
let mut op = ParsedOp::simple("a", "noop");
op.metrics.insert(
"m".into(),
MetricSpec {
cell: Default::default(),
value: "existing_wire".into(),
family: None,
kind: None,
unit: None,
format: None,
},
);
p.ops.push(op);
phases.insert("p".into(), p);
let mut tree = ScopeTree::build("default", &[ScenarioNode::Phase("p".into())]);
tree.extend_with_op_templates(&phases);
mark_with(&mut tree, &BindingsDef::default(), &HashMap::new(), &phases);
let op_idx = tree
.iter_dfs()
.find(|(_, n)| {
matches!(&n.kind,
crate::scope_tree::ScopeKind::OpTemplate { name } if name == "a")
})
.map(|(i, _)| i)
.expect("op-template node");
assert_eq!(tree.nodes[op_idx].materialised, Some(true));
}
#[test]
fn same_op_name_in_different_phases_classifies_per_phase() {
use nmbrs_workload::model::{MetricSpec, ParsedOp};
let mut with_metrics = empty_phase();
let mut op_with = ParsedOp::simple("select_ann", "noop");
op_with.metrics.insert(
"m".into(),
MetricSpec {
cell: Default::default(),
value: "existing_wire".into(),
family: None,
kind: None,
unit: None,
format: None,
},
);
with_metrics.ops.push(op_with);
let mut without_metrics = empty_phase();
without_metrics
.ops
.push(ParsedOp::simple("select_ann", "noop"));
let mut phases = HashMap::new();
phases.insert("ann_query".into(), with_metrics);
phases.insert("pvs_metadata_query".into(), without_metrics);
let mut tree = ScopeTree::build(
"default",
&[
ScenarioNode::Phase("ann_query".into()),
ScenarioNode::Phase("pvs_metadata_query".into()),
],
);
tree.extend_with_op_templates(&phases);
mark_with(&mut tree, &BindingsDef::default(), &HashMap::new(), &phases);
for (idx, node) in tree.iter_dfs() {
let crate::scope_tree::ScopeKind::OpTemplate { name } = &node.kind else {
continue;
};
if name != "select_ann" {
continue;
}
let mut cursor = node.parent;
let mut phase = None;
while let Some(p) = cursor {
if let crate::scope_tree::ScopeKind::Phase { name } = &tree.nodes[p].kind {
phase = Some(name.clone());
break;
}
cursor = tree.nodes[p].parent;
}
let expected_materialised = match phase.as_deref() {
Some("ann_query") => Some(true),
Some("pvs_metadata_query") => Some(false),
other => panic!("unexpected owning phase: {other:?}"),
};
assert_eq!(
tree.nodes[idx].materialised, expected_materialised,
"op-template {name} under phase {phase:?} should materialise={expected_materialised:?}"
);
}
}
#[test]
fn elision_summary_dumps_dfs_order() {
let mut phases = HashMap::new();
phases.insert("p".into(), empty_phase());
let mut tree = ScopeTree::build("default", &[ScenarioNode::Phase("p".into())]);
mark_with(&mut tree, &BindingsDef::default(), &HashMap::new(), &phases);
let summary = elision_summary(&tree);
assert_eq!(summary.len(), 4);
assert_eq!(summary[0].3, "");
assert_eq!(summary[1].3, "workload");
assert_eq!(summary[2].3, "workload.scenario.default");
assert_eq!(summary[3].3, "workload.scenario.default.phase.p");
for (_, _, mat, _, _) in &summary {
assert!(mat.is_some());
}
}
}