From 82a455b4d7080dca02f3af51320c398c1275d59b Mon Sep 17 00:00:00 2001
From: Stephen Chung <schungx@live.com>
Date: Tue, 13 Dec 2022 13:35:23 +0800
Subject: [PATCH] Add Dynamic::deep_scan.

---
 CHANGELOG.md         |  1 +
 src/module/mod.rs    | 28 ++++++----------------------
 src/types/dynamic.rs | 21 +++++++++++++++++++++
 3 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fb68e3f4..db5d1fec 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -63,6 +63,7 @@ Enhancements
 * Expression nesting levels is refined such that it grows less excessively for common patterns.
 * The traits `Index` and `IndexMut` are added to `FnPtr`.
 * `FnPtr::iter_curry` and `FnPtr::iter_curry_mut` are added.
+* `Dynamic::deep_scan` is added to recursively scan for `Dynamic` values.
 
 
 Version 1.11.0
diff --git a/src/module/mod.rs b/src/module/mod.rs
index 241ded4f..c1f8f54f 100644
--- a/src/module/mod.rs
+++ b/src/module/mod.rs
@@ -9,7 +9,7 @@ use crate::func::{
 };
 use crate::types::{dynamic::Variant, BloomFilterU64, CustomTypesCollection};
 use crate::{
-    calc_fn_hash, calc_fn_hash_full, Dynamic, FnArgsVec, Identifier, ImmutableString,
+    calc_fn_hash, calc_fn_hash_full, Dynamic, FnArgsVec, FnPtr, Identifier, ImmutableString,
     NativeCallContext, RhaiResultOf, Shared, SharedModule, SmartString,
 };
 use bitflags::bitflags;
@@ -2147,29 +2147,13 @@ impl Module {
             constants,
         });
 
-        fn update_encapsulated_environ(
-            value: &mut Dynamic,
-            environ: &Shared<crate::func::EncapsulatedEnviron>,
-        ) {
-            match value.0 {
-                #[cfg(not(feature = "no_index"))]
-                crate::types::dynamic::Union::Array(ref mut a, _, _) => a
-                    .iter_mut()
-                    .for_each(|v| update_encapsulated_environ(v, environ)),
-                #[cfg(not(feature = "no_object"))]
-                crate::types::dynamic::Union::Map(ref mut map, _, _) => map
-                    .values_mut()
-                    .for_each(|v| update_encapsulated_environ(v, environ)),
-                crate::types::dynamic::Union::FnPtr(ref mut fn_ptr, _, _) => {
-                    fn_ptr.set_encapsulated_environ(Some(environ.clone()))
-                }
-                _ => (),
-            }
-        }
-
         // Variables with an alias left in the scope become module variables
         for (_name, mut value, mut aliases) in scope {
-            update_encapsulated_environ(&mut value, &environ);
+            value.deep_scan(|v| {
+                if let Some(fn_ptr) = v.downcast_mut::<FnPtr>() {
+                    fn_ptr.set_encapsulated_environ(Some(environ.clone()));
+                }
+            });
 
             match aliases.len() {
                 0 => (),
diff --git a/src/types/dynamic.rs b/src/types/dynamic.rs
index 4c660af8..47423128 100644
--- a/src/types/dynamic.rs
+++ b/src/types/dynamic.rs
@@ -2084,6 +2084,27 @@ impl Dynamic {
             _ => Err(self.type_name()),
         }
     }
+
+    /// Recursively scan for [`Dynamic`] values within this [`Dynamic`] (e.g. items in an array or map),
+    /// calling a filter function on each.
+    ///
+    /// Shared values are _NOT_ scanned.
+    #[inline]
+    pub fn deep_scan(&mut self, mut filter: impl FnMut(&mut Self)) {
+        fn scan_inner(value: &mut Dynamic, filter: &mut impl FnMut(&mut Dynamic)) {
+            match &mut value.0 {
+                #[cfg(not(feature = "no_index"))]
+                Union::Array(a, ..) => a.iter_mut().for_each(|v| scan_inner(v, filter)),
+                #[cfg(not(feature = "no_object"))]
+                Union::Map(m, ..) => m.values_mut().for_each(|v| scan_inner(v, filter)),
+                Union::FnPtr(f, ..) => f.iter_curry_mut().for_each(|v| scan_inner(v, filter)),
+                _ => (),
+            }
+        }
+
+        filter(self);
+        scan_inner(self, &mut filter);
+    }
 }
 
 impl From<()> for Dynamic {