From dd58cdad523a5fb3f501eb2f505ed7136e55d150 Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Mon, 27 Jul 2020 20:06:53 +0100
Subject: [PATCH] Add Function::dump() to dump lua function to a binary chunk

---
 src/ffi/mod.rs    |  4 ++--
 src/function.rs   | 40 ++++++++++++++++++++++++++++++++++++++--
 src/lua.rs        | 11 ++++++++---
 tests/function.rs | 26 +++++++++++++++++++++++++-
 4 files changed, 73 insertions(+), 8 deletions(-)

diff --git a/src/ffi/mod.rs b/src/ffi/mod.rs
index 78059a2..771dd2f 100644
--- a/src/ffi/mod.rs
+++ b/src/ffi/mod.rs
@@ -211,8 +211,8 @@ pub use self::lua::{
     LUA_HOOKCOUNT, LUA_HOOKLINE, LUA_HOOKRET, LUA_HOOKTAILCALL, LUA_MASKCALL, LUA_MASKCOUNT,
     LUA_MASKLINE, LUA_MASKRET, LUA_MINSTACK, LUA_MULTRET, LUA_OK, LUA_OPADD, LUA_OPDIV, LUA_OPEQ,
     LUA_OPLE, LUA_OPLT, LUA_OPMOD, LUA_OPMUL, LUA_OPPOW, LUA_OPSUB, LUA_OPUNM, LUA_REGISTRYINDEX,
-    LUA_TBOOLEAN, LUA_TFUNCTION, LUA_TLIGHTUSERDATA, LUA_TNIL, LUA_TNONE, LUA_TNUMBER, LUA_TSTRING,
-    LUA_TTABLE, LUA_TTHREAD, LUA_TUSERDATA, LUA_YIELD,
+    LUA_SIGNATURE, LUA_TBOOLEAN, LUA_TFUNCTION, LUA_TLIGHTUSERDATA, LUA_TNIL, LUA_TNONE,
+    LUA_TNUMBER, LUA_TSTRING, LUA_TTABLE, LUA_TTHREAD, LUA_TUSERDATA, LUA_YIELD,
 };
 
 #[cfg(any(feature = "lua54", feature = "lua53"))]
diff --git a/src/function.rs b/src/function.rs
index fe98c5a..45ad4e1 100644
--- a/src/function.rs
+++ b/src/function.rs
@@ -1,5 +1,5 @@
-use std::os::raw::c_int;
-use std::ptr;
+use std::os::raw::{c_int, c_void};
+use std::{ptr, slice};
 
 use crate::error::{Error, Result};
 use crate::ffi;
@@ -205,6 +205,42 @@ impl<'lua> Function<'lua> {
             Ok(Function(lua.pop_ref()))
         }
     }
+
+    /// Dumps the function as a binary chunk.
+    ///
+    /// If `strip` is true, the binary representation may not include all debug information
+    /// about the function, to save space.
+    pub fn dump(&self, strip: bool) -> Result<Vec<u8>> {
+        unsafe extern "C" fn writer(
+            _state: *mut ffi::lua_State,
+            buf: *const c_void,
+            buf_len: usize,
+            data: *mut c_void,
+        ) -> c_int {
+            let data = &mut *(data as *mut Vec<u8>);
+            let buf = slice::from_raw_parts(buf as *const u8, buf_len);
+            data.extend_from_slice(buf);
+            0
+        }
+
+        let lua = self.0.lua;
+        let mut data: Vec<u8> = Vec::new();
+        unsafe {
+            let _sg = StackGuard::new(lua.state);
+            assert_stack(lua.state, 1);
+            lua.push_ref(&self.0);
+            let strip = if strip { 1 } else { 0 };
+            ffi::lua_dump(
+                lua.state,
+                writer,
+                &mut data as *mut Vec<u8> as *mut c_void,
+                strip,
+            );
+            ffi::lua_pop(lua.state, 1);
+        }
+
+        Ok(data)
+    }
 }
 
 impl<'lua> PartialEq for Function<'lua> {
diff --git a/src/lua.rs b/src/lua.rs
index 6658fc9..665b602 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -1869,10 +1869,13 @@ impl<'lua, 'a> Chunk<'lua, 'a> {
     /// the value that it evaluates to.  Otherwise, the chunk is interpreted as a block as normal,
     /// and this is equivalent to calling `exec`.
     pub fn eval<R: FromLuaMulti<'lua>>(self) -> Result<R> {
-        // First, try interpreting the lua as an expression by adding
+        // Bytecode is always interpreted as a statement.
+        // For source code, first try interpreting the lua as an expression by adding
         // "return", then as a statement.  This is the same thing the
         // actual lua repl does.
-        if let Ok(function) = self.lua.load_chunk(
+        if self.source.starts_with(ffi::LUA_SIGNATURE) {
+            self.call(())
+        } else if let Ok(function) = self.lua.load_chunk(
             &self.expression_source(),
             self.name.as_ref(),
             self.env.clone(),
@@ -1896,7 +1899,9 @@ impl<'lua, 'a> Chunk<'lua, 'a> {
         'lua: 'fut,
         R: FromLuaMulti<'lua> + 'fut,
     {
-        if let Ok(function) = self.lua.load_chunk(
+        if self.source.starts_with(ffi::LUA_SIGNATURE) {
+            self.call_async(())
+        } else if let Ok(function) = self.lua.load_chunk(
             &self.expression_source(),
             self.name.as_ref(),
             self.env.clone(),
diff --git a/tests/function.rs b/tests/function.rs
index f884d88..7ed8d33 100644
--- a/tests/function.rs
+++ b/tests/function.rs
@@ -10,7 +10,7 @@
 )]
 extern "system" {}
 
-use mlua::{Function, Lua, Result, String};
+use mlua::{Error, Function, Lua, Result, String};
 
 #[test]
 fn test_function() -> Result<()> {
@@ -87,3 +87,27 @@ fn test_rust_function() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_dump() -> Result<()> {
+    let lua = unsafe { Lua::unsafe_new() };
+
+    let concat_tmp = lua
+        .load(r#"function(arg1, arg2) return arg1 .. arg2 end"#)
+        .eval::<Function>()?;
+    let concat_bytecode = concat_tmp.dump(false)?;
+    let concat = lua.load(&concat_bytecode).into_function()?;
+
+    assert_eq!(concat.call::<_, String>(("foo", "bar"))?, "foobar");
+
+    let lua = Lua::new();
+    match lua.load(&concat_bytecode).exec() {
+        Ok(_) => panic!("expected SyntaxError, got no error"),
+        Err(Error::SyntaxError { message: msg, .. }) => {
+            assert!(msg.contains("attempt to load a binary chunk"))
+        }
+        Err(e) => panic!("expected SyntaxError, got {:?}", e),
+    }
+
+    Ok(())
+}