From d49757c7db27696152b63383194c647b8f13ec63 Mon Sep 17 00:00:00 2001
From: Alexander Polakov <a.polakov@zubr.io>
Date: Tue, 23 Nov 2021 18:20:54 +0300
Subject: [PATCH] Add unload() to remove module

---
 src/lua.rs     | 20 ++++++++++++++++++++
 tests/tests.rs | 31 +++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/src/lua.rs b/src/lua.rs
index f5d4350..544aebc 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -650,6 +650,26 @@ impl Lua {
         T::from_lua(value, self)
     }
 
+    /// Unloads module `modname`
+    /// Internally removes `modname` from package.loaded table.
+    pub fn unload<'lua, S>(&'lua self, modname: &S) -> Result<()>
+    where
+        S: AsRef<[u8]> + ?Sized,
+    {
+        let loaded = unsafe {
+            let _sg = StackGuard::new(self.state);
+            check_stack(self.state, 2)?;
+            protect_lua!(self.state, 0, 1, fn(state) {
+                ffi::luaL_getsubtable(state, ffi::LUA_REGISTRYINDEX, cstr!("_LOADED"));
+            })?;
+            Table(self.pop_ref())
+        };
+
+        let modname = self.create_string(modname)?;
+        loaded.raw_remove(modname)?;
+        Ok(())
+    }
+
     /// Consumes and leaks `Lua` object, returning a static reference `&'static Lua`.
     ///
     /// This function is useful when the `Lua` object is supposed to live for the remainder
diff --git a/tests/tests.rs b/tests/tests.rs
index 528f17d..ab2227b 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -1151,6 +1151,37 @@ fn test_load_from_function() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_unload() -> Result<()> {
+    let lua = Lua::new();
+
+    let i = Arc::new(AtomicU32::new(0));
+    let i2 = i.clone();
+    let func = lua.create_function(move |lua, modname: String| {
+        i2.fetch_add(1, Ordering::Relaxed);
+        let t = lua.create_table()?;
+        t.set("__name", modname)?;
+        Ok(t)
+    })?;
+
+    let t: Table = lua.load_from_function("my_module", func.clone())?;
+    assert_eq!(t.get::<_, String>("__name")?, "my_module");
+    assert_eq!(i.load(Ordering::Relaxed), 1);
+
+    let _: Value = lua.load_from_function("my_module", func.clone())?;
+    assert_eq!(i.load(Ordering::Relaxed), 1);
+
+    let _: () = lua.unload("my_module")?;
+    assert_eq!(i.load(Ordering::Relaxed), 1);
+
+    let _: Value = lua.load_from_function("my_module", func)?;
+    assert_eq!(i.load(Ordering::Relaxed), 2);
+
+    let _: () = lua.unload("my_module42")?;
+
+    Ok(())
+}
+
 #[test]
 fn test_inspect_stack() -> Result<()> {
     let lua = Lua::new();