sys_rs/coverage.rs
1use nix::errno::Errno;
2use std::{
3 collections::{hash_map::Entry, HashMap, HashSet},
4 path::Path,
5};
6
7use crate::{
8 asm::Instruction,
9 debug::{Dwarf, LineInfo},
10 diag::{Error, Result},
11 print::{self, Layout},
12 process,
13 profile::{trace_with, Tracer},
14 progress::{self, ProgressFn, State},
15};
16
17/// Coverage collector that maintains an in-memory mapping from source
18/// locations to execution counts.
19///
20/// `Cached` keeps a small cache mapping instruction addresses to source
21/// `LineInfo` (if available) and a `coverage` map counting visits per file
22/// and line. It also tracks the set of files seen.
23pub struct Cached {
24 cache: HashMap<u64, Option<LineInfo>>,
25 coverage: HashMap<(String, usize), usize>,
26 files: HashSet<String>,
27}
28
29impl Cached {
30 #[must_use]
31 /// Create a new, empty `Cached` collector.
32 ///
33 /// # Returns
34 ///
35 /// A new, empty `Cached` instance.
36 pub fn new() -> Self {
37 Self {
38 cache: HashMap::new(),
39 coverage: HashMap::new(),
40 files: HashSet::new(),
41 }
42 }
43
44 #[must_use]
45 /// Get the coverage count for `path:line` if present.
46 ///
47 /// # Arguments
48 ///
49 /// * `path` - Source file path.
50 /// * `line` - Line number in the source file.
51 ///
52 /// # Returns
53 ///
54 /// `Some(&usize)` with the execution count when present, otherwise
55 /// `None`.
56 pub fn coverage(&self, path: String, line: usize) -> Option<&usize> {
57 let key = (path, line);
58 self.coverage.get(&key)
59 }
60
61 #[must_use]
62 /// Return the set of files observed by the collector.
63 ///
64 /// # Returns
65 ///
66 /// A reference to the `HashSet` of file paths that have been observed.
67 pub fn files(&self) -> &HashSet<String> {
68 &self.files
69 }
70
71 /// Print the source line corresponding to `instruction` when available.
72 ///
73 /// This looks up the DWARF line info for the instruction address and
74 /// prints the source line if the file exists on disk. When `record` is
75 /// true the collector updates its internal coverage counts and file set.
76 ///
77 /// # Arguments
78 ///
79 /// * `instruction` - The disassembled instruction whose source line to print.
80 /// * `dwarf` - DWARF helper used to resolve addresses to source lines.
81 /// * `record` - When true the collector records coverage counts for the
82 /// located source line.
83 ///
84 /// # Errors
85 ///
86 /// Returns an error if DWARF lookup fails.
87 pub fn print_source(
88 &mut self,
89 instruction: &Instruction,
90 dwarf: &Dwarf,
91 record: bool,
92 ) -> Result<Option<String>> {
93 let mut ret = None;
94 let addr = instruction.address();
95 if let Entry::Vacant(_) = self.cache.entry(addr) {
96 let info = dwarf.addr2line(addr)?;
97 self.cache.insert(addr, info);
98 }
99
100 if let Some(line) = self
101 .cache
102 .get(&addr)
103 .ok_or_else(|| Error::from(Errno::ENODATA))?
104 {
105 if Path::new(&line.path()).exists() {
106 if record {
107 let key = (line.path(), line.line());
108 *self.coverage.entry(key).or_insert(0) += 1;
109 self.files.insert(line.path());
110 }
111 let output = format!("{line}");
112 println!("{output}");
113 ret = Some(output);
114 }
115 }
116 Ok(ret)
117 }
118
119 /// Run the tracer and print source lines when available.
120 ///
121 /// This convenience wraps `trace_with` using the collector's
122 /// `print_source` formatter so each executed instruction prints a
123 /// source line when the DWARF information is present.
124 ///
125 /// # Arguments
126 ///
127 /// * `context` - The tracer implementation used to run the program.
128 /// * `process` - Process metadata for the target binary.
129 ///
130 /// # Errors
131 ///
132 /// Returns an error if DWARF construction fails or if tracing fails.
133 pub fn trace_with_source_print(
134 &mut self,
135 context: &Tracer,
136 process: &process::Info,
137 ) -> Result<i32> {
138 let dwarf = Dwarf::build(process)?;
139 let state = State::new(process.pid(), Some(&dwarf));
140 trace_with(
141 context,
142 process,
143 state,
144 |instruction, _| self.print_source(instruction, &dwarf, false),
145 progress::default,
146 )
147 }
148
149 /// Run the tracer using a custom progress function and optionally print
150 /// source lines.
151 ///
152 /// If `dwarf` is `Some`, the tracer will attempt to print source lines
153 /// (when the layout indicates source); otherwise it falls back to the
154 /// default printer. When `record` is true coverage counts are collected.
155 /// The provided `progress` function is used for user interaction between
156 /// instructions.
157 ///
158 /// # Arguments
159 ///
160 /// * `context` - The tracer implementation used to run the program.
161 /// * `process` - Process metadata for the target binary.
162 /// * `dwarf` - Optional DWARF helper; when `Some` source printing is enabled.
163 /// * `record` - If true, coverage counts are recorded.
164 /// * `progress` - Custom progress function called between instructions.
165 ///
166 /// # Errors
167 ///
168 /// Returns an error if tracing or DWARF operations fail.
169 pub fn trace_with_custom_progress(
170 &mut self,
171 context: &Tracer,
172 process: &process::Info,
173 dwarf: Option<&Dwarf>,
174 record: bool,
175 progress: impl ProgressFn,
176 ) -> Result<i32> {
177 let src_available = dwarf.is_some();
178 let state = State::new(process.pid(), dwarf);
179 trace_with(
180 context,
181 process,
182 state,
183 |instruction, layout| match (Layout::from(src_available), layout) {
184 (Layout::Source, Layout::Source) => self.print_source(
185 instruction,
186 dwarf.ok_or_else(|| Error::from(Errno::ENODATA))?,
187 record,
188 ),
189 _ => print::default(instruction, layout),
190 },
191 progress,
192 )
193 }
194
195 /// Run the tracer with the default progress function and recording
196 /// enabled.
197 ///
198 /// This is the standard entry point for coverage collection: it tries
199 /// to build DWARF info and then calls `trace_with_custom_progress`.
200 ///
201 /// # Arguments
202 ///
203 /// * `context` - The tracer implementation used to run the program.
204 /// * `process` - Process metadata for the target binary.
205 ///
206 /// # Errors
207 ///
208 /// Returns an error if DWARF construction or tracing fails.
209 pub fn trace_with_default_progress(
210 &mut self,
211 context: &Tracer,
212 process: &process::Info,
213 ) -> Result<i32> {
214 let dwarf = Dwarf::build(process);
215 self.trace_with_custom_progress(
216 context,
217 process,
218 dwarf.as_ref().ok(),
219 true,
220 progress::default,
221 )
222 }
223}
224
225impl Default for Cached {
226 fn default() -> Self {
227 Self::new()
228 }
229}
230
231#[cfg(test)]
232mod tests {
233 use super::*;
234
235 #[test]
236 fn test_cached_new() {
237 let cached = Cached::new();
238 assert_eq!(cached.coverage.len(), 0);
239 assert_eq!(cached.files.len(), 0);
240 }
241
242 #[test]
243 fn test_cached_coverage() {
244 let mut cached = Cached::new();
245 cached.coverage.insert(("file1".to_string(), 10), 5);
246 cached.coverage.insert(("file2".to_string(), 20), 10);
247
248 assert_eq!(cached.coverage("file1".to_string(), 10), Some(&5));
249 assert_eq!(cached.coverage("file2".to_string(), 20), Some(&10));
250 assert_eq!(cached.coverage("file3".to_string(), 30), None);
251 }
252
253 #[test]
254 fn test_cached_files() {
255 let mut cached = Cached::new();
256 cached.files.insert("file1".to_string());
257 cached.files.insert("file2".to_string());
258
259 assert_eq!(cached.files().len(), 2);
260 assert!(cached.files().contains("file1"));
261 assert!(cached.files().contains("file2"));
262 assert!(!cached.files().contains("file3"));
263 }
264}