Start work on require overhaul

CHANGELOG.md

@@ -8,6 +8,16 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+
+### Added
+
+-   Builtin modules such as `fs`, `net` and others can now be imported using `require("@lune/fs")`, `require("@lune/net")` ..
+
+### Removed
+
+-   Removed option to preserve default Luau require behavior
+
 ## `0.5.6` - March 11th, 2023
 
 ### Added

packages/lib/src/globals/require.rs

@@ -1,137 +1,185 @@
 use std::{
-    env::{self, current_dir},
-    fs,
+    cell::RefCell,
+    collections::HashMap,
+    env::current_dir,
     path::{self, PathBuf},
 };
 
 use dunce::canonicalize;
 use mlua::prelude::*;
+use tokio::{fs, sync::oneshot};
 
-use crate::lua::table::TableBuilder;
+use crate::lua::{
+    table::TableBuilder,
+    task::{TaskScheduler, TaskSchedulerScheduleExt},
+};
 
 const REQUIRE_IMPL_LUA: &str = r#"
 local source = info(1, "s")
 if source == '[string "require"]' then
     source = info(2, "s")
 end
-local absolute, relative = paths(source, ...)
-if loaded[absolute] ~= true then
-    local first, second = load(absolute, relative)
-    if first == nil or second ~= nil then
-        error("Module did not return exactly one value")
-    end
-    loaded[absolute] = true
-    cache[absolute] = first
-    return first
-else
-    return cache[absolute]
-end
+local absolute, relative = importer:paths(source, ...)
+return importer:load(thread(), absolute, relative)
 "#;
 
