Fixes Numeric type propagation

To fix the Numeric type propagation we now do the type checking of code
blocks twice, first to collect all the unification and second to unify
the types of a previous namespace with the variable declarations types
of the current namespace.

This changes incurs a significative performance drop in the compilation time.
The test in release goes from 3 secs to 5 secs.

This can be improved significatively when we start using LexicalScopes
instead of cloning the whole namespace on each scope.

Fixes #6371

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

@@ -143,7 +143,13 @@ impl TypeCheckAnalysis for TyExpression {
                             continue;
                         }
 
-                        unify.unify(handler, element.return_type, *elem_type, &element.span)
+                        unify.unify(
+                            handler,
+                            element.return_type,
+                            *elem_type,
+                            &element.span,
+                            true,
+                        )
                     }
                 }
             }

sway-core/src/semantic_analysis/ast_node/code_block.rs

@@ -23,21 +23,70 @@ impl ty::TyCodeBlock {
 
     pub(crate) fn type_check(
         handler: &Handler,
-        ctx: TypeCheckContext,
+        mut ctx: TypeCheckContext,
         code_block: &CodeBlock,
+        clear_unifications: bool,
     ) -> Result<Self, ErrorEmitted> {
-        ctx.scoped(|mut ctx| {
-            let evaluated_contents = code_block
-                .contents
-                .iter()
-                .filter_map(|node| ty::TyAstNode::type_check(handler, ctx.by_ref(), node).ok())
-                .collect::<Vec<ty::TyAstNode>>();
+        // We should clear unifications only in the root code blocks
+        if clear_unifications {
+            ctx.engines.te().clear_unifications();
+        }
+
+        // While collecting unifications we type_check the code block once without error handling.
+        if ctx.collecting_unifications() {
+            let code_block_result = ctx.by_ref().scoped(|mut ctx| {
+                let evaluated_contents = code_block
+                    .contents
+                    .iter()
+                    .filter_map(|node| {
+                        ty::TyAstNode::type_check(&Handler::default(), ctx.by_ref(), node).ok()
+                    })
+                    .collect::<Vec<ty::TyAstNode>>();
+                Ok(ty::TyCodeBlock {
+                    contents: evaluated_contents,
+                    whole_block_span: code_block.whole_block_span.clone(),
+                })
+            })?;
+
+            return Ok(code_block_result);
+        }
 
-            Ok(ty::TyCodeBlock {
-                contents: evaluated_contents,
-                whole_block_span: code_block.whole_block_span.clone(),
+        // We are typechecking the code block AST nodes twice.
+        // The first pass does all the unifications to the variables types.
+        // In the second pass we use the previous_namespace on variable declaration to unify directly with the result of the first pass.
+        // This is required to fix the test case numeric_type_propagation and issue #6371
+        let (_, previous_namespace) = ctx
+            .by_ref()
+            .with_collecting_unifications()
+            .scoped_and_namespace(|mut ctx| {
+                let evaluated_contents = code_block
+                    .contents
+                    .iter()
+                    .filter_map(|node| {
+                        ty::TyAstNode::type_check(&Handler::default(), ctx.by_ref(), node).ok()
+                    })
+                    .collect::<Vec<ty::TyAstNode>>();
+                Ok(ty::TyCodeBlock {
+                    contents: evaluated_contents,
+                    whole_block_span: code_block.whole_block_span.clone(),
+                })
+            })?;
+
+        ctx.engines.te().reapply_unifications(ctx.engines());
+
+        ctx.by_ref()
+            .scoped_with_previous_namespace(&previous_namespace, |mut ctx| {
+                let evaluated_contents = code_block
+                    .contents
+                    .iter()
+                    .filter_map(|node| ty::TyAstNode::type_check(handler, ctx.by_ref(), node).ok())
+                    .collect::<Vec<ty::TyAstNode>>();
+
+                Ok(ty::TyCodeBlock {
+                    contents: evaluated_contents,
+                    whole_block_span: code_block.whole_block_span.clone(),
+                })
             })
-        })
     }
 
     pub fn compute_return_type_and_span(

sway-core/src/semantic_analysis/ast_node/declaration/declaration.rs

@@ -11,7 +11,7 @@ use crate::{
         ty::{self, FunctionDecl, TyCodeBlock, TyDecl, TyStorageField},
         CallPath,
     },
-    namespace::{IsExtendingExistingImpl, IsImplSelf},
+    namespace::{IsExtendingExistingImpl, IsImplSelf, ResolvedDeclaration},
     semantic_analysis::{
         symbol_collection_context::SymbolCollectionContext,
         type_check_context::EnforceTypeArguments, ConstShadowingMode, GenericShadowingMode,
@@ -162,6 +162,29 @@ impl TyDecl {
                     },
                 };
 
+                if let Some(previous_namespace) = ctx.previous_namespace() {
+                    let previous_symbol = previous_namespace
+                        .module(engines)
+                        .current_items()
+                        .check_symbol_unique(&var_decl.name.clone())
+                        .ok();
+
+                    if let Some(ResolvedDeclaration::Typed(ty::TyDecl::VariableDecl(
+                        variable_decl,
+                    ))) = previous_symbol
+                    {
+                        type_engine.unify(
+                            handler,
+                            engines,
+                            body.return_type,
+                            variable_decl.body.return_type,
+                            &decl.span(engines),
+                            "",
+                            None,
+                        );
+                    }
+                }
+
                 let typed_var_decl = ty::TyDecl::VariableDecl(Box::new(ty::TyVariableDecl {
                     name: var_decl.name.clone(),
                     body,

sway-core/src/semantic_analysis/ast_node/declaration/function.rs

@@ -215,7 +215,7 @@ impl ty::TyFunctionDecl {
                     .with_type_annotation(return_type.type_id)
                     .with_function_type_annotation(return_type.type_id);
 
-                let body = ty::TyCodeBlock::type_check(handler, ctx.by_ref(), body)
+                let body = ty::TyCodeBlock::type_check(handler, ctx.by_ref(), body, true)
                     .unwrap_or_else(|_err| ty::TyCodeBlock::default());
 
                 ty_fn_decl.body = body;

sway-core/src/semantic_analysis/ast_node/expression/typed_expression.rs

@@ -636,7 +636,7 @@ impl ty::TyExpression {
         let engines = ctx.engines();
 
         let (typed_block, block_return_type) =
-            match ty::TyCodeBlock::type_check(handler, ctx.by_ref(), contents) {
+            match ty::TyCodeBlock::type_check(handler, ctx.by_ref(), contents, false) {
                 Ok(res) => {
                     let (block_type, _span) = TyCodeBlock::compute_return_type_and_span(&ctx, &res);
                     (res, block_type)
@@ -2038,7 +2038,7 @@ impl ty::TyExpression {
                  assigning it to a mutable variable declared outside of the loop \
                  instead.",
         );
-        let typed_body = ty::TyCodeBlock::type_check(handler, ctx.by_ref(), body)?;
+        let typed_body = ty::TyCodeBlock::type_check(handler, ctx.by_ref(), body, false)?;
 
         let exp = ty::TyExpression {
             expression: ty::TyExpressionVariant::WhileLoop {

sway-core/src/semantic_analysis/namespace/lexical_scope.rs

@@ -17,7 +17,7 @@ use sway_error::{
     error::{CompileError, ShadowingSource, StructFieldUsageContext},
     handler::{ErrorEmitted, Handler},
 };
-use sway_types::{span::Span, Spanned};
+use sway_types::{span::Span, IdentUnique, Spanned};
 
 use std::sync::Arc;
 
@@ -36,6 +36,7 @@ impl ResolvedFunctionDecl {
 }
 
 pub(super) type SymbolMap = im::OrdMap<Ident, ResolvedDeclaration>;
+pub(super) type SymbolUniqueMap = im::OrdMap<IdentUnique, ResolvedDeclaration>;
 
 type SourceIdent = Ident;
 
@@ -76,6 +77,8 @@ pub struct LexicalScope {
 pub struct Items {
     /// An ordered map from `Ident`s to their associated declarations.
     pub(crate) symbols: SymbolMap,
+    /// An ordered map from `IdentUnique`s to their associated declarations.
+    pub(crate) symbols_unique: SymbolUniqueMap,
     pub(crate) implemented_traits: TraitMap,
     /// Contains symbols imported using star imports (`use foo::*`.).
     ///
@@ -629,6 +632,8 @@ impl Items {
             );
         }
 
+        self.symbols_unique
+            .insert(name.clone().into(), item.clone());
         self.symbols.insert(name, item);
 
         Ok(())
@@ -684,6 +689,19 @@ impl Items {
             })
     }
 
+    pub(crate) fn check_symbol_unique(
+        &self,
+        name: &Ident,
+    ) -> Result<ResolvedDeclaration, CompileError> {
+        self.symbols_unique
+            .get(&name.into())
+            .cloned()
+            .ok_or_else(|| CompileError::SymbolNotFound {
+                name: name.clone(),
+                span: name.span(),
+            })
+    }
+
     pub fn get_items_for_type(
         &self,
         engines: &Engines,

sway-core/src/semantic_analysis/type_check_context.rs

@@ -94,6 +94,13 @@ pub struct TypeCheckContext<'a> {
 
     /// Indicates when semantic analysis is type checking storage declaration.
     storage_declaration: bool,
+
+    /// The namespace context accumulated throughout a previous iteration of type-checking.
+    /// This is required because we are typechecking code blocks twice and relying on the
+    /// previous_namespace to unify the types of variable declarations on the second pass.
+    previous_namespace: Option<&'a Namespace>,
+
+    collecting_unifications: bool,
 }
 
 impl<'a> TypeCheckContext<'a> {
@@ -118,7 +125,9 @@ impl<'a> TypeCheckContext<'a> {
             kind: TreeType::Contract,
             disallow_functions: false,
             storage_declaration: false,
+            previous_namespace: None,
             experimental,
+            collecting_unifications: false,
         }
     }
 
@@ -157,7 +166,9 @@ impl<'a> TypeCheckContext<'a> {
             kind: TreeType::Contract,
             disallow_functions: false,
             storage_declaration: false,
+            previous_namespace: None,
             experimental,
+            collecting_unifications: false,
         }
     }
 
@@ -185,7 +196,9 @@ impl<'a> TypeCheckContext<'a> {
             engines: self.engines,
             disallow_functions: self.disallow_functions,
             storage_declaration: self.storage_declaration,
+            previous_namespace: self.previous_namespace,
             experimental: self.experimental,
+            collecting_unifications: self.collecting_unifications,
         }
     }
 
@@ -210,7 +223,9 @@ impl<'a> TypeCheckContext<'a> {
             engines: self.engines,
             disallow_functions: self.disallow_functions,
             storage_declaration: self.storage_declaration,
+            previous_namespace: None,
             experimental: self.experimental,
+            collecting_unifications: self.collecting_unifications,
         };
         with_scoped_ctx(ctx)
     }
@@ -236,11 +251,41 @@ impl<'a> TypeCheckContext<'a> {
             engines: self.engines,
             disallow_functions: self.disallow_functions,
             storage_declaration: self.storage_declaration,
+            previous_namespace: None,
             experimental: self.experimental,
+            collecting_unifications: self.collecting_unifications,
         };
         Ok((with_scoped_ctx(ctx)?, namespace))
     }
 
+    pub fn scoped_with_previous_namespace<T>(
+        self,
+        previous_namespace: &Namespace,
+        with_scoped_ctx: impl FnOnce(TypeCheckContext) -> Result<T, ErrorEmitted>,
+    ) -> Result<T, ErrorEmitted> {
+        let mut namespace = self.namespace.clone();
+        let ctx = TypeCheckContext {
+            namespace: &mut namespace,
+            type_annotation: self.type_annotation,
+            function_type_annotation: self.function_type_annotation,
+            unify_generic: self.unify_generic,
+            self_type: self.self_type,
+            type_subst: self.type_subst,
+            abi_mode: self.abi_mode,
+            const_shadowing_mode: self.const_shadowing_mode,
+            generic_shadowing_mode: self.generic_shadowing_mode,
+            help_text: self.help_text,
+            kind: self.kind,
+            engines: self.engines,
+            disallow_functions: self.disallow_functions,
+            storage_declaration: self.storage_declaration,
+            previous_namespace: Some(previous_namespace),
+            experimental: self.experimental,
+            collecting_unifications: self.collecting_unifications,
+        };
+        with_scoped_ctx(ctx)
+    }
+
     /// Enter the submodule with the given name and produce a type-check context ready for
     /// type-checking its content.
     ///
@@ -348,6 +393,13 @@ impl<'a> TypeCheckContext<'a> {
         Self { self_type, ..self }
     }
 
+    pub(crate) fn with_collecting_unifications(self) -> Self {
+        Self {
+            collecting_unifications: true,
+            ..self
+        }
+    }
+
     /// Map this `TypeCheckContext` instance to a new one with
     /// `disallow_functions` set to `true`.
     pub(crate) fn disallow_functions(self) -> Self {
@@ -419,6 +471,14 @@ impl<'a> TypeCheckContext<'a> {
         self.storage_declaration
     }
 
+    pub(crate) fn previous_namespace(&self) -> Option<&Namespace> {
+        self.previous_namespace
+    }
+
+    pub(crate) fn collecting_unifications(&self) -> bool {
+        self.collecting_unifications
+    }
+
     // Provide some convenience functions around the inner context.
 
     /// Short-hand for calling the `monomorphize` function in the type engine
@@ -1111,6 +1171,14 @@ impl<'a> TypeCheckContext<'a> {
 
         // default numeric types to u64
         if type_engine.contains_numeric(decl_engine, type_id) {
+            // While collecting unification we don't decay numeric and will ignore this error.
+            if self.collecting_unifications {
+                return Err(handler.emit_err(CompileError::MethodNotFound {
+                    method_name: method_name.clone(),
+                    type_name: self.engines.help_out(type_id).to_string(),
+                    span: method_name.span(),
+                }));
+            }
             type_engine.decay_numeric(handler, self.engines, type_id, &method_name.span())?;
         }
 

sway-core/src/type_system/ast_elements/type_parameter.rs

@@ -352,6 +352,7 @@ impl TypeParameter {
 
         // Trait constraints mutate so we replace the previous type id associated TypeInfo.
         type_engine.replace(
+            ctx.engines(),
             type_parameter.type_id,
             TypeSourceInfo {
                 type_info: TypeInfo::UnknownGeneric {
@@ -441,6 +442,7 @@ impl TypeParameter {
                         }
 
                         ctx.engines.te().replace(
+                            ctx.engines(),
                             *type_id,
                             TypeSourceInfo {
                                 type_info: TypeInfo::UnknownGeneric {