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

mq_lang/
module.rs
