1  use crate::analysis::{ModuleResolver, WorkspaceIndex};
  2  use crate::languages::LanguageRegistry;
  3  use crate::types::{ExportResolution, ModuleExport};
  4  use compact_str::CompactString;
  5  use rustc_hash::FxHashSet;
  6  use std::sync::Arc;
  7  use tower_lsp::lsp_types::{Range, Url};
  8  
  9  const MAX_RESOLUTION_DEPTH: usize = 10;
 10  
 11  #[derive(Debug, Clone)]
 12  pub enum CrossModuleResolution {
 13      EnvVar {
 14          name: CompactString,
 15  
 16          defining_file: Url,
 17  
 18          declaration_range: Range,
 19      },
 20  
 21      EnvObject {
 22          canonical_name: CompactString,
 23  
 24          defining_file: Url,
 25      },
 26  
 27      Unresolved,
 28  }
 29  
 30  pub struct CrossModuleResolver {
 31      workspace_index: Arc<WorkspaceIndex>,
 32  
 33      module_resolver: Arc<ModuleResolver>,
 34  
 35      languages: Arc<LanguageRegistry>,
 36  }
 37  
 38  impl CrossModuleResolver {
 39      pub fn new(
 40          workspace_index: Arc<WorkspaceIndex>,
 41          module_resolver: Arc<ModuleResolver>,
 42          languages: Arc<LanguageRegistry>,
 43      ) -> Self {
 44          Self {
 45              workspace_index,
 46              module_resolver,
 47              languages,
 48          }
 49      }
 50  
 51      pub fn resolve_import(
 52          &self,
 53          importer_uri: &Url,
 54          module_specifier: &str,
 55          imported_name: &str,
 56          is_default: bool,
 57      ) -> CrossModuleResolution {
 58          let source_uri = match self.resolve_module_specifier(importer_uri, module_specifier) {
 59              Some(uri) => uri,
 60              None => return CrossModuleResolution::Unresolved,
 61          };
 62  
 63          let mut visited = FxHashSet::default();
 64          self.resolve_recursive(&source_uri, imported_name, is_default, &mut visited, 0)
 65      }
 66  
 67      fn resolve_module_specifier(&self, from_uri: &Url, specifier: &str) -> Option<Url> {
 68          if let Some(cached) = self
 69              .workspace_index
 70              .cached_module_resolution(from_uri, specifier)
 71          {
 72              return cached;
 73          }
 74  
 75          let language = self.languages.get_for_uri(from_uri)?;
 76          let resolved = self
 77              .module_resolver
 78              .resolve_to_uri(specifier, from_uri, language.as_ref());
 79  
 80          self.workspace_index
 81              .cache_module_resolution(from_uri, specifier, resolved.clone());
 82  
 83          resolved
 84      }
 85  
 86      fn resolve_recursive(
 87          &self,
 88          source_uri: &Url,
 89          name: &str,
 90          is_default: bool,
 91          visited: &mut FxHashSet<(Url, String)>,
 92          depth: usize,
 93      ) -> CrossModuleResolution {
 94          if depth >= MAX_RESOLUTION_DEPTH {
 95              return CrossModuleResolution::Unresolved;
 96          }
 97  
 98          let key = (source_uri.clone(), name.to_string());
 99          if visited.contains(&key) {
100              return CrossModuleResolution::Unresolved;
101          }
102          visited.insert(key);
103  
104          let exports = match self.workspace_index.get_exports(source_uri) {
105              Some(e) => e,
106              None => return CrossModuleResolution::Unresolved,
107          };
108  
109          let export = if is_default {
110              exports.default_export.as_ref()
111          } else {
112              exports.get_export(name)
113          };
114  
115          if let Some(export) = export {
116              return self.resolve_export(export, source_uri, visited, depth);
117          }
118  
119          for wildcard_source in &exports.wildcard_reexports {
120              if let Some(wildcard_uri) = self.resolve_module_specifier(source_uri, wildcard_source) {
121                  let result = self.resolve_recursive(&wildcard_uri, name, false, visited, depth + 1);
122                  if !matches!(result, CrossModuleResolution::Unresolved) {
123                      return result;
124                  }
125              }
126          }
127  
128          CrossModuleResolution::Unresolved
129      }
130  
131      fn resolve_export(
132          &self,
133          export: &ModuleExport,
134          source_uri: &Url,
135          visited: &mut FxHashSet<(Url, String)>,
136          depth: usize,
137      ) -> CrossModuleResolution {
138          match &export.resolution {
139              ExportResolution::EnvVar { name } => CrossModuleResolution::EnvVar {
140                  name: name.clone(),
141                  defining_file: source_uri.clone(),
142                  declaration_range: export.declaration_range,
143              },
144  
145              ExportResolution::EnvObject { canonical_name } => CrossModuleResolution::EnvObject {
146                  canonical_name: canonical_name.clone(),
147                  defining_file: source_uri.clone(),
148              },
149  
150              ExportResolution::ReExport {
151                  source_module,
152                  original_name,
153              } => {
154                  if let Some(reexport_uri) = self.resolve_module_specifier(source_uri, source_module)
155                  {
156                      self.resolve_recursive(&reexport_uri, original_name, false, visited, depth + 1)
157                  } else {
158                      CrossModuleResolution::Unresolved
159                  }
160              }
161  
162              ExportResolution::LocalChain { symbol_id: _ } => CrossModuleResolution::Unresolved,
163  
164              ExportResolution::Unknown => CrossModuleResolution::Unresolved,
165          }
166      }
167  
168      pub fn files_exporting_env_var(&self, env_var_name: &str) -> Vec<Url> {
169          self.workspace_index.files_exporting_env_var(env_var_name)
170      }
171  
172      pub fn resolve_namespace_import(
173          &self,
174          importer_uri: &Url,
175          module_specifier: &str,
176      ) -> Vec<(CompactString, CompactString)> {
177          let source_uri = match self.resolve_module_specifier(importer_uri, module_specifier) {
178              Some(uri) => uri,
179              None => return Vec::new(),
180          };
181  
182          let exports = match self.workspace_index.get_exports(&source_uri) {
183              Some(e) => e,
184              None => return Vec::new(),
185          };
186  
187          let mut results = Vec::new();
188          let mut visited = FxHashSet::default();
189  
190          for (name, export) in &exports.named_exports {
191              if let CrossModuleResolution::EnvVar { name: env_name, .. } =
192                  self.resolve_export(export, &source_uri, &mut visited, 0)
193              {
194                  results.push((name.clone(), env_name));
195              }
196          }
197  
198          results
199      }
200  }
201  
202  impl CrossModuleResolver {
203      pub fn can_resolve(&self, from_uri: &Url, specifier: &str) -> bool {
204          self.resolve_module_specifier(from_uri, specifier).is_some()
205      }
206  
207      pub fn workspace_index(&self) -> &Arc<WorkspaceIndex> {
208          &self.workspace_index
209      }
210  
211      pub fn module_resolver(&self) -> &Arc<ModuleResolver> {
212          &self.module_resolver
213      }
214  }
215  
216  #[cfg(test)]
217  mod tests {
218      use super::*;
219      use crate::types::{FileExportEntry, ModuleExport};
220      use tempfile::TempDir;
221  
222      fn setup_test_environment() -> (TempDir, Arc<WorkspaceIndex>, Arc<ModuleResolver>) {
223          let temp_dir = TempDir::new().unwrap();
224          let workspace_root = temp_dir.path().to_path_buf();
225  
226          let workspace_index = Arc::new(WorkspaceIndex::new());
227          let module_resolver = Arc::new(ModuleResolver::new(workspace_root));
228  
229          (temp_dir, workspace_index, module_resolver)
230      }
231  
232      fn create_mock_registry() -> Arc<LanguageRegistry> {
233          use crate::languages::javascript::JavaScript;
234          let mut registry = LanguageRegistry::new();
235          registry.register(Arc::new(JavaScript));
236          Arc::new(registry)
237      }
238  
239      #[test]
240      fn test_resolve_direct_env_export() {
241          let (_temp, workspace_index, module_resolver) = setup_test_environment();
242          let languages = create_mock_registry();
243  
244          let config_uri = Url::parse("file:///workspace/src/config.ts").unwrap();
245          let mut exports = FileExportEntry::new();
246  
247          exports.named_exports.insert(
248              "dbUrl".into(),
249              ModuleExport {
250                  exported_name: "dbUrl".into(),
251                  local_name: None,
252                  resolution: ExportResolution::EnvVar {
253                      name: "DATABASE_URL".into(),
254                  },
255                  declaration_range: Range::default(),
256                  is_default: false,
257              },
258          );
259  
260          workspace_index.update_exports(&config_uri, exports);
261  
262          let resolver = CrossModuleResolver::new(workspace_index, module_resolver, languages);
263  
264          let result = resolver.resolve_import(
265              &Url::parse("file:///workspace/src/api.ts").unwrap(),
266              "./config",
267              "dbUrl",
268              false,
269          );
270  
271          assert!(matches!(result, CrossModuleResolution::Unresolved));
272      }
273  
274      #[test]
275      fn test_max_depth_prevents_infinite_loop() {
276          let (_temp, workspace_index, module_resolver) = setup_test_environment();
277          let languages = create_mock_registry();
278  
279          let uri_a = Url::parse("file:///workspace/a.ts").unwrap();
280          let uri_b = Url::parse("file:///workspace/b.ts").unwrap();
281  
282          let mut exports_a = FileExportEntry::new();
283          exports_a.named_exports.insert(
284              "foo".into(),
285              ModuleExport {
286                  exported_name: "foo".into(),
287                  local_name: None,
288                  resolution: ExportResolution::ReExport {
289                      source_module: "./b".into(),
290                      original_name: "foo".into(),
291                  },
292                  declaration_range: Range::default(),
293                  is_default: false,
294              },
295          );
296  
297          let mut exports_b = FileExportEntry::new();
298          exports_b.named_exports.insert(
299              "foo".into(),
300              ModuleExport {
301                  exported_name: "foo".into(),
302                  local_name: None,
303                  resolution: ExportResolution::ReExport {
304                      source_module: "./a".into(),
305                      original_name: "foo".into(),
306                  },
307                  declaration_range: Range::default(),
308                  is_default: false,
309              },
310          );
311  
312          workspace_index.update_exports(&uri_a, exports_a);
313          workspace_index.update_exports(&uri_b, exports_b);
314  
315          let resolver = CrossModuleResolver::new(workspace_index, module_resolver, languages);
316  
317          let result = resolver.resolve_import(&uri_a, "./b", "foo", false);
318          assert!(matches!(result, CrossModuleResolution::Unresolved));
319      }
320  }