1use std::collections::BTreeSet;
11
12use thiserror::Error;
13
14use crate::manifest::{levenshtein, Contract, ContractName, Manifest};
15
16#[derive(Debug, Error)]
17pub enum ReferenceError {
18 #[error("contract reference '{reference}' is ambiguous; use one of: {candidates}")]
19 Ambiguous {
20 reference: String,
21 candidates: String,
22 },
23 #[error("unknown contract '{reference}'; declared contracts: {candidates}")]
24 Unknown {
25 reference: String,
26 candidates: String,
27 },
28}
29
30#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct ContractSlice {
33 pub contract: ContractName,
34 pub symbols: Vec<String>,
37}
38
39pub fn resolve_reference<'m>(
47 manifest: &'m Manifest,
48 reference: &str,
49) -> Result<&'m ContractName, ReferenceError> {
50 if let Some(exact) = manifest
51 .contracts
52 .iter()
53 .find(|contract| contract.name.as_str() == reference)
54 {
55 return Ok(&exact.name);
56 }
57 let matches: Vec<&Contract> = manifest
58 .contracts
59 .iter()
60 .filter(|contract| shorthand_matches(contract, reference))
61 .collect();
62 match matches.as_slice() {
63 [single] => Ok(&single.name),
64 [] => Err(ReferenceError::Unknown {
65 reference: reference.to_owned(),
66 candidates: nearest_contracts(manifest, reference),
67 }),
68 several => Err(ReferenceError::Ambiguous {
69 reference: reference.to_owned(),
70 candidates: several
71 .iter()
72 .map(|contract| contract.name.as_str())
73 .collect::<Vec<_>>()
74 .join(", "),
75 }),
76 }
77}
78
79fn shorthand_matches(contract: &Contract, reference: &str) -> bool {
80 contract.name.as_str().starts_with(reference)
81 || contract.source == reference
82 || contract.source.rsplit('/').next() == Some(reference)
83}
84
85fn nearest_contracts(manifest: &Manifest, reference: &str) -> String {
86 let mut scored: Vec<(usize, &str)> = manifest
87 .contracts
88 .iter()
89 .map(|contract| {
90 (
91 levenshtein(reference, contract.name.as_str()),
92 contract.name.as_str(),
93 )
94 })
95 .collect();
96 scored.sort_unstable();
97 scored
98 .into_iter()
99 .map(|(_, name)| name)
100 .collect::<Vec<_>>()
101 .join(", ")
102}
103
104#[must_use]
108pub fn slices_for_write(
109 manifest: &Manifest,
110 path: &str,
111 content: &str,
112 contract_sources: &[(&str, &str)],
113) -> Vec<ContractSlice> {
114 let Some(mapping) = manifest.mapping_for(path) else {
115 return Vec::new();
116 };
117 let referenced = identifiers(language_for_path(path), content);
118 mapping
119 .contracts
120 .iter()
121 .map(|name| ContractSlice {
122 contract: name.clone(),
123 symbols: touched_symbols(manifest, name, contract_sources, &referenced),
124 })
125 .collect()
126}
127
128fn touched_symbols(
131 manifest: &Manifest,
132 name: &ContractName,
133 contract_sources: &[(&str, &str)],
134 referenced: &BTreeSet<String>,
135) -> Vec<String> {
136 let source_language = manifest
137 .contracts
138 .iter()
139 .find(|contract| contract.name == *name)
140 .and_then(|contract| language_for_path(&contract.source));
141 let Some(source_text) = contract_sources
142 .iter()
143 .find(|(source_name, _)| *source_name == name.as_str())
144 .map(|(_, text)| *text)
145 else {
146 return Vec::new();
147 };
148 exported_symbols(source_language, source_text)
149 .into_iter()
150 .filter(|symbol| referenced.contains(symbol))
151 .collect()
152}
153
154#[derive(Debug, Clone, Copy, PartialEq, Eq)]
157enum SourceLanguage {
158 TypeScript,
159 Python,
160 Rust,
161 Sql,
162}
163
164fn language_for_path(path: &str) -> Option<SourceLanguage> {
165 let extension = path.rsplit('.').next()?;
166 match extension {
167 "ts" | "tsx" | "mts" | "cts" => Some(SourceLanguage::TypeScript),
168 "py" => Some(SourceLanguage::Python),
169 "rs" => Some(SourceLanguage::Rust),
170 "sql" => Some(SourceLanguage::Sql),
171 _ => None,
172 }
173}
174
175fn grammar(language: SourceLanguage) -> tree_sitter::Language {
176 match language {
177 SourceLanguage::TypeScript => tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
178 SourceLanguage::Python => tree_sitter_python::LANGUAGE.into(),
179 SourceLanguage::Rust => tree_sitter_rust::LANGUAGE.into(),
180 SourceLanguage::Sql => tree_sitter_sequel::LANGUAGE.into(),
181 }
182}
183
184fn parse(language: Option<SourceLanguage>, text: &str) -> Option<tree_sitter::Tree> {
187 let mut parser = tree_sitter::Parser::new();
188 parser.set_language(&grammar(language?)).ok()?;
189 parser.parse(text, None)
190}
191
192fn descendants(root: tree_sitter::Node<'_>) -> Vec<tree_sitter::Node<'_>> {
194 let mut nodes = vec![root];
195 let mut index = 0;
196 while index < nodes.len() {
197 let node = nodes[index];
198 let mut cursor = node.walk();
199 nodes.extend(node.children(&mut cursor));
200 index += 1;
201 }
202 nodes
203}
204
205const IDENTIFIER_KINDS: &[&str] = &[
206 "identifier",
207 "type_identifier",
208 "property_identifier",
209 "shorthand_property_identifier",
210];
211
212fn identifiers(language: Option<SourceLanguage>, text: &str) -> BTreeSet<String> {
214 let Some(tree) = parse(language, text) else {
215 return BTreeSet::new();
216 };
217 descendants(tree.root_node())
218 .into_iter()
219 .filter(|node| IDENTIFIER_KINDS.contains(&node.kind()))
220 .filter_map(|node| node.utf8_text(text.as_bytes()).ok())
221 .map(str::to_owned)
222 .collect()
223}
224
225const DECLARATION_KINDS: &[&str] = &[
227 "variable_declarator",
229 "function_declaration",
230 "class_declaration",
231 "interface_declaration",
232 "type_alias_declaration",
233 "enum_declaration",
234 "export_specifier",
235 "function_definition",
237 "class_definition",
238 "function_item",
240 "struct_item",
241 "enum_item",
242 "const_item",
243 "static_item",
244 "type_item",
245];
246
247fn exported_symbols(language: Option<SourceLanguage>, text: &str) -> Vec<String> {
251 let Some(tree) = parse(language, text) else {
252 return Vec::new();
253 };
254 let mut symbols = Vec::new();
255 for node in descendants(tree.root_node()) {
256 if !DECLARATION_KINDS.contains(&node.kind()) || !is_exported(language, node) {
257 continue;
258 }
259 let named = node
260 .child_by_field_name("name")
261 .and_then(|name| name.utf8_text(text.as_bytes()).ok());
262 if let Some(symbol) = named {
263 if !symbols.iter().any(|existing| existing == symbol) {
264 symbols.push(symbol.to_owned());
265 }
266 }
267 }
268 symbols
269}
270
271fn is_exported(language: Option<SourceLanguage>, node: tree_sitter::Node<'_>) -> bool {
273 match language {
274 Some(SourceLanguage::TypeScript) => {
275 let mut ancestor = node.parent();
276 while let Some(current) = ancestor {
277 if current.kind() == "export_statement" {
278 return true;
279 }
280 ancestor = current.parent();
281 }
282 false
283 }
284 Some(SourceLanguage::Python) => true,
285 Some(SourceLanguage::Rust) => node
286 .child(0)
287 .is_some_and(|first| first.kind() == "visibility_modifier"),
288 Some(SourceLanguage::Sql) | None => false,
289 }
290}