From 9616452c00401f3f33e73b8cacea3a0c4eb58fa6 Mon Sep 17 00:00:00 2001
From: Stephen Chung <Stephen.Chung@intexact.com>
Date: Fri, 29 May 2020 00:53:30 +0800
Subject: [PATCH] Fix bug in calling script function in method style.

---
 src/engine.rs      | 28 +++++++++++++++++++++++++++-
 tests/functions.rs | 17 +++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/src/engine.rs b/src/engine.rs
index b843c85d..8903e4a3 100644
--- a/src/engine.rs
+++ b/src/engine.rs
@@ -618,8 +618,31 @@ impl Engine {
         // First search in script-defined functions (can override built-in)
         if !native_only {
             if let Some(fn_def) = lib.get(&hashes.1) {
+                let mut this_copy: Option<Dynamic>;
+                let mut this_pointer: Option<&mut Dynamic> = None;
+
+                if is_ref && !args.is_empty() {
+                    // Clone the original value.  It'll be consumed because the function
+                    // is pure and doesn't know that the first value is a reference (i.e. `is_ref`)
+                    this_copy = Some(args[0].clone());
+
+                    // Replace the first reference with a reference to the clone, force-casting the lifetime.
+                    // Keep the original reference.  Must remember to restore it before existing this function.
+                    this_pointer = Some(mem::replace(
+                        args.get_mut(0).unwrap(),
+                        unsafe_mut_cast_to_lifetime(this_copy.as_mut().unwrap()),
+                    ));
+                }
+
+                // Run scripted function
                 let result =
                     self.call_script_fn(scope, state, lib, fn_name, fn_def, args, pos, level)?;
+
+                // Restore the original reference
+                if let Some(this_pointer) = this_pointer {
+                    mem::replace(args.get_mut(0).unwrap(), this_pointer);
+                }
+
                 return Ok((result, false));
             }
         }
@@ -1577,6 +1600,7 @@ impl Engine {
                 // Normal function call - except for eval (handled above)
                 let mut arg_values: StaticVec<Dynamic>;
                 let mut args: StaticVec<_>;
+                let mut is_ref = false;
 
                 if args_expr.is_empty() {
                     // No arguments
@@ -1602,6 +1626,8 @@ impl Engine {
                             }
 
                             args = once(target).chain(arg_values.iter_mut()).collect();
+
+                            is_ref = true;
                         }
                         // func(..., ...)
                         _ => {
@@ -1617,7 +1643,7 @@ impl Engine {
 
                 let args = args.as_mut();
                 self.exec_fn_call(
-                    state, lib, name, *native, *hash, args, false, def_val, *pos, level,
+                    state, lib, name, *native, *hash, args, is_ref, def_val, *pos, level,
                 )
                 .map(|(v, _)| v)
             }
diff --git a/tests/functions.rs b/tests/functions.rs
index 26d9b387..ae12e052 100644
--- a/tests/functions.rs
+++ b/tests/functions.rs
@@ -7,12 +7,23 @@ fn test_functions() -> Result<(), Box<EvalAltResult>> {
 
     assert_eq!(engine.eval::<INT>("fn add(x, n) { x + n } add(40, 2)")?, 42);
 
+    assert_eq!(
+        engine.eval::<INT>("fn add(x, n) { x + n } let a = 40; add(a, 2); a")?,
+        40
+    );
+
     #[cfg(not(feature = "no_object"))]
     assert_eq!(
         engine.eval::<INT>("fn add(x, n) { x + n } let x = 40; x.add(2)")?,
         42
     );
 
+    #[cfg(not(feature = "no_object"))]
+    assert_eq!(
+        engine.eval::<INT>("fn add(x, n) { x += n; } let x = 40; x.add(2); x")?,
+        40
+    );
+
     assert_eq!(engine.eval::<INT>("fn mul2(x) { x * 2 } mul2(21)")?, 42);
 
     #[cfg(not(feature = "no_object"))]
@@ -21,5 +32,11 @@ fn test_functions() -> Result<(), Box<EvalAltResult>> {
         42
     );
 
+    #[cfg(not(feature = "no_object"))]
+    assert_eq!(
+        engine.eval::<INT>("fn mul2(x) { x *= 2; } let x = 21; x.mul2(); x")?,
+        21
+    );
+
     Ok(())
 }