Extract trait SymbolResolver

aya/src/maps/stack_trace.rs

@@ -1,7 +1,12 @@
 //! A hash map of kernel or user space stack traces.
 //!
 //! See [`StackTraceMap`] for documentation and examples.
-use std::{borrow::Borrow, collections::BTreeMap, fs, io, mem, path::Path, str::FromStr};
+use std::{
+    borrow::{Borrow, Cow},
+    fs, io, mem,
+    path::Path,
+    str::FromStr,
+};
 
 use crate::{
     maps::{IterableMap, MapData, MapError, MapIter, MapKeys},
@@ -157,6 +162,12 @@ impl<'a, T: Borrow<MapData>> IntoIterator for &'a StackTraceMap<T> {
     }
 }
 
+/// A resolver for symbols based on an address obtained from a stack trace.
+pub trait SymbolResolver {
+    /// Resolve a symbol for a given address, if possible.
+    fn resolve_symbol(&self, addr: u64) -> Option<Cow<'_, str>>;
+}
+
 /// A kernel or user space stack trace.
 ///
 /// See the [`StackTraceMap`] documentation for examples.
@@ -172,12 +183,9 @@ impl StackTrace {
     /// You can use [`util::kernel_symbols()`](crate::util::kernel_symbols) to load kernel symbols. For
     /// user-space traces you need to provide the symbols, for example loading
     /// them from debug info.
-    pub fn resolve(&mut self, symbols: &BTreeMap<u64, String>) -> &StackTrace {
+    pub fn resolve<R: SymbolResolver>(&mut self, symbols: &R) -> &StackTrace {
         for frame in self.frames.iter_mut() {
-            frame.symbol_name = symbols
-                .range(..=frame.ip)
-                .next_back()
-                .map(|(_, s)| s.clone())
+            frame.symbol_name = symbols.resolve_symbol(frame.ip).map(|s| s.into_owned())
         }
 
         self

aya/src/util.rs

@@ -1,5 +1,6 @@
 //! Utility functions.
 use std::{
+    borrow::Cow,
     collections::BTreeMap,
     error::Error,
     ffi::{CStr, CString},
@@ -11,6 +12,7 @@ use std::{
 
 use crate::{
     generated::{TC_H_MAJ_MASK, TC_H_MIN_MASK},
+    maps::stack_trace::SymbolResolver,
     Pod,
 };
 
@@ -201,16 +203,25 @@ fn parse_cpu_ranges(data: &str) -> Result<Vec<u32>, ()> {
     Ok(cpus)
 }
 
+/// The simplest resolver: a direct map from addresses to strings.
+pub type SimpleSymbolResolver = BTreeMap<u64, String>;
+
+impl SymbolResolver for SimpleSymbolResolver {
+    fn resolve_symbol(&self, addr: u64) -> Option<Cow<'_, str>> {
+        self.range(..=addr).next_back().map(|(_, s)| s.into())
+    }
+}
+
 /// Loads kernel symbols from `/proc/kallsyms`.
 ///
 /// The symbols can be passed to [`StackTrace::resolve`](crate::maps::stack_trace::StackTrace::resolve).
-pub fn kernel_symbols() -> Result<BTreeMap<u64, String>, io::Error> {
+pub fn kernel_symbols() -> Result<SimpleSymbolResolver, io::Error> {
     let mut reader = BufReader::new(File::open("/proc/kallsyms")?);
     parse_kernel_symbols(&mut reader)
 }
 
-fn parse_kernel_symbols(reader: impl BufRead) -> Result<BTreeMap<u64, String>, io::Error> {
-    let mut syms = BTreeMap::new();
+fn parse_kernel_symbols(reader: impl BufRead) -> Result<SimpleSymbolResolver, io::Error> {
+    let mut syms = SimpleSymbolResolver::new();
 
     for line in reader.lines() {
         let line = line?;

xtask/public-api/aya.txt

@@ -734,7 +734,7 @@ pub struct aya::maps::stack_trace::StackTrace
 pub aya::maps::stack_trace::StackTrace::id: u32
 impl aya::maps::stack_trace::StackTrace
 pub fn aya::maps::stack_trace::StackTrace::frames(&self) -> &[aya::maps::stack_trace::StackFrame]
-pub fn aya::maps::stack_trace::StackTrace::resolve(&mut self, symbols: &alloc::collections::btree::map::BTreeMap<u64, alloc::string::String>) -> &aya::maps::stack_trace::StackTrace
+pub fn aya::maps::stack_trace::StackTrace::resolve<R: aya::maps::stack_trace::SymbolResolver>(&mut self, symbols: &R) -> &aya::maps::stack_trace::StackTrace
 impl<T: core::borrow::Borrow<aya::maps::MapData>> aya::maps::IterableMap<u32, aya::maps::stack_trace::StackTrace> for aya::maps::stack_trace::StackTraceMap<T>
 pub fn aya::maps::stack_trace::StackTraceMap<T>::get(&self, index: &u32) -> core::result::Result<aya::maps::stack_trace::StackTrace, aya::maps::MapError>
 pub fn aya::maps::stack_trace::StackTraceMap<T>::map(&self) -> &aya::maps::MapData
@@ -803,6 +803,10 @@ impl<T> core::borrow::BorrowMut<T> for aya::maps::stack_trace::StackTraceMap<T>
 pub fn aya::maps::stack_trace::StackTraceMap<T>::borrow_mut(&mut self) -> &mut T
 impl<T> core::convert::From<T> for aya::maps::stack_trace::StackTraceMap<T>
 pub fn aya::maps::stack_trace::StackTraceMap<T>::from(t: T) -> T
+pub trait aya::maps::stack_trace::SymbolResolver
+pub fn aya::maps::stack_trace::SymbolResolver::resolve_symbol(&self, addr: u64) -> core::option::Option<alloc::borrow::Cow<'_, str>>
+impl aya::maps::stack_trace::SymbolResolver for aya::util::SimpleSymbolResolver
+pub fn aya::util::SimpleSymbolResolver::resolve_symbol(&self, addr: u64) -> core::option::Option<alloc::borrow::Cow<'_, str>>
 pub enum aya::maps::Map
 pub aya::maps::Map::Array(aya::maps::MapData)
 pub aya::maps::Map::BloomFilter(aya::maps::MapData)
@@ -6933,10 +6937,11 @@ impl<T> core::borrow::BorrowMut<T> for aya::util::KernelVersion where T: core::m
 pub fn aya::util::KernelVersion::borrow_mut(&mut self) -> &mut T
 impl<T> core::convert::From<T> for aya::util::KernelVersion
 pub fn aya::util::KernelVersion::from(t: T) -> T
-pub fn aya::util::kernel_symbols() -> core::result::Result<alloc::collections::btree::map::BTreeMap<u64, alloc::string::String>, std::io::error::Error>
+pub fn aya::util::kernel_symbols() -> core::result::Result<aya::util::SimpleSymbolResolver, std::io::error::Error>
 pub fn aya::util::nr_cpus() -> core::result::Result<usize, std::io::error::Error>
 pub fn aya::util::online_cpus() -> core::result::Result<alloc::vec::Vec<u32>, std::io::error::Error>
 pub fn aya::util::syscall_prefix() -> core::result::Result<&'static str, std::io::error::Error>
+pub type aya::util::SimpleSymbolResolver = alloc::collections::btree::map::BTreeMap<u64, alloc::string::String>
 pub macro aya::include_bytes_aligned!
 pub enum aya::BpfError
 pub aya::BpfError::BtfError(aya_obj::btf::btf::BtfError)