From 4e3e44f7db90cfe71ffe4cf53b288a43fcf34eca Mon Sep 17 00:00:00 2001
From: mitchmindtree <mitchell.nordine@fuel.sh>
Date: Fri, 28 Oct 2022 02:54:41 +1100
Subject: [PATCH] Fix `TyModule::submodules_recursive` iterator (#3167)

This fixes the submodule iterator added in #3152.

Previously this wasn't actually returning any submodules, only
traversing them! This resulted in `test_fns` method not working. This
commit fixes the behaviour, tested locally as a part of #2985.
---
 sway-core/src/language/ty/module.rs | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/sway-core/src/language/ty/module.rs b/sway-core/src/language/ty/module.rs
index 96e4fb25cc9..e428f37326f 100644
--- a/sway-core/src/language/ty/module.rs
+++ b/sway-core/src/language/ty/module.rs
@@ -23,7 +23,10 @@ pub struct TySubmodule {
 /// Used rather than `impl Iterator` to enable recursive submodule iteration.
 pub struct SubmodulesRecursive<'module> {
     submods: std::slice::Iter<'module, (DepName, TySubmodule)>,
-    current: Option<Box<SubmodulesRecursive<'module>>>,
+    current: Option<(
+        &'module (DepName, TySubmodule),
+        Box<SubmodulesRecursive<'module>>,
+    )>,
 }
 
 impl TyModule {
@@ -56,17 +59,19 @@ impl<'module> Iterator for SubmodulesRecursive<'module> {
     type Item = &'module (DepName, TySubmodule);
     fn next(&mut self) -> Option<Self::Item> {
         loop {
-            self.current = match self.current.as_mut() {
+            self.current = match self.current.take() {
                 None => match self.submods.next() {
                     None => return None,
-                    Some((_, submod)) => Some(Box::new(submod.module.submodules_recursive())),
+                    Some(submod) => {
+                        Some((submod, Box::new(submod.1.module.submodules_recursive())))
+                    }
                 },
-                Some(submod) => match submod.next() {
-                    Some(submod) => return Some(submod),
-                    None => {
-                        self.current = None;
-                        continue;
+                Some((submod, mut submods)) => match submods.next() {
+                    Some(next) => {
+                        self.current = Some((submod, submods));
+                        return Some(next);
                     }
+                    None => return Some(submod),
                 },
             }
         }