use crate::cot::types::{CapabilityAdvertisement, CapabilityInfo};
use crate::distribution::directive::CapabilityFilter;
use crate::models::domain::SensorType;
pub struct CapabilityMatcher;
impl CapabilityMatcher {
pub fn matches(adv: &CapabilityAdvertisement, filter: &CapabilityFilter) -> bool {
if let Some(req) = filter.min_gpu_memory_mb {
if !hardware_bound_satisfied(adv, |hw| hw.gpu_memory_mb, req) {
return false;
}
}
if let Some(req) = filter.min_memory_mb {
if !hardware_bound_satisfied(adv, |hw| hw.memory_mb, req) {
return false;
}
}
if let Some(req) = filter.min_storage_mb {
if !hardware_bound_satisfied(adv, |hw| hw.storage_mb, req) {
return false;
}
}
if !filter.custom.is_empty() {
let adv_custom = adv.hardware.as_ref().map(|hw| &hw.custom);
for (k, v) in &filter.custom {
let matches_value = adv_custom
.and_then(|c| c.get(k))
.map(|av| av == v)
.unwrap_or(false);
if !matches_value {
return false;
}
}
}
for req in &filter.required_capabilities {
if !capability_matches_any(req, &adv.capabilities) {
return false;
}
}
true
}
}
fn normalize_capability_string(s: &str) -> String {
s.trim().to_ascii_uppercase()
}
fn capability_matches_any(required: &str, advertised: &[CapabilityInfo]) -> bool {
let req_norm = normalize_capability_string(required);
if advertised
.iter()
.any(|c| normalize_capability_string(&c.capability_type) == req_norm)
{
return true;
}
if let Some(sensor) = sensor_type_from_string(&req_norm) {
let code = sensor.code();
if advertised
.iter()
.any(|c| c.capability_type.eq_ignore_ascii_case(code))
{
return true;
}
}
false
}
fn sensor_type_from_string(normalized: &str) -> Option<SensorType> {
Some(match normalized {
"EO" | "ELECTRO_OPTICAL" | "ELECTROOPTICAL" | "ELECTRO-OPTICAL" => {
SensorType::ElectroOptical
}
"IR" | "INFRARED" => SensorType::Infrared,
"RAD" | "RADAR" => SensorType::Radar,
"SON" | "SONAR" => SensorType::Sonar,
"ACO" | "ACOUSTIC" => SensorType::Acoustic,
"SIGINT" => SensorType::Sigint,
"MAD" => SensorType::Mad,
_ => return None,
})
}
fn hardware_bound_satisfied(
adv: &CapabilityAdvertisement,
extract: impl FnOnce(&crate::cot::types::HardwareSpec) -> Option<u64>,
req: u64,
) -> bool {
adv.hardware
.as_ref()
.and_then(extract)
.map(|have| have >= req)
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cot::types::{
CapabilityAdvertisement, CapabilityInfo, HardwareSpec, OperationalStatus, Position,
};
use std::collections::HashMap;
fn make_advert(caps: Vec<&str>) -> CapabilityAdvertisement {
let mut adv = CapabilityAdvertisement::new(
"p-1".into(),
"UAV".into(),
Position::new(0.0, 0.0),
OperationalStatus::Active,
1.0,
);
for ct in caps {
adv = adv.with_capability(CapabilityInfo {
capability_type: ct.to_string(),
model_name: "test".into(),
version: "1.0".into(),
precision: 1.0,
status: OperationalStatus::Active,
});
}
adv
}
fn empty_filter() -> CapabilityFilter {
CapabilityFilter::default()
}
#[test]
fn empty_filter_matches_any_advert() {
let adv = make_advert(vec![]);
assert!(CapabilityMatcher::matches(&adv, &empty_filter()));
}
#[test]
fn direct_capability_string_match_case_insensitive() {
let adv = make_advert(vec!["OBJECT_TRACKING"]);
let filter = CapabilityFilter {
required_capabilities: vec!["object_tracking".into()],
..empty_filter()
};
assert!(CapabilityMatcher::matches(&adv, &filter));
}
#[test]
fn required_capability_missing_fails() {
let adv = make_advert(vec!["COMPUTE"]);
let filter = CapabilityFilter {
required_capabilities: vec!["OBJECT_TRACKING".into()],
..empty_filter()
};
assert!(!CapabilityMatcher::matches(&adv, &filter));
}
#[test]
fn all_required_capabilities_must_match() {
let adv = make_advert(vec!["COMPUTE", "OBJECT_TRACKING"]);
let filter_partial_present = CapabilityFilter {
required_capabilities: vec!["COMPUTE".into(), "MISSING".into()],
..empty_filter()
};
assert!(!CapabilityMatcher::matches(&adv, &filter_partial_present));
let filter_all_present = CapabilityFilter {
required_capabilities: vec!["COMPUTE".into(), "OBJECT_TRACKING".into()],
..empty_filter()
};
assert!(CapabilityMatcher::matches(&adv, &filter_all_present));
}
#[test]
fn sensor_type_bridge_matches_short_code_advert_via_long_name_filter() {
let adv = make_advert(vec!["EO"]);
let filter = CapabilityFilter {
required_capabilities: vec!["ELECTRO_OPTICAL".into()],
..empty_filter()
};
assert!(CapabilityMatcher::matches(&adv, &filter));
}
#[test]
fn sensor_type_bridge_matches_all_seven_sensors_via_long_aliases() {
let pairs = [
("EO", "electro_optical"),
("IR", "infrared"),
("RAD", "radar"),
("SON", "sonar"),
("ACO", "acoustic"),
("SIGINT", "sigint"),
("MAD", "mad"),
];
for (advert_code, filter_alias) in pairs {
let adv = make_advert(vec![advert_code]);
let filter = CapabilityFilter {
required_capabilities: vec![filter_alias.into()],
..empty_filter()
};
assert!(
CapabilityMatcher::matches(&adv, &filter),
"expected {advert_code} advert to match filter alias '{filter_alias}'"
);
}
}
#[test]
fn hardware_bound_missing_advert_field_fails() {
let adv = make_advert(vec![]);
let filter = CapabilityFilter {
min_gpu_memory_mb: Some(8_192),
..empty_filter()
};
assert!(!CapabilityMatcher::matches(&adv, &filter));
}
#[test]
fn hardware_bound_with_present_field_satisfied() {
let adv = make_advert(vec![]).with_hardware(HardwareSpec {
gpu_memory_mb: Some(16_384),
memory_mb: Some(32_768),
storage_mb: Some(512_000),
custom: HashMap::new(),
});
let filter = CapabilityFilter {
min_gpu_memory_mb: Some(8_192),
min_memory_mb: Some(16_384),
min_storage_mb: Some(100_000),
..empty_filter()
};
assert!(CapabilityMatcher::matches(&adv, &filter));
}
#[test]
fn hardware_bound_below_required_fails() {
let adv = make_advert(vec![]).with_hardware(HardwareSpec {
gpu_memory_mb: Some(4_096),
..Default::default()
});
let filter = CapabilityFilter {
min_gpu_memory_mb: Some(8_192),
..empty_filter()
};
assert!(!CapabilityMatcher::matches(&adv, &filter));
}
#[test]
fn hardware_bound_at_exactly_required_passes() {
let adv = make_advert(vec![]).with_hardware(HardwareSpec {
memory_mb: Some(16_384),
..Default::default()
});
let filter = CapabilityFilter {
min_memory_mb: Some(16_384),
..empty_filter()
};
assert!(CapabilityMatcher::matches(&adv, &filter));
}
#[test]
fn custom_field_match_and_mismatch() {
let mut adv_custom = HashMap::new();
adv_custom.insert("cpu_arch".to_string(), "aarch64".to_string());
adv_custom.insert("tensorrt_version".to_string(), "10.0".to_string());
let adv = make_advert(vec![]).with_hardware(HardwareSpec {
custom: adv_custom,
..Default::default()
});
let mut filter_match = empty_filter();
filter_match
.custom
.insert("cpu_arch".into(), "aarch64".into());
assert!(CapabilityMatcher::matches(&adv, &filter_match));
let mut filter_value_mismatch = empty_filter();
filter_value_mismatch
.custom
.insert("cpu_arch".into(), "x86_64".into());
assert!(!CapabilityMatcher::matches(&adv, &filter_value_mismatch));
let mut filter_key_missing = empty_filter();
filter_key_missing
.custom
.insert("gpu_compute_capability".into(), "8.9".into());
assert!(!CapabilityMatcher::matches(&adv, &filter_key_missing));
}
#[test]
fn custom_field_advert_has_no_hardware_fails() {
let adv = make_advert(vec![]);
let mut filter = empty_filter();
filter.custom.insert("cpu_arch".into(), "aarch64".into());
assert!(!CapabilityMatcher::matches(&adv, &filter));
}
#[test]
fn combined_filter_all_constraints_must_pass() {
let mut adv_custom = HashMap::new();
adv_custom.insert("cpu_arch".to_string(), "aarch64".to_string());
let adv = make_advert(vec!["OBJECT_TRACKING", "EO"]).with_hardware(HardwareSpec {
gpu_memory_mb: Some(16_384),
memory_mb: Some(32_768),
custom: adv_custom,
..Default::default()
});
let mut filter_pass = CapabilityFilter {
min_gpu_memory_mb: Some(8_192),
min_memory_mb: Some(16_384),
required_capabilities: vec!["OBJECT_TRACKING".into(), "ELECTRO_OPTICAL".into()],
..empty_filter()
};
filter_pass
.custom
.insert("cpu_arch".into(), "aarch64".into());
assert!(CapabilityMatcher::matches(&adv, &filter_pass));
let mut filter_break_gpu = filter_pass.clone();
filter_break_gpu.min_gpu_memory_mb = Some(32_768);
assert!(!CapabilityMatcher::matches(&adv, &filter_break_gpu));
let mut filter_break_cap = filter_pass.clone();
filter_break_cap
.required_capabilities
.push("NONEXISTENT".into());
assert!(!CapabilityMatcher::matches(&adv, &filter_break_cap));
let mut filter_break_custom = filter_pass.clone();
filter_break_custom
.custom
.insert("cpu_arch".into(), "x86_64".into());
assert!(!CapabilityMatcher::matches(&adv, &filter_break_custom));
}
}