From d9c1689066996158314da6c4f072844732fac1e0 Mon Sep 17 00:00:00 2001
From: Tamir Duberstein <tamird@gmail.com>
Date: Tue, 1 Oct 2024 10:13:50 -0400
Subject: [PATCH] Avoid intermediate allocations in parse_cpu_ranges
---
aya/src/util.rs | 42 +++++++++++++++++++++++-------------------
1 file changed, 23 insertions(+), 19 deletions(-)
diff --git a/aya/src/util.rs b/aya/src/util.rs
index 10fc15ef..28efe3f0 100644
--- a/aya/src/util.rs
+++ b/aya/src/util.rs
@@ -202,30 +202,34 @@ pub fn nr_cpus() -> Result<usize, (&'static str, io::Error)> {
fn read_cpu_ranges(path: &'static str) -> Result<Vec<u32>, (&'static str, io::Error)> {
(|| {
let data = fs::read_to_string(path)?;
- parse_cpu_ranges(&data).map_err(|range| io::Error::new(io::ErrorKind::InvalidData, range))
+ parse_cpu_ranges(&data)
})()
.map_err(|error| (path, error))
}
-fn parse_cpu_ranges(data: &str) -> Result<Vec<u32>, &str> {
- let mut cpus = Vec::new();
- for range in data.split(',') {
- cpus.extend({
- match range
- .splitn(2, '-')
- .map(u32::from_str)
- .collect::<Result<Vec<_>, _>>()
- .map_err(|ParseIntError { .. }| range)?
- .as_slice()
- {
- &[] | &[_, _, _, ..] => return Err(range),
- &[start] => start..=start,
- &[start, end] => start..=end,
- }
+fn parse_cpu_ranges(data: &str) -> Result<Vec<u32>, io::Error> {
+ data.split(',')
+ .map(|range| {
+ let mut iter = range
+ .split('-')
+ .map(|s| s.parse::<u32>().map_err(|ParseIntError { .. }| range));
+ let start = iter.next().unwrap()?; // str::split always returns at least one element.
+ let end = match iter.next() {
+ None => start,
+ Some(end) => {
+ if iter.next().is_some() {
+ return Err(range);
+ }
+ end?
+ }
+ };
+ Ok(start..=end)
+ })
+ .try_fold(Vec::new(), |mut cpus, range| {
+ let range = range.map_err(|range| io::Error::new(io::ErrorKind::InvalidData, range))?;
+ cpus.extend(range);
+ Ok(cpus)
})
- }
-
- Ok(cpus)
}
/// Loads kernel symbols from `/proc/kallsyms`.