Add set_warning_function/remove_warning_function/warning functions to Lua for 5.4

This utilizes Lua 5.4 warnings system (https://www.lua.org/manual/5.4/manual.html#pdf-warn)

4 files changed, 119 insertions, 2 deletions

diff --git a/src/ffi/lua.rs b/src/ffi/lua.rs
index 41ef9c6..12ec1cc 100644
--- a/src/ffi/lua.rs
+++ b/src/ffi/lua.rs
@@ -574,7 +574,7 @@ pub unsafe fn lua_resume(
 // warning-related functions
 #[cfg(feature = "lua54")]
 extern "C" {
-    pub fn lua_setwarnf(L: *mut lua_State, f: lua_WarnFunction, ud: *mut c_void);
+    pub fn lua_setwarnf(L: *mut lua_State, f: Option<lua_WarnFunction>, ud: *mut c_void);
     pub fn lua_warning(L: *mut lua_State, msg: *const c_char, tocont: c_int);
 }
 
diff --git a/src/lua.rs b/src/lua.rs
index ebb940c..7d8b0eb 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -38,7 +38,10 @@ use crate::value::{FromLua, FromLuaMulti, MultiValue, Nil, ToLua, ToLuaMulti, Va
 #[cfg(not(feature = "lua54"))]
 use crate::util::push_userdata;
 #[cfg(feature = "lua54")]
-use crate::{userdata::USER_VALUE_MAXSLOT, util::push_userdata_uv};
+use {
+    crate::{types::WarnCallback, userdata::USER_VALUE_MAXSLOT, util::push_userdata_uv},
+    std::ffi::CStr,
+};
 
 #[cfg(not(feature = "send"))]
 use std::rc::Rc;
@@ -101,6 +104,8 @@ struct ExtraData {
     ref_waker_idx: c_int,
 
     hook_callback: Option<HookCallback>,
+    #[cfg(feature = "lua54")]
+    warn_callback: Option<WarnCallback>,
 }
 
 #[cfg_attr(any(feature = "lua51", feature = "luajit"), allow(dead_code))]
@@ -523,6 +528,8 @@ impl Lua {
             #[cfg(feature = "async")]
             ref_waker_idx,
             hook_callback: None,
+            #[cfg(feature = "lua54")]
+            warn_callback: None,
         }));
 
         mlua_expect!(
@@ -780,6 +787,64 @@ impl Lua {
         }
     }
 
+    /// Sets the warning function to be used by Lua to emit warnings.
+    ///
+    /// Requires `feature = "lua54"`
+    #[cfg(feature = "lua54")]
+    pub fn set_warning_function<F>(&self, callback: F)
+    where
+        F: 'static + MaybeSend + Fn(&Lua, &CStr, bool) -> Result<()>,
+    {
+        unsafe extern "C" fn warn_proc(ud: *mut c_void, msg: *const c_char, tocont: c_int) {
+            let state = ud as *mut ffi::lua_State;
+            let lua = match Lua::make_from_ptr(state) {
+                Some(lua) => lua,
+                None => return,
+            };
+            let extra = lua.extra.get();
+            callback_error_ext(state, extra, move |_| {
+                let cb = mlua_expect!(
+                    (*lua.extra.get()).warn_callback.as_ref(),
+                    "no warning callback set"
+                );
+                let msg = CStr::from_ptr(msg);
+                cb(&lua, msg, tocont != 0)
+            });
+        }
+
+        let state = self.main_state.unwrap_or(self.state);
+        unsafe {
+            (*self.extra.get()).warn_callback = Some(Box::new(callback));
+            ffi::lua_setwarnf(state, Some(warn_proc), state as *mut c_void);
+        }
+    }
+
+    /// Removes warning function previously set by `set_warning_function`.
+    ///
+    /// This function has no effect if a warning function was not previously set.
+    ///
+    /// Requires `feature = "lua54"`
+    #[cfg(feature = "lua54")]
+    pub fn remove_warning_function(&self) {
+        let state = self.main_state.unwrap_or(self.state);
+        unsafe {
+            (*self.extra.get()).warn_callback = None;
+            ffi::lua_setwarnf(state, None, ptr::null_mut());
+        }
+    }
+
+    /// Emits a warning with the given message.
+    ///
+    /// A message in a call with `tocont` set to `true` should be continued in another call to this function.
+    ///
+    /// Requires `feature = "lua54"`
+    #[cfg(feature = "lua54")]
+    pub fn warning<S: Into<Vec<u8>>>(&self, msg: S, tocont: bool) -> Result<()> {
+        let msg = CString::new(msg).map_err(|err| Error::RuntimeError(err.to_string()))?;
+        unsafe { ffi::lua_warning(self.state, msg.as_ptr(), if tocont { 1 } else { 0 }) };
+        Ok(())
+    }
+
     /// Gets information about the interpreter runtime stack.
     ///
     /// This function returns [`Debug`] structure that can be used to get information about the function
diff --git a/src/types.rs b/src/types.rs
index 2cb5af0..0f07e8c 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -4,6 +4,9 @@ use std::os::raw::{c_int, c_void};
 use std::sync::{Arc, Mutex};
 use std::{fmt, mem, ptr};
 
+#[cfg(feature = "lua54")]
+use std::ffi::CStr;
+
 #[cfg(feature = "async")]
 use futures_core::future::LocalBoxFuture;
 
@@ -53,6 +56,12 @@ pub(crate) type HookCallback = Arc<RefCell<dyn FnMut(&Lua, Debug) -> Result<()>
 #[cfg(not(feature = "send"))]
 pub(crate) type HookCallback = Arc<RefCell<dyn FnMut(&Lua, Debug) -> Result<()>>>;
 
+#[cfg(all(feature = "send", feature = "lua54"))]
+pub(crate) type WarnCallback = Box<dyn Fn(&Lua, &CStr, bool) -> Result<()> + Send>;
+
+#[cfg(all(not(feature = "send"), feature = "lua54"))]
+pub(crate) type WarnCallback = Box<dyn Fn(&Lua, &CStr, bool) -> Result<()>>;
+
 #[cfg(feature = "send")]
 pub trait MaybeSend: Send {}
 #[cfg(feature = "send")]
diff --git a/tests/tests.rs b/tests/tests.rs
index db5340d..528f17d 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -1205,3 +1205,46 @@ fn test_multi_states() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+#[cfg(feature = "lua54")]
+fn test_warnings() -> Result<()> {
+    let lua = Lua::new();
+    lua.set_app_data::<Vec<(StdString, bool)>>(Vec::new());
+
+    lua.set_warning_function(|lua, msg, tocont| {
+        let msg = msg.to_string_lossy().to_string();
+        lua.app_data_mut::<Vec<(StdString, bool)>>()
+            .unwrap()
+            .push((msg, tocont));
+        Ok(())
+    });
+
+    lua.warning("native warning ...", true)?;
+    lua.warning("finish", false)?;
+    lua.load(r#"warn("lua warning", "continue")"#).exec()?;
+
+    lua.remove_warning_function();
+    lua.warning("one more warning", false)?;
+
+    let messages = lua.app_data_ref::<Vec<(StdString, bool)>>().unwrap();
+    assert_eq!(
+        *messages,
+        vec![
+            ("native warning ...".to_string(), true),
+            ("finish".to_string(), false),
+            ("lua warning".to_string(), true),
+            ("continue".to_string(), false),
+        ]
+    );
+
+    // Trigger error inside warning
+    lua.set_warning_function(|_, _, _| Err(Error::RuntimeError("warning error".to_string())));
+    assert!(matches!(
+        lua.load(r#"warn("test")"#).exec(),
+        Err(Error::CallbackError { cause, .. })
+            if matches!(*cause, Error::RuntimeError(ref err) if err == "warning error")
+    ));
+
+    Ok(())
+}