use crate::attrs::num_attr;
use cel::{Context, Program};
use code_ranker_plugin_api::{attrs::AttrValue, node::Node};
use std::collections::BTreeMap;
mod model;
pub use model::{MetricDef, RegistryError, Scope};
pub(crate) use model::{default_omit_at, default_value_type};
mod eval;
pub use eval::Populations;
use eval::{exec_f64, register_agg, topo_order};
pub(crate) use eval::{references, register_math};
#[cfg(test)]
use eval::{percentile, reduce};
pub fn apply_to_node(node: &mut Node, defs: &BTreeMap<String, MetricDef>, engine: &Engine) {
let mut attrs: BTreeMap<String, f64> = BTreeMap::new();
let mut strings: BTreeMap<String, String> = BTreeMap::new();
for (k, v) in &node.attrs {
match v {
AttrValue::Int(i) => {
attrs.insert(k.clone(), *i as f64);
}
AttrValue::Float(f) => {
attrs.insert(k.clone(), *f);
}
AttrValue::Str(s) => {
strings.insert(k.clone(), s.clone());
}
AttrValue::Bool(_) => {}
}
}
for (k, v) in crate::nodepath::path_fields(node) {
strings.insert(k.to_string(), v);
}
for (key, value) in engine.eval_node(&attrs, &strings) {
let omit = defs.get(&key).map(|d| d.omit_at).unwrap_or(0.0);
let a = num_attr(value);
if a == num_attr(omit) {
node.attrs.remove(&key);
} else {
node.attrs.insert(key, a);
}
}
}
pub struct Engine {
node_programs: Vec<(String, Program)>,
graph_programs: Vec<(String, Program)>,
}
impl Engine {
pub fn compile(defs: &BTreeMap<String, MetricDef>) -> Result<Engine, RegistryError> {
Ok(Engine {
node_programs: compile_scope(defs, Scope::Node)?,
graph_programs: compile_scope(defs, Scope::Graph)?,
})
}
pub fn has_graph_metrics(&self) -> bool {
!self.graph_programs.is_empty()
}
pub fn eval_node(
&self,
attrs: &BTreeMap<String, f64>,
strings: &BTreeMap<String, String>,
) -> BTreeMap<String, f64> {
let mut ctx = Context::default();
register_math(&mut ctx);
for (k, v) in attrs {
let _ = ctx.add_variable(k.as_str(), *v);
}
for (k, v) in strings {
let _ = ctx.add_variable(k.as_str(), v.clone());
}
let mut produced: BTreeMap<String, f64> = BTreeMap::new();
for (key, program) in &self.node_programs {
if let Some(v) = exec_f64(program, &ctx) {
let _ = ctx.add_variable(key.as_str(), v); produced.insert(key.clone(), v);
}
}
produced
}
pub fn eval_graph(&self, pops: &Populations) -> BTreeMap<String, f64> {
let mut ctx = Context::default();
register_math(&mut ctx);
register_agg(&mut ctx, std::sync::Arc::new(pops.clone()));
let mut produced: BTreeMap<String, f64> = BTreeMap::new();
for (key, program) in &self.graph_programs {
if let Some(v) = exec_f64(program, &ctx) {
let _ = ctx.add_variable(key.as_str(), v); produced.insert(key.clone(), v);
}
}
produced
}
}
fn compile_scope(
defs: &BTreeMap<String, MetricDef>,
scope: Scope,
) -> Result<Vec<(String, Program)>, RegistryError> {
let scoped: Vec<(&String, &MetricDef)> =
defs.iter().filter(|(_, d)| d.scope == scope).collect();
let keys: Vec<String> = scoped.iter().map(|(k, _)| (*k).clone()).collect();
let order = topo_order(&scoped, &keys)?;
let mut programs = Vec::with_capacity(order.len());
for key in order {
let def = defs.get(&key).expect("key from defs");
let program = Program::compile(&def.formula_cel).map_err(|e| RegistryError::Parse {
key: key.clone(),
message: e.to_string(),
})?;
programs.push((key, program));
}
Ok(programs)
}
#[cfg(test)]
#[path = "registry_test.rs"]
mod tests;
#[cfg(test)]
#[path = "registry_cover_test.rs"]
mod cover_tests;