From b1fbfcbc073d552f68e91e8d1fe94995ddbbcb30 Mon Sep 17 00:00:00 2001
From: Stephen Chung <schungx@live.com>
Date: Sat, 13 May 2023 09:31:57 +0800
Subject: [PATCH] Add take function.

---
 CHANGELOG.md              |  1 +
 src/packages/lang_core.rs | 24 ++++++++++++++++++++++
 tests/internal_fn.rs      | 42 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0291fb8d..a8567f9b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@ Enhancements
 ------------
 
 * Expressions involving `this` should now run slightly faster due to a dedicated `AST` node `ThisPtr`.
+* A `take` function is added to the standard library to take ownership of any data (replacing with `()`) in order to avoid cloning.
 
 
 Version 1.14.0
diff --git a/src/packages/lang_core.rs b/src/packages/lang_core.rs
index 26815f43..6020c7ee 100644
--- a/src/packages/lang_core.rs
+++ b/src/packages/lang_core.rs
@@ -26,6 +26,30 @@ def_package! {
 
 #[export_module]
 mod core_functions {
+    /// Take ownership of the data in a `Dynamic` value and return it.
+    /// The data is _NOT_ cloned.
+    ///
+    /// The original value is replaced with `()`.
+    ///
+    /// # Example
+    ///
+    /// ```rhai
+    /// let x = 42;
+    ///
+    /// print(take(x));         // prints 42
+    ///
+    /// print(x);               // prints ()
+    /// ```
+    #[rhai_fn(return_raw)]
+    pub fn take(value: &mut Dynamic) -> RhaiResultOf<Dynamic> {
+        if value.is_read_only() {
+            return Err(
+                ERR::ErrorNonPureMethodCallOnConstant("take".to_string(), Position::NONE).into(),
+            );
+        }
+
+        Ok(std::mem::take(value))
+    }
     /// Return the _tag_ of a `Dynamic` value.
     ///
     /// # Example
diff --git a/tests/internal_fn.rs b/tests/internal_fn.rs
index f440e479..38c06741 100644
--- a/tests/internal_fn.rs
+++ b/tests/internal_fn.rs
@@ -64,6 +64,48 @@ fn test_internal_fn() -> Result<(), Box<EvalAltResult>> {
     Ok(())
 }
 
+#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Default)]
+struct TestStruct(INT);
+
+impl Clone for TestStruct {
+    fn clone(&self) -> Self {
+        Self(self.0 + 1)
+    }
+}
+
+#[test]
+fn test_internal_fn_take() -> Result<(), Box<EvalAltResult>> {
+    let mut engine = Engine::new();
+
+    engine
+        .register_type_with_name::<TestStruct>("TestStruct")
+        .register_fn("new_ts", |x: INT| TestStruct(x));
+
+    assert_eq!(
+        engine.eval::<TestStruct>(
+            "
+                let x = new_ts(0);
+                for n in 0..41 { x = x }
+                x
+            ",
+        )?,
+        TestStruct(42)
+    );
+
+    assert_eq!(
+        engine.eval::<TestStruct>(
+            "
+                let x = new_ts(0);
+                for n in 0..41 { x = take(x) }
+                take(x)
+            ",
+        )?,
+        TestStruct(0)
+    );
+
+    Ok(())
+}
+
 #[test]
 fn test_internal_fn_big() -> Result<(), Box<EvalAltResult>> {
     let engine = Engine::new();