From e3aa2c19ce3ae98d993f5c2929a3157eeb1b9e9e Mon Sep 17 00:00:00 2001
From: Stephen Chung <schungx@live.com>
Date: Fri, 31 Dec 2021 15:59:13 +0800
Subject: [PATCH] Make for loops for efficient.

---
 scripts/for3.rhai | 26 +++++++++++++++++++++++++
 src/engine.rs     | 49 ++++++++++++++++++++++++++++++-----------------
 tests/for.rs      | 28 +++++++++++++++++++++++++++
 3 files changed, 85 insertions(+), 18 deletions(-)
 create mode 100644 scripts/for3.rhai

diff --git a/scripts/for3.rhai b/scripts/for3.rhai
new file mode 100644
index 00000000..de799684
--- /dev/null
+++ b/scripts/for3.rhai
@@ -0,0 +1,26 @@
+const MAX = 100;
+const CHECK = ((MAX - 1) ** 2) * MAX;
+
+print("Ready... Go!");
+
+let now = timestamp();
+
+print(`Creating ${MAX} closures...`);
+
+let list = [];
+
+for i in 0..MAX {
+    list.push(|| i ** 2);
+}
+
+print(`Time = ${now.elapsed} seconds...`);
+print(`Summing ${MAX} closures...`);
+
+let sum = 0;
+
+for f in list {
+    sum += f.call();
+}
+
+print(`Sum = ${sum} (should be ${CHECK})`);
+print(`Finished. Total run time = ${now.elapsed} seconds.`);
diff --git a/src/engine.rs b/src/engine.rs
index fd4104f8..90f06b62 100644
--- a/src/engine.rs
+++ b/src/engine.rs
@@ -2892,16 +2892,18 @@ impl Engine {
                 if let Some(func) = func {
                     // Add the loop variables
                     let orig_scope_len = scope.len();
-                    let counter_index = counter.as_ref().map(|Ident { name, .. }| {
-                        scope.push(unsafe_cast_var_name_to_lifetime(name), 0 as INT);
+                    let counter_index = if let Some(counter) = counter {
+                        scope.push(unsafe_cast_var_name_to_lifetime(&counter.name), 0 as INT);
                         scope.len() - 1
-                    });
+                    } else {
+                        usize::MAX
+                    };
                     scope.push(unsafe_cast_var_name_to_lifetime(name), ());
                     let index = scope.len() - 1;
 
                     for (x, iter_value) in func(iter_obj).enumerate() {
                         // Increment counter
-                        if let Some(c) = counter_index {
+                        if counter_index < usize::MAX {
                             #[cfg(not(feature = "unchecked"))]
                             if x > INT::MAX as usize {
                                 return Err(ERR::ErrorArithmetic(
@@ -2911,26 +2913,37 @@ impl Engine {
                                 .into());
                             }
 
-                            let mut counter_var = scope
-                                .get_mut_by_index(c)
-                                .write_lock::<INT>()
-                                .expect("`INT`");
-                            *counter_var = x as INT;
+                            let index_value = (x as INT).into();
+
+                            #[cfg(not(feature = "no_closure"))]
+                            {
+                                let index_var = scope.get_mut_by_index(counter_index);
+                                if index_var.is_shared() {
+                                    *index_var.write_lock().expect("`Dynamic`") = index_value;
+                                } else {
+                                    *index_var = index_value;
+                                }
+                            }
+                            #[cfg(feature = "no_closure")]
+                            {
+                                *scope.get_mut_by_index(counter_index) = index_value;
+                            }
                         }
 
-                        let loop_var = scope.get_mut_by_index(index);
                         let value = iter_value.flatten();
 
                         #[cfg(not(feature = "no_closure"))]
-                        let loop_var_is_shared = loop_var.is_shared();
+                        {
+                            let loop_var = scope.get_mut_by_index(index);
+                            if loop_var.is_shared() {
+                                *loop_var.write_lock().expect("`Dynamic`") = value;
+                            } else {
+                                *loop_var = value;
+                            }
+                        }
                         #[cfg(feature = "no_closure")]
-                        let loop_var_is_shared = false;
-
-                        if loop_var_is_shared {
-                            let mut value_ref = loop_var.write_lock().expect("`Dynamic`");
-                            *value_ref = value;
-                        } else {
-                            *loop_var = value;
+                        {
+                            *scope.get_mut_by_index(index) = value;
                         }
 
                         #[cfg(not(feature = "unchecked"))]
diff --git a/tests/for.rs b/tests/for.rs
index 59e7223a..7cbf401a 100644
--- a/tests/for.rs
+++ b/tests/for.rs
@@ -371,3 +371,31 @@ fn test_for_module_iterator() -> Result<(), Box<EvalAltResult>> {
 
     Ok(())
 }
+
+#[test]
+#[cfg(not(feature = "no_closure"))]
+fn test_for_capture() -> Result<(), Box<EvalAltResult>> {
+    let engine = Engine::new();
+
+    assert_eq!(
+        engine.eval::<INT>(
+            "
+                let a = [];
+
+                for (x, i) in 100..110 {
+                    a += || i + x;
+                }
+
+                let sum = 0;
+
+                for fp in a {
+                    sum += fp.call();
+                }
+
+                sum
+            "
+        )?,
+        1180
+    );
+    Ok(())
+}