Add method to TyProgram for retrieving all #[test] functions (Fue…

…lLabs#3152)

*This PR is based on FuelLabs#3151, and will be rebased on `master` once it
lands.*

Adds a method for producing an iterator that yields all test function
declarations throughout the program.

This is another step towards FuelLabs#2985.

sway-core/src/language/ty/module.rs

@@ -1,6 +1,9 @@
 use sway_types::Ident;
 
-use crate::{language::ty::*, language::DepName, semantic_analysis::namespace};
+use crate::{
+    declaration_engine::de_get_function, language::ty::*, language::DepName,
+    semantic_analysis::namespace,
+};
 
 #[derive(Clone, Debug)]
 pub struct TyModule {
@@ -14,3 +17,58 @@ pub struct TySubmodule {
     pub library_name: Ident,
     pub module: TyModule,
 }
+
+/// Iterator type for iterating over submodules.
+///
+/// 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>>>,
+}
+
+impl TyModule {
+    /// An iterator yielding all submodules recursively, depth-first.
+    pub fn submodules_recursive(&self) -> SubmodulesRecursive {
+        SubmodulesRecursive {
+            submods: self.submodules.iter(),
+            current: None,
+        }
+    }
+
+    /// All test functions within this module.
+    pub fn test_fns(&self) -> impl '_ + Iterator<Item = TyFunctionDeclaration> {
+        self.all_nodes.iter().filter_map(|node| {
+            if let TyAstNodeContent::Declaration(TyDeclaration::FunctionDeclaration(ref decl_id)) =
+                node.content
+            {
+                let fn_decl = de_get_function(decl_id.clone(), &node.span)
+                    .expect("no function declaration for ID");
+                if fn_decl.is_test() {
+                    return Some(fn_decl);
+                }
+            }
+            None
+        })
+    }
+}
+
+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() {
+                None => match self.submods.next() {
+                    None => return None,
+                    Some((_, submod)) => Some(Box::new(submod.module.submodules_recursive())),
+                },
+                Some(submod) => match submod.next() {
+                    Some(submod) => return Some(submod),
+                    None => {
+                        self.current = None;
+                        continue;
+                    }
+                },
+            }
+        }
+    }
+}

sway-core/src/language/ty/program.rs

@@ -377,6 +377,14 @@ impl TyProgram {
             })
             .collect()
     }
+
+    /// All test function declarations within the program.
+    pub fn test_fns(&self) -> impl '_ + Iterator<Item = TyFunctionDeclaration> {
+        self.root
+            .submodules_recursive()
+            .flat_map(|(_, submod)| submod.module.test_fns())
+            .chain(self.root.test_fns())
+    }
 }
 
 #[derive(Clone, Debug)]