1use serde_json::Value;
2use std::collections::HashMap;
3use tower_lsp::lsp_types::{Location,Position,Range, Url};
4
5
6#[derive(Debug, Clone)]
7pub struct NodeInfo {
8 pub src: String,
9 pub name_location: Option<String>,
10 pub name_locations: Vec<String>,
11 pub referenced_declaration: Option<u64>,
12 pub node_type: Option<String>,
13 pub member_location: Option<String>,
14 pub absolute_path: Option<String>,
15}
16
17fn push_if_node_or_array<'a>(tree: &'a Value, key: &str, stack: &mut Vec<&'a Value>) {
18 if let Some(value) = tree.get(key) {
19 match value {
20 Value::Array(arr) => {
21 stack.extend(arr);
22 }
23 Value::Object(_) => {
24 stack.push(value);
25 }
26 _ => {}
27 }
28 }
29}
30
31pub fn cache_ids(
32 sources: &Value,
33) -> (
34 HashMap<String, HashMap<u64, NodeInfo>>,
35 HashMap<String, String>,
36) {
37 let mut nodes: HashMap<String, HashMap<u64, NodeInfo>> = HashMap::new();
38 let mut path_to_abs: HashMap<String, String> = HashMap::new();
39
40 if let Some(sources_obj) = sources.as_object() {
41 for (path, contents) in sources_obj {
42 if let Some(contents_array) = contents.as_array()
43 && let Some(first_content) = contents_array.first()
44 && let Some(source_file) = first_content.get("source_file")
45 && let Some(ast) = source_file.get("ast")
46 {
47 let abs_path = ast
49 .get("absolutePath")
50 .and_then(|v| v.as_str())
51 .unwrap_or(path)
52 .to_string();
53
54 path_to_abs.insert(path.clone(), abs_path.clone());
55
56 if !nodes.contains_key(&abs_path) {
58 nodes.insert(abs_path.clone(), HashMap::new());
59 }
60
61 if let Some(id) = ast.get("id").and_then(|v| v.as_u64())
62 && let Some(src) = ast.get("src").and_then(|v| v.as_str())
63 {
64 nodes.get_mut(&abs_path).unwrap().insert(
65 id,
66 NodeInfo {
67 src: src.to_string(),
68 name_location: None,
69 name_locations: vec![],
70 referenced_declaration: None,
71 node_type: ast
72 .get("nodeType")
73 .and_then(|v| v.as_str())
74 .map(|s| s.to_string()),
75 member_location: None,
76 absolute_path: ast
77 .get("absolutePath")
78 .and_then(|v| v.as_str())
79 .map(|s| s.to_string()),
80 },
81 );
82 }
83
84 let mut stack = vec![ast];
85
86 while let Some(tree) = stack.pop() {
87 if let Some(id) = tree.get("id").and_then(|v| v.as_u64())
88 && let Some(src) = tree.get("src").and_then(|v| v.as_str())
89 {
90 let mut name_location = tree
92 .get("nameLocation")
93 .and_then(|v| v.as_str())
94 .map(|s| s.to_string());
95
96 if name_location.is_none()
100 && let Some(name_locations) = tree.get("nameLocations")
101 && let Some(locations_array) = name_locations.as_array()
102 && !locations_array.is_empty()
103 {
104 let node_type = tree.get("nodeType").and_then(|v| v.as_str());
105 if node_type == Some("IdentifierPath") {
106 name_location = locations_array
107 .last()
108 .and_then(|v| v.as_str())
109 .map(|s| s.to_string());
110 } else {
111 name_location = locations_array[0].as_str().map(|s| s.to_string());
112 }
113 }
114
115 let name_locations = if let Some(name_locations) = tree.get("nameLocations")
116 && let Some(locations_array) = name_locations.as_array()
117 {
118 locations_array
119 .iter()
120 .filter_map(|v| v.as_str().map(|s| s.to_string()))
121 .collect()
122 } else {
123 vec![]
124 };
125
126 let mut final_name_location = name_location;
127 if final_name_location.is_none()
128 && let Some(member_loc) = tree.get("memberLocation").and_then(|v| v.as_str()) {
129 final_name_location = Some(member_loc.to_string());
130 }
131
132 let node_info = NodeInfo {
133 src: src.to_string(),
134 name_location: final_name_location,
135 name_locations,
136 referenced_declaration: tree
137 .get("referencedDeclaration")
138 .and_then(|v| v.as_u64()),
139 node_type: tree
140 .get("nodeType")
141 .and_then(|v| v.as_str())
142 .map(|s| s.to_string()),
143 member_location: tree
144 .get("memberLocation")
145 .and_then(|v| v.as_str())
146 .map(|s| s.to_string()),
147 absolute_path: tree
148 .get("absolutePath")
149 .and_then(|v| v.as_str())
150 .map(|s| s.to_string()),
151 };
152
153 nodes.get_mut(&abs_path).unwrap().insert(id, node_info);
154 }
155
156 push_if_node_or_array(tree, "arguments", &mut stack);
157 push_if_node_or_array(tree, "arguments", &mut stack);
158 push_if_node_or_array(tree, "baseContracts", &mut stack);
159 push_if_node_or_array(tree, "baseContracts", &mut stack);
160 push_if_node_or_array(tree, "baseExpression", &mut stack);
161 push_if_node_or_array(tree, "baseName", &mut stack);
162 push_if_node_or_array(tree, "baseType", &mut stack);
163 push_if_node_or_array(tree, "block", &mut stack);
164 push_if_node_or_array(tree, "body", &mut stack);
165 push_if_node_or_array(tree, "components", &mut stack);
166 push_if_node_or_array(tree, "components", &mut stack);
167 push_if_node_or_array(tree, "condition", &mut stack);
168 push_if_node_or_array(tree, "declarations", &mut stack);
169 push_if_node_or_array(tree, "endExpression", &mut stack);
170 push_if_node_or_array(tree, "errorCall", &mut stack);
171 push_if_node_or_array(tree, "eventCall", &mut stack);
172 push_if_node_or_array(tree, "expression", &mut stack);
173 push_if_node_or_array(tree, "externalCall", &mut stack);
174 push_if_node_or_array(tree, "falseBody", &mut stack);
175 push_if_node_or_array(tree, "falseExpression", &mut stack);
176 push_if_node_or_array(tree, "file", &mut stack);
177 push_if_node_or_array(tree, "foreign", &mut stack);
178 push_if_node_or_array(tree, "indexExpression", &mut stack);
179 push_if_node_or_array(tree, "initialValue", &mut stack);
180 push_if_node_or_array(tree, "initialValue", &mut stack);
181 push_if_node_or_array(tree, "initializationExpression", &mut stack);
182 push_if_node_or_array(tree, "keyType", &mut stack);
183 push_if_node_or_array(tree, "leftExpression", &mut stack);
184 push_if_node_or_array(tree, "leftHandSide", &mut stack);
185 push_if_node_or_array(tree, "libraryName", &mut stack);
186 push_if_node_or_array(tree, "literals", &mut stack);
187 push_if_node_or_array(tree, "loopExpression", &mut stack);
188 push_if_node_or_array(tree, "members", &mut stack);
189 push_if_node_or_array(tree, "modifierName", &mut stack);
190 push_if_node_or_array(tree, "modifiers", &mut stack);
191 push_if_node_or_array(tree, "name", &mut stack);
192 push_if_node_or_array(tree, "names", &mut stack);
193 push_if_node_or_array(tree, "nodes", &mut stack);
194 push_if_node_or_array(tree, "options", &mut stack);
195 push_if_node_or_array(tree, "options", &mut stack);
196 push_if_node_or_array(tree, "options", &mut stack);
197 push_if_node_or_array(tree, "overrides", &mut stack);
198 push_if_node_or_array(tree, "overrides", &mut stack);
199 push_if_node_or_array(tree, "parameters", &mut stack);
200 push_if_node_or_array(tree, "parameters", &mut stack);
201 push_if_node_or_array(tree, "pathNode", &mut stack);
202 push_if_node_or_array(tree, "returnParameters", &mut stack);
203 push_if_node_or_array(tree, "returnParameters", &mut stack);
204 push_if_node_or_array(tree, "rightExpression", &mut stack);
205 push_if_node_or_array(tree, "rightHandSide", &mut stack);
206 push_if_node_or_array(tree, "startExpression", &mut stack);
207 push_if_node_or_array(tree, "statements", &mut stack);
208 push_if_node_or_array(tree, "statements", &mut stack);
209 push_if_node_or_array(tree, "storageLayout", &mut stack);
210 push_if_node_or_array(tree, "subExpression", &mut stack);
211 push_if_node_or_array(tree, "subdenomination", &mut stack);
212 push_if_node_or_array(tree, "symbolAliases", &mut stack);
213 push_if_node_or_array(tree, "trueBody", &mut stack);
214 push_if_node_or_array(tree, "trueExpression", &mut stack);
215 push_if_node_or_array(tree, "typeName", &mut stack);
216 push_if_node_or_array(tree, "unitAlias", &mut stack);
217 push_if_node_or_array(tree, "value", &mut stack);
218 push_if_node_or_array(tree, "valueType", &mut stack);
219 }
220 }
221 }
222 }
223
224 (nodes, path_to_abs)
225}
226
227pub fn pos_to_bytes(source_bytes: &[u8], position: Position) -> usize {
228 let text = String::from_utf8_lossy(source_bytes);
229 let lines: Vec<&str> = text.lines().collect();
230
231 let mut byte_offset = 0;
232
233 for (line_num, line_text) in lines.iter().enumerate() {
234 if line_num < position.line as usize {
235 byte_offset += line_text.len() + 1; } else if line_num == position.line as usize {
237 let char_offset = std::cmp::min(position.character as usize, line_text.len());
238 byte_offset += char_offset;
239 break;
240 }
241 }
242
243 byte_offset
244}
245
246pub fn bytes_to_pos(source_bytes: &[u8], byte_offset: usize) -> Option<Position> {
247 let text = String::from_utf8_lossy(source_bytes);
248 let mut curr_offset = 0;
249
250 for (line_num, line_text) in text.lines().enumerate() {
251 let line_bytes = line_text.len() + 1; if curr_offset + line_bytes > byte_offset {
253 let col = byte_offset - curr_offset;
254 return Some(Position::new(line_num as u32, col as u32));
255 }
256 curr_offset += line_bytes;
257 }
258
259 None
260}
261
262pub fn goto_bytes(
263 nodes: &HashMap<String, HashMap<u64, NodeInfo>>,
264 path_to_abs: &HashMap<String, String>,
265 id_to_path: &HashMap<String, String>,
266 uri: &str,
267 position: usize,
268) -> Option<(String, usize)> {
269 let path = match uri.starts_with("file://") {
270 true => &uri[7..],
271 false => uri,
272 };
273
274 let abs_path = path_to_abs.get(path)?;
276
277 let current_file_nodes = nodes.get(abs_path)?;
279
280 let mut refs = HashMap::new();
281
282 for (id, content) in current_file_nodes {
284 if content.referenced_declaration.is_none() {
285 continue;
286 }
287
288 let src_parts: Vec<&str> = content.src.split(':').collect();
289 if src_parts.len() != 3 {
290 continue;
291 }
292
293 let start_b: usize = src_parts[0].parse().ok()?;
294 let length: usize = src_parts[1].parse().ok()?;
295 let end_b = start_b + length;
296
297 if start_b <= position && position < end_b {
298 let diff = end_b - start_b;
299 if !refs.contains_key(&diff) || refs[&diff] <= *id {
300 refs.insert(diff, *id);
301 }
302 }
303 }
304
305 if refs.is_empty() {
306 let tmp = current_file_nodes.iter();
309 for (_id, content) in tmp {
310 if content.node_type == Some("ImportDirective".to_string()) {
311 let src_parts: Vec<&str> = content.src.split(':').collect();
312 if src_parts.len() != 3 {
313 continue;
314 }
315
316 let start_b: usize = src_parts[0].parse().ok()?;
317 let length: usize = src_parts[1].parse().ok()?;
318 let end_b = start_b + length;
319
320 if start_b <= position && position < end_b && let Some(import_path) = &content.absolute_path {
321 return Some((import_path.clone(), 0));
322 }
323 }
324 }
325 return None;
326 }
327
328 let min_diff = *refs.keys().min()?;
330 let chosen_id = refs[&min_diff];
331 let ref_id = current_file_nodes[&chosen_id].referenced_declaration?;
332
333 let mut target_node: Option<&NodeInfo> = None;
335 for file_nodes in nodes.values() {
336 if let Some(node) = file_nodes.get(&ref_id) {
337 target_node = Some(node);
338 break;
339 }
340 }
341
342 let node = target_node?;
343
344 let (location_str, file_id) = if let Some(name_location) = &node.name_location {
346 let parts: Vec<&str> = name_location.split(':').collect();
347 if parts.len() == 3 {
348 (parts[0], parts[2])
349 } else {
350 return None;
351 }
352 } else {
353 let parts: Vec<&str> = node.src.split(':').collect();
354 if parts.len() == 3 {
355 (parts[0], parts[2])
356 } else {
357 return None;
358 }
359 };
360
361 let location: usize = location_str.parse().ok()?;
362 let file_path = id_to_path.get(file_id)?.clone();
363
364 Some((file_path, location))
365}
366
367pub fn goto_declaration(
368 ast_data: &Value,
369 file_uri: &Url,
370 position: Position,
371 source_bytes: &[u8]
372) -> Option<Location> {
373 let sources = ast_data.get("sources")?;
374 let build_infos = ast_data.get("build_infos")?.as_array()?;
375 let first_build_info = build_infos.first()?;
376 let id_to_path = first_build_info.get("source_id_to_path")?.as_object()?;
377
378 let id_to_path_map: HashMap<String, String> = id_to_path
379 .iter()
380 .map(|(k,v)| (k.clone(), v.as_str().unwrap_or("").to_string()))
381 .collect();
382
383 let (nodes, path_to_abs) = cache_ids(sources);
384 let byte_position = pos_to_bytes(source_bytes, position);
385
386 if let Some((file_path, location_bytes)) = goto_bytes(
387 &nodes,
388 &path_to_abs,
389 &id_to_path_map,
390 file_uri.as_ref(),
391 byte_position,
392 ) {
393 let target_file_path = std::path::Path::new(&file_path);
394 let absolute_path = if target_file_path.is_absolute() {
395 target_file_path.to_path_buf()
396 } else {
397 std::env::current_dir().ok()?.join(target_file_path)
398 };
399
400 if let Ok(target_source_bytes) = std::fs::read(&absolute_path)
401 && let Some(target_position) = bytes_to_pos(&target_source_bytes, location_bytes)
402 && let Ok(target_uri) = Url::from_file_path(&absolute_path)
403 {
404 return Some(Location {
405 uri: target_uri,
406 range: Range {
407 start: target_position,
408 end: target_position,
409 }
410 });
411 }
412
413 };
414
415 None
416
417
418}