Skip to main content

shape_ast/
module_utils.rs

1//! Shared module resolution utilities.
2//!
3//! Types and functions used by both `shape-runtime` (module loader) and
4//! `shape-vm` (import inlining) to inspect module exports and manipulate
5//! AST item lists during import resolution.
6
7use crate::ast::{ExportItem, Item, Program, Span};
8use crate::error::{Result, ShapeError};
9
10// ---------------------------------------------------------------------------
11// Types
12// ---------------------------------------------------------------------------
13
14/// High-level kind of an exported symbol.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum ModuleExportKind {
17    Function,
18    BuiltinFunction,
19    TypeAlias,
20    BuiltinType,
21    Interface,
22    Enum,
23    Annotation,
24    Value,
25}
26
27/// Exported symbol metadata discovered from a module's AST.
28#[derive(Debug, Clone)]
29pub struct ModuleExportSymbol {
30    /// Original symbol name in module scope.
31    pub name: String,
32    /// Alias if exported as `name as alias`.
33    pub alias: Option<String>,
34    /// High-level symbol kind.
35    pub kind: ModuleExportKind,
36    /// Source span for navigation/diagnostics.
37    pub span: Span,
38}
39
40// ---------------------------------------------------------------------------
41// direct_export_target
42// ---------------------------------------------------------------------------
43
44/// Map a direct (non-`Named`) export item to its name and kind.
45///
46/// Returns `None` for `ExportItem::Named`, which requires scope-level
47/// resolution handled by [`collect_exported_symbols`].
48pub fn direct_export_target(export_item: &ExportItem) -> Option<(String, ModuleExportKind)> {
49    match export_item {
50        ExportItem::Function(function) => {
51            Some((function.name.clone(), ModuleExportKind::Function))
52        }
53        ExportItem::BuiltinFunction(function) => {
54            Some((function.name.clone(), ModuleExportKind::BuiltinFunction))
55        }
56        ExportItem::BuiltinType(type_decl) => {
57            Some((type_decl.name.clone(), ModuleExportKind::BuiltinType))
58        }
59        ExportItem::TypeAlias(alias) => Some((alias.name.clone(), ModuleExportKind::TypeAlias)),
60        ExportItem::Enum(enum_def) => Some((enum_def.name.clone(), ModuleExportKind::Enum)),
61        ExportItem::Struct(struct_def) => {
62            Some((struct_def.name.clone(), ModuleExportKind::TypeAlias))
63        }
64        ExportItem::Interface(interface) => {
65            Some((interface.name.clone(), ModuleExportKind::Interface))
66        }
67        ExportItem::Trait(trait_def) => {
68            Some((trait_def.name.clone(), ModuleExportKind::Interface))
69        }
70        ExportItem::Annotation(annotation) => {
71            Some((annotation.name.clone(), ModuleExportKind::Annotation))
72        }
73        ExportItem::ForeignFunction(function) => {
74            Some((function.name.clone(), ModuleExportKind::Function))
75        }
76        ExportItem::Named(_) => None,
77    }
78}
79
80// ---------------------------------------------------------------------------
81// strip_import_items
82// ---------------------------------------------------------------------------
83
84/// Remove all `Item::Import` entries from a list of AST items.
85///
86/// Used when inlining module contents into a consumer program — the module's
87/// own imports have already been resolved and should not pollute the
88/// consumer's import set.
89pub fn strip_import_items(items: Vec<Item>) -> Vec<Item> {
90    items
91        .into_iter()
92        .filter(|item| !matches!(item, Item::Import(..)))
93        .collect()
94}
95
96// ---------------------------------------------------------------------------
97// collect_exported_symbols
98// ---------------------------------------------------------------------------
99
100/// Internal scope-symbol kind mirroring [`ModuleExportKind`] for scope
101/// resolution of `export { name }` statements.
102#[derive(Debug, Clone, Copy, PartialEq, Eq)]
103enum ScopeSymbolKind {
104    Function,
105    BuiltinFunction,
106    TypeAlias,
107    BuiltinType,
108    Interface,
109    Enum,
110    Annotation,
111    Value,
112}
113
114fn scope_symbol_kind_to_export(kind: ScopeSymbolKind) -> ModuleExportKind {
115    match kind {
116        ScopeSymbolKind::Function => ModuleExportKind::Function,
117        ScopeSymbolKind::BuiltinFunction => ModuleExportKind::BuiltinFunction,
118        ScopeSymbolKind::TypeAlias => ModuleExportKind::TypeAlias,
119        ScopeSymbolKind::BuiltinType => ModuleExportKind::BuiltinType,
120        ScopeSymbolKind::Interface => ModuleExportKind::Interface,
121        ScopeSymbolKind::Enum => ModuleExportKind::Enum,
122        ScopeSymbolKind::Annotation => ModuleExportKind::Annotation,
123        ScopeSymbolKind::Value => ModuleExportKind::Value,
124    }
125}
126
127/// Lightweight scope used to resolve `export { name }` to the right kind.
128struct ScopeTable {
129    symbols: std::collections::HashMap<String, (ScopeSymbolKind, Span)>,
130}
131
132impl ScopeTable {
133    fn from_program(program: &Program) -> Self {
134        let mut symbols = std::collections::HashMap::new();
135        for item in &program.items {
136            match item {
137                Item::Function(f, span) => {
138                    symbols.insert(f.name.clone(), (ScopeSymbolKind::Function, *span));
139                }
140                Item::BuiltinFunctionDecl(f, span) => {
141                    symbols.insert(f.name.clone(), (ScopeSymbolKind::BuiltinFunction, *span));
142                }
143                Item::BuiltinTypeDecl(t, span) => {
144                    symbols.insert(t.name.clone(), (ScopeSymbolKind::BuiltinType, *span));
145                }
146                Item::TypeAlias(a, span) => {
147                    symbols.insert(a.name.clone(), (ScopeSymbolKind::TypeAlias, *span));
148                }
149                Item::Enum(e, span) => {
150                    symbols.insert(e.name.clone(), (ScopeSymbolKind::Enum, *span));
151                }
152                Item::StructType(s, span) => {
153                    symbols.insert(s.name.clone(), (ScopeSymbolKind::TypeAlias, *span));
154                }
155                Item::Interface(i, span) => {
156                    symbols.insert(i.name.clone(), (ScopeSymbolKind::Interface, *span));
157                }
158                Item::Trait(t, span) => {
159                    symbols.insert(t.name.clone(), (ScopeSymbolKind::Interface, *span));
160                }
161                Item::VariableDecl(decl, span) => {
162                    if let Some(name) = decl.pattern.as_identifier() {
163                        symbols.insert(name.to_string(), (ScopeSymbolKind::Value, *span));
164                    }
165                }
166                Item::AnnotationDef(a, span) => {
167                    symbols.insert(a.name.clone(), (ScopeSymbolKind::Annotation, *span));
168                }
169                _ => {}
170            }
171        }
172        Self { symbols }
173    }
174
175    fn resolve(&self, name: &str) -> Option<(ScopeSymbolKind, Span)> {
176        self.symbols.get(name).copied()
177    }
178}
179
180/// Collect exported symbol metadata from a parsed module AST.
181///
182/// This is the canonical implementation shared by both the runtime module
183/// loader and the VM import inliner. It handles both direct exports
184/// (`pub fn`, `pub type`, etc.) and named re-exports (`export { a, b }`).
185pub fn collect_exported_symbols(program: &Program) -> Result<Vec<ModuleExportSymbol>> {
186    let scope = ScopeTable::from_program(program);
187    let mut symbols = Vec::new();
188
189    for item in &program.items {
190        let Item::Export(export, _) = item else {
191            continue;
192        };
193
194        // Direct exports: the ExportItem already carries name + kind.
195        if let Some((name, kind)) = direct_export_target(&export.item) {
196            let span = match &export.item {
197                ExportItem::Function(f) => f.name_span,
198                ExportItem::BuiltinFunction(f) => f.name_span,
199                ExportItem::Annotation(a) => a.name_span,
200                ExportItem::ForeignFunction(f) => f.name_span,
201                _ => scope
202                    .resolve(&name)
203                    .map(|(_, span)| span)
204                    .unwrap_or_default(),
205            };
206            symbols.push(ModuleExportSymbol {
207                name,
208                alias: None,
209                kind,
210                span,
211            });
212            continue;
213        }
214
215        // Named re-exports: resolve through scope table.
216        if let ExportItem::Named(specs) = &export.item {
217            for spec in specs {
218                match scope.resolve(&spec.name) {
219                    Some((kind, span)) => {
220                        if kind == ScopeSymbolKind::Value {
221                            return Err(ShapeError::ModuleError {
222                                message: format!(
223                                    "Cannot export variable '{}': variable exports are not yet supported. \
224                                     Only functions and types can be exported.",
225                                    spec.name
226                                ),
227                                module_path: None,
228                            });
229                        }
230                        symbols.push(ModuleExportSymbol {
231                            name: spec.name.clone(),
232                            alias: spec.alias.clone(),
233                            kind: scope_symbol_kind_to_export(kind),
234                            span,
235                        });
236                    }
237                    None => {
238                        return Err(ShapeError::ModuleError {
239                            message: format!(
240                                "Cannot export '{}': not found in module scope",
241                                spec.name
242                            ),
243                            module_path: None,
244                        });
245                    }
246                }
247            }
248        }
249    }
250
251    Ok(symbols)
252}
253
254// ---------------------------------------------------------------------------
255// export_kind_description
256// ---------------------------------------------------------------------------
257
258/// Human-readable description of an export kind for diagnostics.
259pub fn export_kind_description(kind: ModuleExportKind) -> &'static str {
260    match kind {
261        ModuleExportKind::Function => "a function",
262        ModuleExportKind::BuiltinFunction => "a builtin function",
263        ModuleExportKind::TypeAlias => "a type",
264        ModuleExportKind::BuiltinType => "a builtin type",
265        ModuleExportKind::Interface => "an interface",
266        ModuleExportKind::Enum => "an enum",
267        ModuleExportKind::Annotation => "an annotation",
268        ModuleExportKind::Value => "a value",
269    }
270}