From 799dd9d9d15514d54637f09b56041d336d8b8b9d Mon Sep 17 00:00:00 2001
From: Stephen Chung <schungx@live.com>
Date: Mon, 22 Aug 2022 13:10:32 +0800
Subject: [PATCH] Add fallible type iterators.

---
 CHANGELOG.md       |  8 ++++++++
 src/eval/stmt.rs   | 39 +++++++++++----------------------------
 src/func/native.rs |  5 +++--
 src/module/mod.rs  | 15 ++++++++++++++-
 4 files changed, 36 insertions(+), 31 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index abcec383..d9aecc83 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,14 @@ Bug fixes
 * API for registering property getters/setters and indexers to an `Engine` now works with functions that take a first parameter of `NativeCallContext`.
 * Missing API function `Module::set_getter_setter_fn` is added.
 
+New features
+------------
+
+Fallible type iterators
+-----------------------
+
+* For very special needs, the ability to register fallible type iterators is added.
+
 
 Version 1.9.0
 =============
diff --git a/src/eval/stmt.rs b/src/eval/stmt.rs
index 8e8ab47c..18d16eba 100644
--- a/src/eval/stmt.rs
+++ b/src/eval/stmt.rs
@@ -641,38 +641,21 @@ impl Engine {
                                     break;
                                 }
 
-                                let index_value = (x as INT).into();
+                                let index_value = Dynamic::from(x as INT);
 
-                                #[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;
-                                }
+                                *scope.get_mut_by_index(counter_index).write_lock().unwrap() =
+                                    index_value;
                             }
 
-                            let value = iter_value.flatten();
-
-                            #[cfg(not(feature = "no_closure"))]
-                            {
-                                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;
+                            let value = match iter_value {
+                                Ok(v) => v.flatten(),
+                                Err(err) => {
+                                    loop_result = Err(err.fill_position(expr.position()));
+                                    break;
                                 }
-                            }
-                            #[cfg(feature = "no_closure")]
-                            {
-                                *scope.get_mut_by_index(index) = value;
-                            }
+                            };
+
+                            *scope.get_mut_by_index(index).write_lock().unwrap() = value;
 
                             #[cfg(not(feature = "unchecked"))]
                             if let Err(err) = self
diff --git a/src/func/native.rs b/src/func/native.rs
index 6b3d46de..26801793 100644
--- a/src/func/native.rs
+++ b/src/func/native.rs
@@ -426,10 +426,11 @@ pub type FnBuiltin = fn(NativeCallContext, &mut FnCallArgs) -> RhaiResult;
 
 /// Function that gets an iterator from a type.
 #[cfg(not(feature = "sync"))]
-pub type IteratorFn = dyn Fn(Dynamic) -> Box<dyn Iterator<Item = Dynamic>>;
+pub type IteratorFn = dyn Fn(Dynamic) -> Box<dyn Iterator<Item = RhaiResultOf<Dynamic>>>;
 /// Function that gets an iterator from a type.
 #[cfg(feature = "sync")]
-pub type IteratorFn = dyn Fn(Dynamic) -> Box<dyn Iterator<Item = Dynamic>> + Send + Sync;
+pub type IteratorFn =
+    dyn Fn(Dynamic) -> Box<dyn Iterator<Item = RhaiResultOf<Dynamic>>> + Send + Sync;
 
 #[cfg(not(feature = "sync"))]
 pub type FnPlugin = dyn PluginFunction;
diff --git a/src/module/mod.rs b/src/module/mod.rs
index ef496427..79a43c64 100644
--- a/src/module/mod.rs
+++ b/src/module/mod.rs
@@ -2251,11 +2251,24 @@ impl Module {
 
     /// Set a type iterator into the [`Module`].
     #[cfg(not(feature = "sync"))]
-    #[inline]
+    #[inline(always)]
     pub fn set_iter(
         &mut self,
         type_id: TypeId,
         func: impl Fn(Dynamic) -> Box<dyn Iterator<Item = Dynamic>> + 'static,
+    ) -> &mut Self {
+        self.set_iter_result(type_id, move |x| {
+            Box::new(func(x).map(Ok)) as Box<dyn Iterator<Item = RhaiResultOf<Dynamic>>>
+        })
+    }
+
+    /// Set a fallible type iterator into the [`Module`].
+    #[cfg(not(feature = "sync"))]
+    #[inline]
+    pub fn set_iter_result(
+        &mut self,
+        type_id: TypeId,
+        func: impl Fn(Dynamic) -> Box<dyn Iterator<Item = RhaiResultOf<Dynamic>>> + 'static,
     ) -> &mut Self {
         let func = Shared::new(func);
         if self.indexed {