1use gbp_core::{ConformanceClass, MemberId};
9use std::collections::{BTreeSet, HashMap};
10
11#[derive(Default)]
13pub struct CapabilitiesNegotiator {
14 advertised: HashMap<MemberId, BTreeSet<String>>,
15}
16
17impl CapabilitiesNegotiator {
18 pub fn new() -> Self {
20 Self::default()
21 }
22
23 pub fn advertise<I, S>(&mut self, member: MemberId, capabilities: I)
26 where
27 I: IntoIterator<Item = S>,
28 S: Into<String>,
29 {
30 let set: BTreeSet<String> = capabilities.into_iter().map(Into::into).collect();
31 self.advertised.insert(member, set);
32 }
33
34 pub fn forget(&mut self, member: MemberId) {
36 self.advertised.remove(&member);
37 }
38
39 pub fn capabilities_of(&self, member: MemberId) -> Option<&BTreeSet<String>> {
41 self.advertised.get(&member)
42 }
43
44 pub fn group_supports(&self, cap: &str) -> bool {
46 if self.advertised.is_empty() {
47 return false;
48 }
49 self.advertised.values().all(|set| set.contains(cap))
50 }
51
52 pub fn intersection(&self) -> BTreeSet<String> {
55 let mut iter = self.advertised.values();
56 let Some(first) = iter.next() else {
57 return BTreeSet::new();
58 };
59 let mut acc = first.clone();
60 for set in iter {
61 acc.retain(|c| set.contains(c));
62 }
63 acc
64 }
65
66 pub fn union(&self) -> BTreeSet<String> {
68 let mut acc = BTreeSet::new();
69 for set in self.advertised.values() {
70 for c in set {
71 acc.insert(c.clone());
72 }
73 }
74 acc
75 }
76
77 pub fn missing(&self, cap: &str) -> Vec<MemberId> {
79 self.advertised
80 .iter()
81 .filter_map(|(m, set)| if set.contains(cap) { None } else { Some(*m) })
82 .collect()
83 }
84
85 pub fn len(&self) -> usize {
87 self.advertised.len()
88 }
89
90 pub fn is_empty(&self) -> bool {
92 self.advertised.is_empty()
93 }
94
95 pub fn declare_class(&mut self, member: MemberId, class: ConformanceClass) {
102 let entry = self.advertised.entry(member).or_default();
103 for token in class.tokens() {
104 entry.insert((*token).to_string());
105 }
106 }
107
108 pub fn group_class(&self) -> Option<ConformanceClass> {
111 if self.advertised.is_empty() {
112 return None;
113 }
114 let mut class: Option<ConformanceClass> = None;
116 for caps in self.advertised.values() {
117 let member_class = ConformanceClass::from_tokens(caps.iter().map(String::as_str));
118 class = Some(match (class, member_class) {
119 (None, mc) => mc?,
120 (Some(c), Some(mc)) => c.min(mc),
121 (Some(_), None) => return None,
122 });
123 }
124 class
125 }
126
127 pub fn reset_for_epoch(&mut self) {
131 self.advertised.clear();
132 }
133}
134
135#[cfg(test)]
136mod tests {
137 use super::*;
138
139 #[test]
140 fn intersection_is_lowest_common() {
141 let mut n = CapabilitiesNegotiator::new();
142 n.advertise(1, ["opus", "fec", "h264"]);
143 n.advertise(2, ["opus", "fec"]);
144 n.advertise(3, ["opus", "av1"]);
145 let common = n.intersection();
146 assert!(common.contains("opus"));
147 assert!(!common.contains("fec"));
148 assert!(!common.contains("h264"));
149 assert_eq!(n.missing("fec"), vec![3]);
150 }
151
152 #[test]
153 fn declare_class_inserts_tokens() {
154 let mut n = CapabilitiesNegotiator::new();
155 n.declare_class(1, ConformanceClass::C);
156 assert!(n.group_supports(ConformanceClass::TOKEN_A));
157 assert!(n.group_supports(ConformanceClass::TOKEN_B));
158 assert!(n.group_supports(ConformanceClass::TOKEN_C));
159 }
160
161 #[test]
162 fn group_class_returns_minimum() {
163 let mut n = CapabilitiesNegotiator::new();
164 n.declare_class(1, ConformanceClass::C);
165 n.declare_class(2, ConformanceClass::B);
166 n.declare_class(3, ConformanceClass::A);
167 assert_eq!(n.group_class(), Some(ConformanceClass::A));
168 }
169
170 #[test]
171 fn group_class_none_when_member_missing_tokens() {
172 let mut n = CapabilitiesNegotiator::new();
173 n.declare_class(1, ConformanceClass::B);
174 n.advertise(2, ["opus"]); assert_eq!(n.group_class(), None);
176 }
177
178 #[test]
179 fn group_class_none_when_empty() {
180 let n = CapabilitiesNegotiator::new();
181 assert_eq!(n.group_class(), None);
182 }
183
184 #[test]
185 fn group_supports_requires_everyone() {
186 let mut n = CapabilitiesNegotiator::new();
187 n.advertise(1, ["opus"]);
188 n.advertise(2, ["opus"]);
189 assert!(n.group_supports("opus"));
190 n.advertise(3, [] as [&str; 0]);
191 assert!(!n.group_supports("opus"));
192 }
193}