Refactor lexical scope walking into a new walk_scope_chain.

FuelLabs · Dec 2, 2024 · e18803c · e18803c
1 parent d07f1af
commit e18803c
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 55 deletions.
diff --git a/sway-core/src/semantic_analysis/ast_node/expression/typed_expression.rs b/sway-core/src/semantic_analysis/ast_node/expression/typed_expression.rs
@@ -2416,31 +2416,26 @@ impl ty::TyExpression {
         base_name: &Ident,
         projections: &[ty::ProjectionKind],
     ) -> Result<(TypeId, TypeId), ErrorEmitted> {
-        let mut lexical_scope_opt = Some(module.current_lexical_scope());
-        while let Some(lexical_scope) = lexical_scope_opt {
-            let result = Self::find_subfield_type_helper(
+        let ret = module.walk_scope_chain(|lexical_scope| {
+            Self::find_subfield_type_helper(
                 lexical_scope,
                 handler,
                 engines,
                 namespace,
                 base_name,
                 projections,
-            )?;
-            if let Some(result) = result {
-                return Ok(result);
-            }
-            if let Some(parent_scope_id) = lexical_scope.parent {
-                lexical_scope_opt = module.get_lexical_scope(parent_scope_id);
-            } else {
-                lexical_scope_opt = None;
-            }
-        }
+            )
+        })?;
 
-        // Symbol not found
-        Err(handler.emit_err(CompileError::UnknownVariable {
-            var_name: base_name.clone(),
-            span: base_name.span(),
-        }))
+        if let Some(ret) = ret {
+            Ok(ret)
+        } else {
+            // Symbol not found
+            Err(handler.emit_err(CompileError::UnknownVariable {
+                var_name: base_name.clone(),
+                span: base_name.span(),
+            }))
+        }
     }
 
     /// Returns a tuple where the first element is the [TypeId] of the actual expression, and

diff --git a/sway-core/src/semantic_analysis/namespace/module.rs b/sway-core/src/semantic_analysis/namespace/module.rs
@@ -299,26 +299,40 @@ impl Module {
         self.current_lexical_scope_id = parent_scope_id.unwrap_or(0);
     }
 
+    pub fn walk_scope_chain<T>(
+        &self,
+        mut f: impl FnMut(&LexicalScope) -> Result<Option<T>, ErrorEmitted>,
+    ) -> Result<Option<T>, ErrorEmitted> {
+        let mut lexical_scope_opt = Some(self.current_lexical_scope());
+        while let Some(lexical_scope) = lexical_scope_opt {
+            let result = f(lexical_scope)?;
+            if let Some(result) = result {
+                return Ok(Some(result));
+            }
+            if let Some(parent_scope_id) = lexical_scope.parent {
+                lexical_scope_opt = self.get_lexical_scope(parent_scope_id);
+            } else {
+                lexical_scope_opt = None;
+            }
+        }
+        Ok(None)
+    }
+
     pub fn get_items_for_type(
         &self,
         engines: &Engines,
         type_id: TypeId,
     ) -> Vec<ResolvedTraitImplItem> {
-        let mut lexical_scope_opt = Some(self.current_lexical_scope());
         let mut vec = vec![];
-        while let Some(lexical_scope) = lexical_scope_opt {
+        let _ = self.walk_scope_chain(|lexical_scope| {
             vec.extend(
                 lexical_scope
                     .items
                     .implemented_traits
                     .get_items_for_type(engines, type_id),
             );
-            if let Some(parent_scope_id) = lexical_scope.parent {
-                lexical_scope_opt = self.get_lexical_scope(parent_scope_id);
-            } else {
-                lexical_scope_opt = None;
-            }
-        }
+            Ok(Some(()))
+        });
         vec
     }
 
@@ -328,26 +342,19 @@ impl Module {
         engines: &Engines,
         symbol: &Ident,
     ) -> Result<ResolvedDeclaration, ErrorEmitted> {
-        let mut lexical_scope_opt = Some(self.current_lexical_scope());
-        while let Some(lexical_scope) = lexical_scope_opt {
-            let result = lexical_scope
-                .items
-                .resolve_symbol(handler, engines, symbol)?;
-            if let Some(result) = result {
-                return Ok(result);
-            }
-            if let Some(parent_scope_id) = lexical_scope.parent {
-                lexical_scope_opt = self.get_lexical_scope(parent_scope_id);
-            } else {
-                lexical_scope_opt = None;
-            }
+        let ret = self.walk_scope_chain(|lexical_scope| {
+            lexical_scope.items.resolve_symbol(handler, engines, symbol)
+        })?;
+
+        if let Some(ret) = ret {
+            Ok(ret)
+        } else {
+            // Symbol not found
+            Err(handler.emit_err(CompileError::SymbolNotFound {
+                name: symbol.clone(),
+                span: symbol.span(),
+            }))
         }
-
-        // Symbol not found
-        Err(handler.emit_err(CompileError::SymbolNotFound {
-            name: symbol.clone(),
-            span: symbol.span(),
-        }))
     }
 
     pub fn get_methods_for_type(
@@ -358,7 +365,7 @@ impl Module {
         self.get_items_for_type(engines, type_id)
             .into_iter()
             .filter_map(|item| match item {
-                ResolvedTraitImplItem::Parsed(_) => todo!(),
+                ResolvedTraitImplItem::Parsed(_) => unreachable!(),
                 ResolvedTraitImplItem::Typed(item) => match item {
                     ty::TyTraitItem::Fn(decl_ref) => Some(ResolvedFunctionDecl::Typed(decl_ref)),
                     ty::TyTraitItem::Constant(_decl_ref) => None,

diff --git a/sway-core/src/semantic_analysis/namespace/trait_map.rs b/sway-core/src/semantic_analysis/namespace/trait_map.rs
@@ -1390,22 +1390,16 @@ impl TraitMap {
         type_id: TypeId,
         engines: &Engines,
     ) -> BTreeSet<(Ident, TypeId)> {
-        let mut lexical_scope_opt = Some(module.current_lexical_scope());
         let mut all_impld_traits: BTreeSet<(Ident, TypeId)> = Default::default();
-        while let Some(lexical_scope) = lexical_scope_opt {
+        let _ = module.walk_scope_chain(|lexical_scope| {
             all_impld_traits.extend(
                 lexical_scope
                     .items
                     .implemented_traits
                     .get_implemented_traits(type_id, engines),
             );
-            if let Some(parent_scope_id) = lexical_scope.parent {
-                lexical_scope_opt = module.get_lexical_scope(parent_scope_id);
-            } else {
-                lexical_scope_opt = None;
-            }
-        }
-
+            Ok(Some(()))
+        });
         all_impld_traits
     }