1pub struct RangeListIter<'a> {
2 current_range: Option<std::ops::RangeInclusive<u32>>,
3 remaining: &'a str,
4}
5
6impl Iterator for RangeListIter<'_> {
7 type Item = u32;
8
9 fn next(&mut self) -> Option<Self::Item> {
10 loop {
11 if let Some(current_range) = &mut self.current_range {
13 let Some(rv) = current_range.next() else {
15 self.current_range = None;
17 continue;
18 };
19
20 break Some(rv);
22 } else if !self.remaining.is_empty() {
24 let (next_range, remaining) = match self.remaining.split_once(',') {
25 Some(x) => x,
27 None => (self.remaining, ""),
29 };
30
31 self.remaining = remaining;
32
33 if next_range.is_empty() {
34 continue;
35 }
36
37 let mut split = next_range.split('-');
38 let start = split.next().unwrap();
39 let end = split.next();
40 assert!(split.next().is_none());
41
42 let start = start.parse().unwrap();
43 let end = end.map(|x| x.parse().unwrap()).unwrap_or(start);
44
45 self.current_range = Some(std::ops::RangeInclusive::new(start, end));
46
47 continue;
48 } else {
50 break None;
52 }
53 }
54 }
55}
56
57pub fn parse_range_list(range_list: &str) -> RangeListIter {
67 RangeListIter {
68 current_range: None,
69 remaining: range_list,
70 }
71}
72
73pub fn nodes() -> Vec<u32> {
75 let name = "/sys/devices/system/node/possible";
76 parse_range_list(std::fs::read_to_string(name).unwrap().trim()).collect()
77}
78
79pub fn cpus(node: u32) -> Vec<u32> {
81 let name = format!("/sys/devices/system/node/node{node}/cpulist");
82 parse_range_list(std::fs::read_to_string(name).unwrap().trim()).collect()
83}
84
85pub fn core(cpu: u32) -> u32 {
87 let name = format!("/sys/devices/system/cpu/cpu{cpu}/topology/core_id");
88 std::fs::read_to_string(name)
89 .unwrap()
90 .trim()
91 .parse()
92 .unwrap()
93}
94
95pub fn online() -> Vec<u32> {
97 let name = "/sys/devices/system/cpu/online";
98 parse_range_list(std::fs::read_to_string(name).unwrap().trim()).collect()
99}
100
101pub fn count_physical_cores() -> u32 {
104 let affinity = nix::sched::sched_getaffinity(nix::unistd::Pid::from_raw(0)).unwrap();
105
106 let mut physical_cores = std::collections::HashSet::new();
107
108 for cpu in online() {
109 if affinity.is_set(cpu.try_into().unwrap()).unwrap() {
110 physical_cores.insert(core(cpu));
111 }
112 }
113
114 assert!(!physical_cores.is_empty());
115 physical_cores.len().try_into().unwrap()
116}
117
118#[cfg(test)]
119mod tests {
120 use super::*;
121
122 fn check(list: &str, array: &[u32]) {
123 let list: Vec<_> = parse_range_list(list).collect();
124 assert_eq!(list, array);
125 }
126
127 #[test]
128 fn test_range_list() {
129 check("", &[]);
130 check("1", &[1]);
131 check("1,2", &[1, 2]);
132 check("1-2", &[1, 2]);
133 check("1-1", &[1]);
134 check("1,2,3", &[1, 2, 3]);
135 check("1-3", &[1, 2, 3]);
136 check("1,2-3,4", &[1, 2, 3, 4]);
137 check("1,2-4,5", &[1, 2, 3, 4, 5]);
138 check(
139 "0-5,7-9,13,15-19",
140 &[0, 1, 2, 3, 4, 5, 7, 8, 9, 13, 15, 16, 17, 18, 19],
141 );
142 check("1,,5", &[1, 5]);
143 check("1,1,5", &[1, 1, 5]);
144 check("1-1,5", &[1, 5]);
145 check("1-1,0,5", &[1, 0, 5]);
146 check("1-0", &[]);
147 }
148}