polkavm/
source_cache.rs

1use std::borrow::Cow;
2use std::collections::HashMap;
3use std::path::{Path, PathBuf};
4
5#[derive(Default)]
6pub struct SourceCache {
7    cache: HashMap<String, Option<(String, Vec<usize>)>>,
8    home_path: Option<Option<PathBuf>>,
9    rustc_sources: Option<Vec<(String, PathBuf)>>,
10}
11
12impl SourceCache {
13    #[allow(clippy::unused_self)]
14    fn home_path_uncached(&self) -> Option<PathBuf> {
15        let buffer = std::env::var("HOME").ok()?;
16        if buffer.is_empty() {
17            return None;
18        }
19
20        Some(buffer.into())
21    }
22
23    fn rustc_sources_uncached(&mut self) -> Option<Vec<(String, PathBuf)>> {
24        let mut list = Vec::new();
25        let toolchains_path = {
26            let mut buffer = self.home_path()?;
27            buffer.push(".rustup");
28            buffer.push("toolchains");
29            buffer
30        };
31
32        let iter = match std::fs::read_dir(&toolchains_path) {
33            Ok(iter) => iter,
34            Err(error) => {
35                log::warn!("Error reading {toolchains_path:?}: {error}");
36                return None;
37            }
38        };
39
40        for entry in iter {
41            let entry = match entry {
42                Ok(entry) => entry,
43                Err(error) => {
44                    log::warn!("Error reading entry in {toolchains_path:?}: {error}");
45                    continue;
46                }
47            };
48
49            let root_path = entry.path();
50            let rustc_path = {
51                let mut buffer = root_path.clone();
52                buffer.push("bin");
53                buffer.push("rustc");
54                buffer
55            };
56
57            if !rustc_path.exists() {
58                continue;
59            }
60
61            let sources_path = {
62                let mut buffer = root_path;
63                buffer.push("lib");
64                buffer.push("rustlib");
65                buffer.push("src");
66                buffer.push("rust");
67                buffer
68            };
69
70            if !sources_path.exists() {
71                continue;
72            }
73
74            let output = match std::process::Command::new(&rustc_path).args(["--version"]).output() {
75                Ok(output) => output,
76                Err(error) => {
77                    log::warn!("Error extracting version from {rustc_path:?}: {error}");
78                    continue;
79                }
80            };
81
82            if !output.status.success() {
83                log::warn!(
84                    "Error extracting version from {rustc_path:?}: non successful status: {}",
85                    output.status
86                );
87                continue;
88            }
89
90            let Ok(version_string) = String::from_utf8(output.stdout) else {
91                log::warn!("Error extracting version from {rustc_path:?}: returned version is not valid UTF-8");
92                continue;
93            };
94
95            // For example: "rustc 1.70.0-nightly (61863f31c 2023-08-15)"
96            let p = &version_string[version_string.find('(')? + 1..];
97            let p = &p[..p.find(' ')?];
98
99            log::debug!("Found Rust sources for hash '{p}' at: {sources_path:?}");
100            list.push((p.to_owned(), sources_path));
101        }
102
103        Some(list)
104    }
105
106    fn home_path(&mut self) -> Option<PathBuf> {
107        if let Some(ref path) = self.home_path {
108            return path.clone();
109        }
110
111        let result = self.home_path_uncached();
112        if let Some(ref path) = result {
113            log::debug!("Found HOME at: {path:?}");
114        } else {
115            log::debug!("HOME not found!");
116        }
117        self.home_path = Some(result.clone());
118        result
119    }
120
121    fn rustc_sources(&mut self) -> &[(String, PathBuf)] {
122        if self.rustc_sources.is_none() {
123            self.rustc_sources = Some(self.rustc_sources_uncached().unwrap_or_default());
124        }
125
126        self.rustc_sources.as_ref().unwrap()
127    }
128
129    fn read_source_file(&mut self, path: &str) -> Option<String> {
130        const HOME_PREFIX: &str = "~/";
131        const RUSTC_PREFIX: &str = "/rustc/";
132        let filesystem_path = if let Some(relative_path) = path.strip_prefix(HOME_PREFIX) {
133            // Example of a path like this:
134            //   "~/.cargo/registry/src/github.com-1ecc6299db9ec823/compiler_builtins-0.1.91/src/macros.rs"
135            let mut buffer = self.home_path()?;
136            buffer.push(relative_path);
137            Cow::Owned(buffer)
138        } else if let Some(relative_path) = path.strip_prefix(RUSTC_PREFIX) {
139            // Example of a path like this:
140            //   "/rustc/61863f31ccd4783186a5e839e6298d166d27368c/library/alloc/src/fmt.rs"
141            let p = relative_path;
142            let index = p.find('/')?;
143            let hash = &p[..index];
144            let relative_path = &p[index + 1..];
145            if relative_path.is_empty() || hash.is_empty() {
146                return None;
147            }
148
149            let sources = self.rustc_sources();
150            let (_, sources_root) = sources.iter().find(|(sources_hash, _)| hash.starts_with(sources_hash))?;
151            Cow::Owned(sources_root.join(relative_path))
152        } else {
153            Cow::Borrowed(Path::new(path))
154        };
155
156        match std::fs::read_to_string(&filesystem_path) {
157            Ok(contents) => {
158                log::debug!("Loaded source file: '{path}' (from {filesystem_path:?})");
159                Some(contents)
160            }
161            Err(error) => {
162                log::warn!("Failed to load source file '{path}' from {filesystem_path:?}: {error}");
163                None
164            }
165        }
166    }
167
168    pub fn lookup_source_line(&mut self, path: &str, line: u32) -> Option<&str> {
169        if !self.cache.contains_key(path) {
170            let Some(contents) = self.read_source_file(path) else {
171                self.cache.insert(path.to_owned(), None);
172                return None;
173            };
174
175            let mut line_to_offset = Vec::new();
176            line_to_offset.push(0);
177            for (offset, byte) in contents.bytes().enumerate() {
178                if byte == b'\n' {
179                    line_to_offset.push(offset + 1);
180                }
181            }
182
183            self.cache.insert(path.to_owned(), Some((contents, line_to_offset)));
184        }
185
186        let cached = self.cache.get(path)?;
187        let (contents, line_to_offset) = cached.as_ref()?;
188        let line = (line as usize).wrapping_sub(1);
189        let offset = *line_to_offset.get(line)?;
190        let next_offset = line_to_offset.get(line.wrapping_add(1)).copied().unwrap_or(contents.len());
191        contents.get(offset..next_offset).map(|s| s.trim_end())
192    }
193}