From fc90630781323024263c1d832c63a566986333e3 Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Tue, 22 Aug 2023 15:51:28 -0500
Subject: [PATCH] Implement context preservation for Lune struct

---
 src/lune/mod.rs | 62 +++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 58 insertions(+), 4 deletions(-)

diff --git a/src/lune/mod.rs b/src/lune/mod.rs
index 2eb9270..9af012f 100644
--- a/src/lune/mod.rs
+++ b/src/lune/mod.rs
@@ -1,5 +1,7 @@
 use std::process::ExitCode;
 
+use mlua::{Lua, Table as LuaTable, Value as LuaValue};
+
 mod builtins;
 mod error;
 mod globals;
@@ -9,7 +11,6 @@ mod util;
 use self::scheduler::{LuaSchedulerExt, Scheduler};
 
 pub use error::LuneError;
-use mlua::Lua;
 
 #[derive(Debug, Clone)]
 pub struct Lune {
@@ -20,7 +21,7 @@ pub struct Lune {
 
 impl Lune {
     /**
-        Creates a new Lune script runner.
+        Creates a new Lune runtime, with a new Luau VM and task scheduler.
     */
     #[allow(clippy::new_without_default)]
     pub fn new() -> Self {
@@ -41,7 +42,7 @@ impl Lune {
     }
 
     /**
-        Arguments to give in `process.args` for a Lune script.
+        Sets arguments to give in `process.args` for Lune scripts.
     */
     pub fn with_args<V>(mut self, args: V) -> Self
     where
@@ -53,7 +54,7 @@ impl Lune {
     }
 
     /**
-        Runs a Lune script inside of a new Luau VM.
+        Runs a Lune script inside of the current runtime.
     */
     pub async fn run(
         &mut self,
@@ -66,6 +67,59 @@ impl Lune {
             .set_name(script_name.as_ref());
 
         self.scheduler.push_back(self.lua, main, ())?;
+
         Ok(self.scheduler.run_to_completion(self.lua).await)
     }
+
+    /**
+        Creates a context struct that can be called / ran multiple times,
+        preserving the function environment / context between each run.
+
+        Note that this is slightly slower than using [`run`] directly.
+    */
+    pub fn context(&self, script_name: impl Into<String>) -> Result<LuneContext, LuneError> {
+        let script_name = script_name.into();
+
+        let environment = self.lua.create_table()?;
+        for pair in self.lua.globals().pairs::<LuaValue, LuaValue>() {
+            let (key, value) = pair?;
+            environment.set(key, value)?;
+        }
+
+        Ok(LuneContext {
+            parent: self,
+            script_name,
+            environment,
+        })
+    }
+}
+
+pub struct LuneContext<'a> {
+    parent: &'a Lune,
+    script_name: String,
+    environment: LuaTable<'a>,
+}
+
+impl<'a> LuneContext<'a> {
+    /**
+        Runs a Lune script inside of the current runtime.
+
+        The function environment / context will be preserved between each run.
+    */
+    pub async fn run(&mut self, script_contents: impl AsRef<[u8]>) -> Result<ExitCode, LuneError> {
+        let main = self
+            .parent
+            .lua
+            .load(script_contents.as_ref())
+            .set_name(&self.script_name)
+            .set_environment(self.environment.clone());
+
+        self.parent.scheduler.push_back(self.parent.lua, main, ())?;
+
+        Ok(self
+            .parent
+            .scheduler
+            .run_to_completion(self.parent.lua)
+            .await)
+    }
 }