-pub fn create(lua: &'static Lua) -> LuaResult<LuaTable> {
-    // Preserve original require behavior if we have a special env var set,
-    // returning an empty table since there are no globals to overwrite
-    if env::var_os("LUAU_PWD_REQUIRE").is_some() {
-        return TableBuilder::new(lua)?.build_readonly();
+#[derive(Debug, Clone, Default)]
+struct Importer<'lua> {
+    builtins: HashMap<String, LuaMultiValue<'lua>>,
+    cached: RefCell<HashMap<String, LuaResult<LuaMultiValue<'lua>>>>,
+    pwd: String,
+}
+
+impl<'lua> Importer<'lua> {
+    pub fn new() -> Self {
+        let mut pwd = current_dir()
+            .expect("Failed to access current working directory")
+            .to_string_lossy()
+            .to_string();
+        if !pwd.ends_with(path::MAIN_SEPARATOR) {
+            pwd = format!("{pwd}{}", path::MAIN_SEPARATOR)
+        }
+        Self {
+            pwd,
+            ..Default::default()
+        }
     }
-    // Store the current pwd, and make the functions for path conversions & loading a file
-    let mut require_pwd = current_dir()?.to_string_lossy().to_string();
-    if !require_pwd.ends_with(path::MAIN_SEPARATOR) {
-        require_pwd = format!("{require_pwd}{}", path::MAIN_SEPARATOR)
-    }
-    let require_info: LuaFunction = lua.named_registry_value("dbg.info")?;
-    let require_error: LuaFunction = lua.named_registry_value("error")?;
-    let require_get_abs_rel_paths = lua
-        .create_function(
-            |_, (require_pwd, require_source, require_path): (String, String, String)| {
-                if require_path.starts_with('@') {
-                    return Ok((require_path.clone(), require_path));
-                }
-                let path_relative_to_pwd = PathBuf::from(
-                    &require_source
-                        .trim_start_matches("[string \"")
-                        .trim_end_matches("\"]"),
-                )
-                .parent()
-                .unwrap()
-                .join(&require_path);
-                // Try to normalize and resolve relative path segments such as './' and '../'
-                let file_path = match (
-                    canonicalize(path_relative_to_pwd.with_extension("luau")),
-                    canonicalize(path_relative_to_pwd.with_extension("lua")),
-                ) {
-                    (Ok(luau), _) => luau,
-                    (_, Ok(lua)) => lua,
-                    _ => {
-                        return Err(LuaError::RuntimeError(format!(
-                            "File does not exist at path '{require_path}'"
-                        )))
-                    }
-                };
-                let absolute = file_path.to_string_lossy().to_string();
-                let relative = absolute.trim_start_matches(&require_pwd).to_string();
-                Ok((absolute, relative))
-            },
-        )?
-        .bind(require_pwd)?;
-    // Note that file loading must be blocking to guarantee the require cache works, if it
-    // were async then one lua script may require a module during the file reading process
-    let require_get_loaded_file = lua.create_function(
-        |lua: &Lua, (path_absolute, path_relative): (String, String)| {
-            // Check if we got a special require path starting with an @
-            if path_absolute == path_relative && path_absolute.starts_with('@') {
-                return match path_absolute {
-                    p if p.starts_with("@lune/") => {
-                        let _module_name = p.strip_prefix("@lune/").unwrap();
-                        // TODO: Return builtin module
-                        Err(LuaError::RuntimeError(
-                            "Builtin require paths prefixed by '@' are not yet supported"
-                                .to_string(),
-                        ))
-                    }
-                    _ => Err(LuaError::RuntimeError(
-                        "Custom require paths prefixed by '@' are not yet supported".to_string(),
-                    )),
-                };
+
+    fn paths(&self, require_source: String, require_path: String) -> LuaResult<(String, String)> {
+        if require_path.starts_with('@') {
+            return Ok((require_path.clone(), require_path));
+        }
+        let path_relative_to_pwd = PathBuf::from(
+            &require_source
+                .trim_start_matches("[string \"")
+                .trim_end_matches("\"]"),
+        )
+        .parent()
+        .unwrap()
+        .join(&require_path);
+        // Try to normalize and resolve relative path segments such as './' and '../'
+        let file_path = match (
+            canonicalize(path_relative_to_pwd.with_extension("luau")),
+            canonicalize(path_relative_to_pwd.with_extension("lua")),
+        ) {
+            (Ok(luau), _) => luau,
+            (_, Ok(lua)) => lua,
+            _ => {
+                return Err(LuaError::RuntimeError(format!(
+                    "File does not exist at path '{require_path}'"
+                )))
             }
-            // Use a name without extensions for loading the chunk, the
-            // above code assumes the require path is without extensions
-            let path_relative_no_extension = path_relative
-                .trim_end_matches(".lua")
-                .trim_end_matches(".luau");
-            // Try to read the wanted file, note that we use bytes instead of reading
-            // to a string since lua scripts are not necessarily valid utf-8 strings
-            match fs::read(path_absolute) {
-                Ok(contents) => lua
+        };
+        let absolute = file_path.to_string_lossy().to_string();
+        let relative = absolute.trim_start_matches(&self.pwd).to_string();
+        Ok((absolute, relative))
+    }
+
+    fn load_builtin(&self, module_name: &str) -> LuaResult<LuaMultiValue> {
+        match self.builtins.get(module_name) {
+            Some(module) => Ok(module.clone()),
+            None => Err(LuaError::RuntimeError(format!(
+                "No builtin module exists with the name '{}'",
+                module_name
+            ))),
+        }
+    }
+
+    async fn load_file(
+        &self,
+        lua: &'lua Lua,
+        absolute_path: String,
+        relative_path: String,
+    ) -> LuaResult<LuaMultiValue> {
+        let cached = { self.cached.borrow().get(&absolute_path).cloned() };
+        match cached {
+            Some(cached) => cached,
+            None => {
+                // Try to read the wanted file, note that we use bytes instead of reading
+                // to a string since lua scripts are not necessarily valid utf-8 strings
+                let contents = fs::read(&absolute_path).await.map_err(LuaError::external)?;
+                // Use a name without extensions for loading the chunk, some
+                // other code assumes the require path is without extensions
+                let path_relative_no_extension = relative_path
+                    .trim_end_matches(".lua")
+                    .trim_end_matches(".luau");
+                // Load the file into a thread
+                let loaded_func = lua
                     .load(&contents)
                     .set_name(path_relative_no_extension)?
-                    .eval::<LuaValue>(),
-                Err(e) => Err(LuaError::external(e)),
+                    .into_function()?;
+                let loaded_thread = lua.create_thread(loaded_func)?;
+                // Run the thread and provide a channel that will
+                // then get its result received when it finishes
+                let (tx, rx) = oneshot::channel();
+                {
+                    let sched = lua.app_data_ref::<&TaskScheduler>().unwrap();
+                    let task = sched.schedule_blocking(loaded_thread, LuaMultiValue::new())?;
+                    sched.set_task_result_sender(task, tx);
+                }
+                // Wait for the thread to finish running, cache + return our result
+                let rets = rx.await.expect("Sender was dropped during require");
+                self.cached.borrow_mut().insert(absolute_path, rets.clone());
+                rets
             }
-        },
-    )?;
-    /*
-        We need to get the source file where require was
-        called to be able to do path-relative requires,
-        so we make a small wrapper to do that here, this
-        will then call our actual async require function
+        }
+    }
 
-        This must be done in lua because due to how our
-        scheduler works mlua can not preserve debug info
-    */
+    async fn load(
+        &self,
+        lua: &'lua Lua,
+        absolute_path: String,
+        relative_path: String,
+    ) -> LuaResult<LuaMultiValue> {
+        if absolute_path == relative_path && absolute_path.starts_with('@') {
+            if let Some(module_name) = absolute_path.strip_prefix("@lune/") {
+                self.load_builtin(module_name)
+            } else {
+                Err(LuaError::RuntimeError(
+                    "Require paths prefixed by '@' are not yet supported".to_string(),
+                ))
+            }
+        } else {
+            self.load_file(lua, absolute_path, relative_path).await
+        }
+    }
+}
+
+impl<'i> LuaUserData for Importer<'i> {
+    fn add_methods<'lua, M: LuaUserDataMethods<'lua, Self>>(methods: &mut M) {
+        methods.add_method(
+            "paths",
+            |_, this, (require_source, require_path): (String, String)| {
+                this.paths(require_source, require_path)
+            },
+        );
+        methods.add_method(
+            "load",
+            |lua, this, (thread, absolute_path, relative_path): (LuaThread, String, String)| {
+                // TODO: Make this work
+                // this.load(lua, absolute_path, relative_path)
+                Ok(())
+            },
+        );
+    }
+}
+
+pub fn create(lua: &'static Lua) -> LuaResult<LuaTable> {
+    let require_importer = Importer::new();
+    let require_thread: LuaFunction = lua.named_registry_value("co.thread")?;
+    let require_info: LuaFunction = lua.named_registry_value("dbg.info")?;
     let require_env = TableBuilder::new(lua)?
-        .with_value("loaded", lua.create_table()?)?
-        .with_value("cache", lua.create_table()?)?
+        .with_value("importer", require_importer)?
+        .with_value("thread", require_thread)?
         .with_value("info", require_info)?
-        .with_value("error", require_error)?
-        .with_value("paths", require_get_abs_rel_paths)?
-        .with_value("load", require_get_loaded_file)?
         .build_readonly()?;
+
     let require_fn_lua = lua
         .load(REQUIRE_IMPL_LUA)
         .set_name("require")?
         .set_environment(require_env)?
         .into_function()?;
+
     TableBuilder::new(lua)?
         .with_value("require", require_fn_lua)?
         .build_readonly()

packages/lib/src/globals/task.rs

@@ -158,9 +158,9 @@ fn coroutine_resume<'lua>(
     let result = match value {
         LuaThreadOrTaskReference::Thread(t) => {
             let task = sched.create_task(TaskKind::Instant, t, None, true)?;
-            sched.resume_task(task, None)
+            sched.resume_task(task)
         }
-        LuaThreadOrTaskReference::TaskReference(t) => sched.resume_task(t, None),
+        LuaThreadOrTaskReference::TaskReference(t) => sched.resume_task(t),
     };
     sched.force_set_current_task(Some(current));
     match result {
@@ -187,7 +187,7 @@ fn coroutine_wrap<'lua>(lua: &'lua Lua, func: LuaFunction) -> LuaResult<LuaFunct
         let result = lua
             .app_data_ref::<&TaskScheduler>()
             .unwrap()
-            .resume_task(task, Some(Ok(args)));
+            .resume_task_override(task, Ok(args));
         sched.force_set_current_task(Some(current));
         result
     })

packages/lib/src/lua/task/ext/resume_ext.rs

@@ -87,9 +87,15 @@ fn resume_next_blocking_task(
         task
     } {
         None => TaskSchedulerState::new(scheduler),
-        Some(task) => match scheduler.resume_task(task, override_args) {
-            Ok(_) => TaskSchedulerState::new(scheduler),
-            Err(task_err) => TaskSchedulerState::err(scheduler, task_err),
+        Some(task) => match override_args {
+            Some(args) => match scheduler.resume_task_override(task, args) {
+                Ok(_) => TaskSchedulerState::new(scheduler),
+                Err(task_err) => TaskSchedulerState::err(scheduler, task_err),
+            },
+            None => match scheduler.resume_task(task) {
+                Ok(_) => TaskSchedulerState::new(scheduler),
+                Err(task_err) => TaskSchedulerState::err(scheduler, task_err),
+            },
         },
     }
 }

packages/lib/src/lua/task/scheduler.rs

@@ -9,11 +9,12 @@ use std::{
 use futures_util::{future::LocalBoxFuture, stream::FuturesUnordered, Future};
 use mlua::prelude::*;
 
-use tokio::sync::{mpsc, Mutex as AsyncMutex};
+use tokio::sync::{mpsc, oneshot, Mutex as AsyncMutex};
 
 use super::scheduler_message::TaskSchedulerMessage;
 pub use super::{task_kind::TaskKind, task_reference::TaskReference};
 
+type TaskResultSender = oneshot::Sender<LuaResult<LuaMultiValue<'static>>>;
 type TaskFutureRets<'fut> = LuaResult<Option<LuaMultiValue<'fut>>>;
 type TaskFuture<'fut> = LocalBoxFuture<'fut, (Option<TaskReference>, TaskFutureRets<'fut>)>;
 
@@ -48,6 +49,7 @@ pub struct TaskScheduler<'fut> {
     pub(super) tasks_count: Cell<usize>,
     pub(super) tasks_current: Cell<Option<TaskReference>>,
     pub(super) tasks_queue_blocking: RefCell<VecDeque<TaskReference>>,
+    pub(super) tasks_result_senders: RefCell<HashMap<TaskReference, TaskResultSender>>,
     pub(super) tasks_current_lua_error: Arc<RefCell<Option<LuaError>>>,
     // Future tasks & objects for waking
     pub(super) futures: AsyncMutex<FuturesUnordered<TaskFuture<'fut>>>,
@@ -77,6 +79,7 @@ impl<'fut> TaskScheduler<'fut> {
             tasks_count: Cell::new(0),
             tasks_current: Cell::new(None),
             tasks_queue_blocking: RefCell::new(VecDeque::new()),
+            tasks_result_senders: RefCell::new(HashMap::new()),
             tasks_current_lua_error,
             futures: AsyncMutex::new(FuturesUnordered::new()),
             futures_tx: tx,
@@ -270,6 +273,8 @@ impl<'fut> TaskScheduler<'fut> {
                     TaskKind::Future => self.futures_count.set(self.futures_count.get() - 1),
                     _ => self.tasks_count.set(self.tasks_count.get() - 1),
                 }
+                // Remove any sender
+                self.tasks_result_senders.borrow_mut().remove(task_ref);
                 // NOTE: We need to close the thread here to
                 // make 100% sure that nothing can resume it
                 let close: LuaFunction = self.lua.named_registry_value("co.close")?;
@@ -291,10 +296,60 @@ impl<'fut> TaskScheduler<'fut> {
 
         This will be a no-op if the task no longer exists.
     */
-    pub fn resume_task<'a>(
+    pub fn resume_task(&self, reference: TaskReference) -> LuaResult<LuaMultiValue> {
+        // Fetch and check if the task was removed, if it got
+        // removed it means it was intentionally cancelled
+        let task = {
+            let mut tasks = self.tasks.borrow_mut();
+            match tasks.remove(&reference) {
+                Some(task) => task,
+                None => return Ok(LuaMultiValue::new()),
+            }
+        };
+        // Decrement the corresponding task counter
+        match task.kind {
+            TaskKind::Future => self.futures_count.set(self.futures_count.get() - 1),
+            _ => self.tasks_count.set(self.tasks_count.get() - 1),
+        }
+        // Fetch and remove the thread to resume + its arguments
+        let thread: LuaThread = self.lua.registry_value(&task.thread)?;
+        let thread_args: Option<LuaMultiValue> = {
+            self.lua
+                .registry_value::<Option<Vec<LuaValue>>>(&task.args)
+                .expect("Failed to get stored args for task")
+                .map(LuaMultiValue::from_vec)
+        };
+        self.lua.remove_registry_value(task.thread)?;
+        self.lua.remove_registry_value(task.args)?;
+        // We got everything we need and our references
+        // were cleaned up properly, resume the thread
+        self.tasks_current.set(Some(reference));
+        let rets = match thread_args {
+            Some(args) => thread.resume(args),
+            None => thread.resume(()),
+        };
+        self.tasks_current.set(None);
+        // If we have a result sender for this task, we should run it if the thread finished
+        if thread.status() != LuaThreadStatus::Resumable {
+            if let Some(sender) = self.tasks_result_senders.borrow_mut().remove(&reference) {
+                let _ = sender.send(rets.clone());
+            }
+        }
+        rets
+    }
+
+    /**
+        Resumes a task, if the task still exists in the scheduler, using the given arguments.
+
+        A task may no longer exist in the scheduler if it has been manually
+        cancelled and removed by calling [`TaskScheduler::cancel_task()`].
+
+        This will be a no-op if the task no longer exists.
+    */
+    pub fn resume_task_override<'a>(
         &self,
         reference: TaskReference,
-        override_args: Option<LuaResult<LuaMultiValue<'a>>>,
+        override_args: LuaResult<LuaMultiValue<'a>>,
     ) -> LuaResult<LuaMultiValue<'a>> {
         // Fetch and check if the task was removed, if it got
         // removed it means it was intentionally cancelled
@@ -312,51 +367,27 @@ impl<'fut> TaskScheduler<'fut> {
         }
         // Fetch and remove the thread to resume + its arguments
         let thread: LuaThread = self.lua.registry_value(&task.thread)?;
-        let args_opt_res = override_args.or_else(|| {
-            Ok(self
-                .lua
-                .registry_value::<Option<Vec<LuaValue>>>(&task.args)
-                .expect("Failed to get stored args for task")
-                .map(LuaMultiValue::from_vec))
-            .transpose()
-        });
         self.lua.remove_registry_value(task.thread)?;
         self.lua.remove_registry_value(task.args)?;
         // We got everything we need and our references
         // were cleaned up properly, resume the thread
         self.tasks_current.set(Some(reference));
-        let rets = match args_opt_res {
-            Some(args_res) => match args_res {
-                Err(e) => {
-                    // NOTE: Setting this error here means that when the thread
-                    // is resumed it will error instantly, so we don't need
-                    // to call it with proper args, empty args is fine
-                    self.tasks_current_lua_error.replace(Some(e));
-                    thread.resume(())
-                }
-                Ok(args) => thread.resume(args),
-            },
-            None => thread.resume(()),
+        let rets = match override_args {
+            Err(e) => {
+                // NOTE: Setting this error here means that when the thread
+                // is resumed it will error instantly, so we don't need
+                // to call it with proper args, empty args is fine
+                self.tasks_current_lua_error.replace(Some(e));
+                thread.resume(())
+            }
+            Ok(args) => thread.resume(args),
         };
         self.tasks_current.set(None);
         rets
     }
 
     /**
-        Queues a new task to run on the task scheduler.
-
-        When we want to schedule a task to resume instantly after the
-        currently running task we should pass `after_current_resume = true`.
-
-        This is useful in cases such as our task.spawn implementation:
-
-        ```lua
-        task.spawn(function()
-            -- This will be a new task, but it should
-            -- also run right away, until the first yield
-        end)
-        -- Here we have either yielded or finished the above task
-        ```
+        Queues a new blocking task to run on the task scheduler.
     */
     pub(crate) fn queue_blocking_task(
         &self,
@@ -414,4 +445,10 @@ impl<'fut> TaskScheduler<'fut> {
         }));
         Ok(task_ref)
     }
+
+    pub(crate) fn set_task_result_sender(&self, task_ref: TaskReference, sender: TaskResultSender) {
+        self.tasks_result_senders
+            .borrow_mut()
+            .insert(task_ref, sender);
+    }
 }