use anyhow::format_err;
use wlan_common::mac::{Aid, MAX_AID};
#[derive(Debug)]
pub struct Map {
aids: [u64; 32],
}
impl Map {
const ELEM_BITS: u16 = 64;
pub fn assign_aid(&mut self) -> Result<Aid, anyhow::Error> {
for (i, bitmap) in self.aids.iter_mut().enumerate() {
if bitmap.count_zeros() > 0 {
let first_unset_bit_pos = (!*bitmap).trailing_zeros() as u16;
let aid = first_unset_bit_pos + Map::ELEM_BITS * (i as u16);
if aid <= MAX_AID {
*bitmap |= 1 << first_unset_bit_pos;
return Ok(aid);
} else {
return Err(format_err!("no available association ID"));
}
}
}
panic!("unexpected error assigning association ID")
}
pub fn release_aid(&mut self, aid: Aid) {
let index = (aid / Map::ELEM_BITS) as usize;
self.aids[index] &= !(1 << (aid % Map::ELEM_BITS));
}
}
impl Default for Map {
fn default() -> Self {
let mut map = Map { aids: [0u64; 32] };
map.aids[0] = 1;
map
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_map_never_assign_zero() {
let mut aid_map: Map = Default::default();
for i in 1..=2007u16 {
assert_eq!(aid_map.assign_aid().unwrap(), i);
}
let result = aid_map.assign_aid();
assert!(result.is_err());
assert_eq!(format!("{}", result.unwrap_err()), "no available association ID");
}
#[allow(
clippy::legacy_numeric_constants,
reason = "mass allow for https://fxbug.dev/381896734"
)]
#[test]
fn test_map_no_available_assoc_id() {
let mut aid_map: Map = Default::default();
for i in 0..31 {
aid_map.aids[i] = u64::max_value();
}
for i in 0..24 {
aid_map.aids[31] += 1 << i;
}
let result = aid_map.assign_aid();
assert!(result.is_err());
assert_eq!(format!("{}", result.unwrap_err()), "no available association ID");
}
#[test]
fn test_map_ascending_available() {
let mut aid_map: Map = Default::default();
for i in 1..=1000u16 {
assert_eq!(aid_map.assign_aid().unwrap(), i);
}
aid_map.release_aid(157);
aid_map.release_aid(792);
aid_map.release_aid(533);
assert_eq!(aid_map.assign_aid().unwrap(), 157);
assert_eq!(aid_map.assign_aid().unwrap(), 533);
assert_eq!(aid_map.assign_aid().unwrap(), 792);
for i in 1001..=2007u16 {
assert_eq!(aid_map.assign_aid().unwrap(), i);
}
let result = aid_map.assign_aid();
assert!(result.is_err());
assert_eq!(format!("{}", result.unwrap_err()), "no available association ID");
aid_map.release_aid(666);
aid_map.release_aid(222);
aid_map.release_aid(111);
assert_eq!(aid_map.assign_aid().unwrap(), 111);
assert_eq!(aid_map.assign_aid().unwrap(), 222);
assert_eq!(aid_map.assign_aid().unwrap(), 666);
}
}