mq_lang/
module.rs

1pub mod error;
2pub mod resolver;
3
4use crate::{
5    Arena, ArenaId, Program, Shared, TokenArena,
6    ast::{node as ast, parser::Parser},
7    lexer::{self, Lexer},
8    module::{
9        error::ModuleError,
10        resolver::{LocalFsModuleResolver, ModuleResolver},
11    },
12};
13use rustc_hash::FxHashMap;
14use smol_str::SmolStr;
15use std::{borrow::Cow, path::PathBuf, sync::LazyLock};
16
17pub type ModuleId = ArenaId<ModuleName>;
18
19type ModuleName = SmolStr;
20type StandardModules = FxHashMap<SmolStr, fn() -> &'static str>;
21
22impl<T: ModuleResolver> Default for ModuleLoader<T> {
23    fn default() -> Self {
24        Self::new(T::default())
25    }
26}
27
28#[derive(Debug, Clone)]
29pub struct ModuleLoader<T: ModuleResolver = LocalFsModuleResolver> {
30    pub(crate) loaded_modules: Arena<ModuleName>,
31    #[cfg(feature = "debugger")]
32    pub(crate) source_code: Option<String>,
33    resolver: T,
34}
35
36#[derive(Debug, Clone, PartialEq)]
37pub struct Module {
38    pub name: String,
39    pub functions: Program,
40    pub modules: Program,
41    pub vars: Program,
42    pub macros: Program,
43}
44
45impl Module {
46    pub const BUILTIN_MODULE: &str = "builtin";
47    pub const TOP_LEVEL_MODULE: &str = "top-level";
48    pub const TOP_LEVEL_MODULE_ID: ArenaId<ModuleName> = ArenaId::new(0);
49}
50
51pub static STANDARD_MODULES: LazyLock<StandardModules> = LazyLock::new(|| {
52    let mut map = FxHashMap::default();
53
54    macro_rules! std_module {
55        ($name:ident) => {
56            fn $name() -> &'static str {
57                include_str!(concat!("../modules/", stringify!($name), ".mq"))
58            }
59            map.insert(SmolStr::new(stringify!($name)), $name as fn() -> &'static str);
60        };
61    }
62
63    std_module!(ast);
64    std_module!(csv);
65    std_module!(fuzzy);
66    std_module!(json);
67    std_module!(section);
68    std_module!(test);
69    std_module!(toml);
70    std_module!(xml);
71    std_module!(yaml);
72
73    map
74});
75
76pub const BUILTIN_FILE: &str = include_str!("../builtin.mq");
77
78impl<T: ModuleResolver> ModuleLoader<T> {
79    pub fn new(resolver: T) -> Self {
80        let mut loaded_modules = Arena::new(10);
81        loaded_modules.alloc(Module::TOP_LEVEL_MODULE.into());
82
83        Self {
84            loaded_modules,
85            #[cfg(feature = "debugger")]
86            source_code: None,
87            resolver,
88        }
89    }
90
91    #[inline(always)]
92    pub fn module_name(&self, module_id: ModuleId) -> Cow<'static, str> {
93        match module_id {
94            Module::TOP_LEVEL_MODULE_ID => Cow::Borrowed(Module::TOP_LEVEL_MODULE),
95            _ => self
96                .loaded_modules
97                .get(module_id)
98                .map(|s| Cow::Owned(s.to_string()))
99                .unwrap_or_else(|| Cow::Borrowed("<unknown>")),
100        }
101    }
102
103    #[cfg(feature = "debugger")]
104    pub fn set_source_code(&mut self, source_code: String) {
105        self.source_code = Some(source_code);
106    }
107
108    pub fn search_paths(&self) -> Vec<PathBuf> {
109        self.resolver.search_paths()
110    }
111
112    pub fn set_search_paths(&mut self, paths: Vec<PathBuf>) {
113        self.resolver.set_search_paths(paths);
114    }
115
116    pub fn load(&mut self, module_name: &str, code: &str, token_arena: TokenArena) -> Result<Module, ModuleError> {
117        if self.loaded_modules.contains(module_name.into()) {
118            return Err(ModuleError::AlreadyLoaded(Cow::Owned(module_name.to_string())));
119        }
120
121        let module_id = self.loaded_modules.len().into();
122        let mut program = Self::parse_program(code, module_id, token_arena)?;
123
124        self.load_from_ast(module_name, &mut program)
125    }
126
127    pub fn load_from_ast(&mut self, module_name: &str, program: &mut Program) -> Result<Module, ModuleError> {
128        if self.loaded_modules.contains(module_name.into()) {
129            return Err(ModuleError::AlreadyLoaded(Cow::Owned(module_name.to_string())));
130        }
131
132        let modules = program
133            .iter()
134            .filter(|node| {
135                matches!(
136                    *node.expr,
137                    ast::Expr::Include(_) | ast::Expr::Module(_, _) | ast::Expr::Import(_)
138                )
139            })
140            .cloned()
141            .collect::<Vec<_>>();
142
143        let functions = program
144            .iter()
145            .filter(|node| matches!(*node.expr, ast::Expr::Def(..)))
146            .cloned()
147            .collect::<Vec<_>>();
148
149        let vars = program
150            .iter()
151            .filter(|node| matches!(*node.expr, ast::Expr::Let(..)))
152            .cloned()
153            .collect::<Vec<_>>();
154
155        let macros = program
156            .iter()
157            .filter(|node| matches!(*node.expr, ast::Expr::Macro(..)))
158            .cloned()
159            .collect::<Vec<_>>();
160
161        let expected_len = functions.len() + modules.len() + vars.len() + macros.len();
162
163        if program.len() != expected_len {
164            return Err(ModuleError::InvalidModule);
165        }
166
167        self.loaded_modules.alloc(module_name.into());
168
169        Ok(Module {
170            name: module_name.to_string(),
171            functions,
172            modules,
173            vars,
174            macros,
175        })
176    }
177
178    pub fn load_from_file(&mut self, module_path: &str, token_arena: TokenArena) -> Result<Module, ModuleError> {
179        let program = self.resolve(module_path)?;
180        self.load(module_path, &program, token_arena)
181    }
182
183    pub fn resolve(&self, module_name: &str) -> Result<String, ModuleError> {
184        if STANDARD_MODULES.contains_key(module_name) {
185            Ok(STANDARD_MODULES.get(module_name).map(|f| f()).unwrap().to_string())
186        } else {
187            self.resolver.resolve(module_name)
188        }
189    }
190
191    pub fn load_builtin(&mut self, token_arena: TokenArena) -> Result<Module, ModuleError> {
192        self.load(Module::BUILTIN_MODULE, BUILTIN_FILE, token_arena)
193    }
194
195    #[cfg(feature = "debugger")]
196    pub fn get_source_code_for_debug(&self, module_id: ModuleId) -> Result<String, ModuleError> {
197        match self.module_name(module_id) {
198            Cow::Borrowed(Module::TOP_LEVEL_MODULE) => Ok(self.source_code.clone().unwrap_or_default()),
199            Cow::Borrowed(Module::BUILTIN_MODULE) => Ok(BUILTIN_FILE.to_string()),
200            Cow::Borrowed(module_name) => self.resolve(module_name),
201            Cow::Owned(module_name) => self.resolve(&module_name),
202        }
203    }
204
205    pub fn get_source_code(&self, module_id: ModuleId, source_code: String) -> Result<String, ModuleError> {
206        match self.module_name(module_id) {
207            Cow::Borrowed(Module::TOP_LEVEL_MODULE) => Ok(source_code),
208            Cow::Borrowed(Module::BUILTIN_MODULE) => Ok(BUILTIN_FILE.to_string()),
209            Cow::Borrowed(module_name) => self.resolve(module_name),
210            Cow::Owned(module_name) => self.resolve(&module_name),
211        }
212    }
213
214    fn parse_program(code: &str, module_id: ModuleId, token_arena: TokenArena) -> Result<Program, ModuleError> {
215        let tokens = Lexer::new(lexer::Options::default()).tokenize(code, module_id)?;
216        let mut token_arena = {
217            #[cfg(not(feature = "sync"))]
218            {
219                token_arena.borrow_mut()
220            }
221
222            #[cfg(feature = "sync")]
223            {
224                token_arena.write().unwrap()
225            }
226        };
227
228        let program = Parser::new(
229            tokens.into_iter().map(Shared::new).collect::<Vec<_>>().iter(),
230            &mut token_arena,
231            module_id,
232        )
233        .parse()?;
234
235        Ok(program)
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use rstest::{fixture, rstest};
242    use smallvec::{SmallVec, smallvec};
243    use smol_str::SmolStr;
244
245    use crate::{
246        Shared, SharedCell, Token, TokenKind,
247        ast::node::{self as ast, IdentWithToken, Param},
248        module::LocalFsModuleResolver,
249        range::{Position, Range},
250    };
251
252    use super::{Module, ModuleError, ModuleLoader};
253
254    #[fixture]
255    fn token_arena() -> Shared<SharedCell<crate::arena::Arena<Shared<Token>>>> {
256        Shared::new(SharedCell::new(crate::arena::Arena::new(10)))
257    }
258
259    #[rstest]
260    #[case::load1("test".to_string(), Err(ModuleError::InvalidModule))]
261    #[case::load2("let test = \"value\"".to_string(), Ok(Module{
262        name: "test".to_string(),
263        functions: Vec::new(),
264        modules: Vec::new(),
265        vars: vec![
266            Shared::new(ast::Node{token_id: 0.into(), expr: Shared::new(ast::Expr::Let(
267                IdentWithToken::new_with_token("test", Some(Shared::new(Token{
268                    kind: TokenKind::Ident(SmolStr::new("test")),
269                    range: Range{start: Position{line: 1, column: 5}, end: Position{line: 1, column: 9}},
270                    module_id: 1.into()
271                }))),
272                Shared::new(ast::Node{token_id: 2.into(), expr: Shared::new(ast::Expr::Literal(ast::Literal::String("value".to_string())))})
273            ))})],
274        macros: Vec::new(),
275    }))]
276    #[case::load3("def test(): 1;".to_string(), Ok(Module{
277        name: "test".to_string(),
278        modules: Vec::new(),
279        functions: vec![
280            Shared::new(ast::Node{token_id: 0.into(), expr: Shared::new(ast::Expr::Def(
281            IdentWithToken::new_with_token("test", Some(Shared::new(Token{
282                kind: TokenKind::Ident(SmolStr::new("test")),
283                range: Range{start: Position{line: 1, column: 5}, end: Position{line: 1, column: 9}},
284                module_id: 1.into()
285            }))),
286            SmallVec::new(),
287            vec![
288                Shared::new(ast::Node{token_id: 2.into(), expr: Shared::new(ast::Expr::Literal(ast::Literal::Number(1.into())))})
289            ]
290            ))})],
291        vars: Vec::new(),
292        macros: Vec::new(),
293    }))]
294    #[case::load4("def test(a, b): add(a, b);".to_string(), Ok(Module{
295        name: "test".to_string(),
296        modules: Vec::new(),
297        functions: vec![
298            Shared::new(ast::Node{token_id: 0.into(), expr: Shared::new(ast::Expr::Def(
299                IdentWithToken::new_with_token("test", Some(Shared::new(Token{kind: TokenKind::Ident(SmolStr::new("test")), range: Range{start: Position{line: 1, column: 5}, end: Position{line: 1, column: 9}}, module_id: 1.into()}))),
300                smallvec![
301                    Param::new(IdentWithToken::new_with_token("a", Some(Shared::new(Token{kind: TokenKind::Ident(SmolStr::new("a")), range: Range{start: Position{line: 1, column: 10}, end: Position{line: 1, column: 11}}, module_id: 1.into()})))),
302                    Param::new(IdentWithToken::new_with_token("b", Some(Shared::new(Token{kind: TokenKind::Ident(SmolStr::new("b")), range: Range{start: Position{line: 1, column: 13}, end: Position{line: 1, column: 14}}, module_id: 1.into()})))),
303                ],
304                vec![
305                    Shared::new(ast::Node{token_id: 4.into(), expr: Shared::new(ast::Expr::Call(
306                    IdentWithToken::new_with_token("add", Some(Shared::new(Token{kind: TokenKind::Ident(SmolStr::new("add")), range: Range{start: Position{line: 1, column: 17}, end: Position{line: 1, column: 20}}, module_id: 1.into()}))),
307                    smallvec![
308                        Shared::new(ast::Node{token_id: 2.into(),
309                            expr: Shared::new(
310                                ast::Expr::Ident(IdentWithToken::new_with_token("a", Some(Shared::new(Token{kind: TokenKind::Ident(SmolStr::new("a")), range: Range{start: Position{line: 1, column: 21}, end: Position{line: 1, column: 22}}, module_id: 1.into()}))))
311                                )}),
312                        Shared::new(ast::Node{token_id: 3.into(),
313                            expr: Shared::new(
314                                ast::Expr::Ident(IdentWithToken::new_with_token("b", Some(Shared::new(Token{kind: TokenKind::Ident(SmolStr::new("b")), range: Range{start: Position{line: 1, column: 24}, end: Position{line: 1, column: 25}}, module_id: 1.into()}))))
315                            )})
316                    ],
317                ))})]
318            ))})],
319        vars: Vec::new(),
320        macros: Vec::new(),
321    }))]
322    fn test_load(
323        token_arena: Shared<SharedCell<crate::arena::Arena<Shared<Token>>>>,
324        #[case] program: String,
325        #[case] expected: Result<Module, ModuleError>,
326    ) {
327        assert_eq!(
328            ModuleLoader::new(LocalFsModuleResolver::default()).load("test", &program, token_arena),
329            expected
330        );
331    }
332
333    #[rstest]
334    #[case::load_standard_csv("csv", Ok(Module {
335        name: "csv".to_string(),
336        functions: Vec::new(),
337        modules: Vec::new(), // Assuming the csv.mq only contains definitions or is empty for this test
338        vars: Vec::new(),
339        macros: Vec::new(),
340    }))]
341    fn test_load_standard_module(
342        token_arena: Shared<SharedCell<crate::arena::Arena<Shared<Token>>>>,
343        #[case] module_name: &str,
344        #[case] expected: Result<Module, ModuleError>,
345    ) {
346        let mut loader = ModuleLoader::new(LocalFsModuleResolver::default());
347        let result = loader.load_from_file(module_name, token_arena.clone());
348        // Only check that loading does not return NotFound error and returns Some(Module)
349        match expected {
350            Ok(_) => {
351                assert!(result.is_ok(), "Expected Ok, got {:?}", result);
352                assert_eq!(result.unwrap().name, module_name);
353            }
354            Err(ref e) => {
355                assert_eq!(result.unwrap_err(), *e);
356            }
357        }
358    }
359
360    #[test]
361    fn test_standard_modules_contains_csv() {
362        assert!(super::STANDARD_MODULES.contains_key("csv"));
363        let csv_content = super::STANDARD_MODULES.get("csv").unwrap()();
364        assert!(csv_content.contains("")); // Just check it's a string, optionally check for expected content
365    }
366}
mq_lang/module.rs

mq_lang/
module.rs
