bjorn3
63 lines
1.5 KiB
Rust
63 lines
1.5 KiB
Rust
use std::collections::btree_map::{BTreeMap, Entry};
|
|
|
|
pub struct PortSet {
|
|
from: u16,
|
|
range: u16,
|
|
next: u16,
|
|
ports: BTreeMap<u16, usize>,
|
|
}
|
|
|
|
impl PortSet {
|
|
pub fn new(from: u16, to: u16) -> Option<PortSet> {
|
|
if from > to {
|
|
return None;
|
|
}
|
|
Some(PortSet {
|
|
from,
|
|
range: to - from + 1,
|
|
next: 0,
|
|
ports: BTreeMap::new(),
|
|
})
|
|
}
|
|
|
|
pub fn get_port(&mut self) -> Option<u16> {
|
|
if self.ports.len() >= self.range as usize {
|
|
return None;
|
|
}
|
|
|
|
let port = loop {
|
|
if let Entry::Vacant(entry) = self.ports.entry(self.next) {
|
|
entry.insert(1);
|
|
let port = self.from + self.next;
|
|
self.next = self.next.wrapping_add(1);
|
|
break port;
|
|
}
|
|
self.next = self.next.wrapping_add(1);
|
|
};
|
|
|
|
Some(port)
|
|
}
|
|
|
|
pub fn claim_port(&mut self, port: u16) -> bool {
|
|
if let Entry::Vacant(entry) = self.ports.entry(port) {
|
|
entry.insert(1);
|
|
true
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
|
|
pub fn acquire_port(&mut self, port: u16) {
|
|
*self.ports.entry(port).or_insert(0) += 1;
|
|
}
|
|
|
|
pub fn release_port(&mut self, port: u16) {
|
|
if let Entry::Occupied(mut entry) = self.ports.entry(port) {
|
|
*entry.get_mut() -= 1;
|
|
if *entry.get() == 0 {
|
|
entry.remove();
|
|
}
|
|
}
|
|
}
|
|
}
|