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