Introduce explicit type check unification pass (#5158)

## Description

This PR refactors code block and function unification into an explicit
pass that we can run as a separate step from type checking.

The main idea is that we want to run monormophization after type
checking, and unification needs to happen after monomorphization, so
this is an incremental refactor as a first step to further move us in
that direction in following PRs.

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [x] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [x] I have added tests that prove my fix is effective or that my
feature works.
- [x] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [x] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [x] I have requested a review from the relevant team or maintainers.

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

@@ -1,15 +1,26 @@
 use std::hash::Hasher;
 
 use sway_error::handler::{ErrorEmitted, Handler};
+use sway_types::Span;
 
 use crate::{
     decl_engine::*, engine_threading::*, language::ty::*, semantic_analysis::TypeCheckContext,
     type_system::*, types::DeterministicallyAborts,
 };
 
-#[derive(Clone, Debug, Default)]
+#[derive(Clone, Debug)]
 pub struct TyCodeBlock {
     pub contents: Vec<TyAstNode>,
+    pub(crate) whole_block_span: Span,
+}
+
+impl Default for TyCodeBlock {
+    fn default() -> Self {
+        Self {
+            contents: Default::default(),
+            whole_block_span: Span::dummy(),
+        }
+    }
 }
 
 impl EqWithEngines for TyCodeBlock {}
@@ -21,7 +32,7 @@ impl PartialEqWithEngines for TyCodeBlock {
 
 impl HashWithEngines for TyCodeBlock {
     fn hash<H: Hasher>(&self, state: &mut H, engines: &Engines) {
-        let TyCodeBlock { contents } = self;
+        let TyCodeBlock { contents, .. } = self;
         contents.hash(state, engines);
     }
 }

sway-core/src/language/ty/declaration/function.rs

@@ -231,9 +231,7 @@ impl TyFunctionDecl {
         TyFunctionDecl {
             purity,
             name,
-            body: TyCodeBlock {
-                contents: Default::default(),
-            },
+            body: TyCodeBlock::default(),
             implementing_type: None,
             span,
             call_path: CallPath::from(Ident::dummy()),

sway-core/src/semantic_analysis.rs

@@ -9,8 +9,10 @@ mod program;
 mod type_check_analysis;
 pub(crate) mod type_check_context;
 mod type_check_finalization;
+mod type_check_unification;
 pub use ast_node::*;
 pub use namespace::Namespace;
 pub(crate) use type_check_analysis::*;
 pub(crate) use type_check_context::TypeCheckContext;
 pub(crate) use type_check_finalization::*;
+pub(crate) use type_check_unification::*;

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

@@ -1,18 +1,18 @@
 use super::*;
 use crate::{
     decl_engine::DeclRef,
-    language::{parsed::CodeBlock, ty},
+    language::{
+        parsed::CodeBlock,
+        ty::{self, TyAstNodeContent, TyCodeBlock},
+    },
 };
 
 impl ty::TyCodeBlock {
     pub(crate) fn type_check(
         handler: &Handler,
         mut ctx: TypeCheckContext,
         code_block: &CodeBlock,
-    ) -> Result<(Self, TypeId), ErrorEmitted> {
-        let decl_engine = ctx.engines.de();
-        let engines = ctx.engines();
-
+    ) -> Result<Self, ErrorEmitted> {
         // Create a temp namespace for checking within the code block scope.
         let mut code_block_namespace = ctx.namespace.clone();
         let evaluated_contents = code_block
@@ -24,11 +24,24 @@ impl ty::TyCodeBlock {
             })
             .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(
+        ctx: &TypeCheckContext,
+        code_block: &TyCodeBlock,
+    ) -> (TypeId, Span) {
+        let engines = ctx.engines();
+        let decl_engine = engines.de();
+
         let implicit_return_span = code_block
             .contents
             .iter()
             .find_map(|x| match &x.content {
-                AstNodeContent::ImplicitReturnExpression(expr) => Some(Some(expr.span())),
+                TyAstNodeContent::ImplicitReturnExpression(expr) => Some(Some(expr.span.clone())),
                 _ => None,
             })
             .flatten();
@@ -38,7 +51,8 @@ impl ty::TyCodeBlock {
         // The fact that there is at most one implicit return is an invariant held by the parser.
         // If any node diverges then the entire block has unknown type.
         let mut node_deterministically_aborts = false;
-        let block_type = evaluated_contents
+        let block_type = code_block
+            .contents
             .iter()
             .find_map(|node| {
                 if node.deterministically_aborts(decl_engine, true) {
@@ -96,13 +110,7 @@ impl ty::TyCodeBlock {
                         .insert(engines, TypeInfo::Tuple(Vec::new()))
                 }
             });
-
-        ctx.unify_with_type_annotation(handler, block_type, &span);
-
-        let typed_code_block = ty::TyCodeBlock {
-            contents: evaluated_contents,
-        };
-        Ok((typed_code_block, block_type))
+        (block_type, span)
     }
 }
 
@@ -133,3 +141,23 @@ impl TypeCheckFinalization for ty::TyCodeBlock {
         })
     }
 }
+
+impl TypeCheckUnification for ty::TyCodeBlock {
+    fn type_check_unify(
+        &mut self,
+        handler: &Handler,
+        ctx: &mut TypeCheckUnificationContext,
+    ) -> Result<(), ErrorEmitted> {
+        handler.scope(|handler| {
+            let type_check_ctx = &ctx.type_check_ctx;
+            let (block_implicit_return, span) =
+                TyCodeBlock::compute_return_type_and_span(type_check_ctx, self);
+            let return_type_id = match ctx.type_id {
+                Some(type_id) => type_id,
+                None => block_implicit_return,
+            };
+            type_check_ctx.unify_with_type_annotation(handler, return_type_id, &span);
+            Ok(())
+        })
+    }
+}

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

@@ -165,9 +165,6 @@ impl ty::TyFunctionDecl {
             ..
         } = ty_fn_decl;
 
-        let type_engine = ctx.engines.te();
-        let engines = ctx.engines();
-
         // create a namespace for the function
         let mut fn_namespace = ctx.namespace.clone();
         let mut ctx = ctx
@@ -191,57 +188,24 @@ impl ty::TyFunctionDecl {
         //
         // If there are no implicit block returns, then we do not want to type check them, so we
         // stifle the errors. If there _are_ implicit block returns, we want to type_check them.
-        let (body, _implicit_block_return) = {
-            let ctx = ctx
-                .by_ref()
-                .with_purity(*purity)
-                .with_help_text("Function body's return type does not match up with its return type annotation.")
-                .with_type_annotation(return_type.type_id);
-            ty::TyCodeBlock::type_check(handler, ctx, body).unwrap_or_else(|err| {
-                (
-                    ty::TyCodeBlock { contents: vec![] },
-                    type_engine.insert(engines, TypeInfo::ErrorRecovery(err)),
-                )
-            })
-        };
-
-        ty_fn_decl.body = body;
-
-        Self::type_check_body_monomorphized(handler, ctx, ty_fn_decl)?;
 
-        Ok(ty_fn_decl.clone())
-    }
+        let mut ctx = ctx
+            .by_ref()
+            .with_purity(*purity)
+            .with_help_text(
+                "Function body's return type does not match up with its return type annotation.",
+            )
+            .with_type_annotation(return_type.type_id);
 
-    pub fn type_check_body_monomorphized(
-        handler: &Handler,
-        mut ctx: TypeCheckContext,
-        fn_decl: &Self,
-    ) -> Result<(), ErrorEmitted> {
-        let return_type = &fn_decl.return_type;
+        let body = ty::TyCodeBlock::type_check(handler, ctx.by_ref(), body)
+            .unwrap_or_else(|_err| ty::TyCodeBlock::default());
 
-        // gather the return statements
-        let return_statements: Vec<&ty::TyExpression> = fn_decl
-            .body
-            .contents
-            .iter()
-            .flat_map(|node| node.gather_return_statements())
-            .collect();
+        ty_fn_decl.body = body;
 
-        unify_return_statements(
-            handler,
-            ctx.by_ref(),
-            &return_statements,
-            return_type.type_id,
-        )?;
+        let mut unification_ctx = TypeCheckUnificationContext::new(ctx.engines, ctx);
+        ty_fn_decl.type_check_unify(handler, &mut unification_ctx)?;
 
-        return_type.type_id.check_type_parameter_bounds(
-            handler,
-            &ctx,
-            &return_type.span,
-            vec![],
-        )?;
-
-        Ok(())
+        Ok(ty_fn_decl.clone())
     }
 }
 
@@ -281,6 +245,46 @@ impl TypeCheckAnalysis for ty::TyFunctionDecl {
     }
 }
 
+impl TypeCheckUnification for ty::TyFunctionDecl {
+    fn type_check_unify(
+        &mut self,
+        handler: &Handler,
+        ctx: &mut TypeCheckUnificationContext,
+    ) -> Result<(), ErrorEmitted> {
+        handler.scope(|handler| {
+            self.body.type_check_unify(handler, ctx)?;
+
+            let type_check_ctx = &mut ctx.type_check_ctx;
+
+            let return_type = &self.return_type;
+
+            // gather the return statements
+            let return_statements: Vec<&ty::TyExpression> = self
+                .body
+                .contents
+                .iter()
+                .flat_map(|node| node.gather_return_statements())
+                .collect();
+
+            unify_return_statements(
+                handler,
+                type_check_ctx.by_ref(),
+                &return_statements,
+                return_type.type_id,
+            )?;
+
+            return_type.type_id.check_type_parameter_bounds(
+                handler,
+                type_check_ctx,
+                &return_type.span,
+                vec![],
+            )?;
+
+            Ok(())
+        })
+    }
+}
+
 impl TypeCheckFinalization for ty::TyFunctionDecl {
     fn type_check_finalize(
         &mut self,
@@ -306,7 +310,7 @@ fn test_function_selector_behavior() {
         purity: Default::default(),
         name: Ident::dummy(),
         implementing_type: None,
-        body: ty::TyCodeBlock { contents: vec![] },
+        body: ty::TyCodeBlock::default(),
         parameters: vec![],
         span: Span::dummy(),
         call_path: CallPath::from(Ident::dummy()),
@@ -329,7 +333,7 @@ fn test_function_selector_behavior() {
         purity: Default::default(),
         name: Ident::new_with_override("bar".into(), Span::dummy()),
         implementing_type: None,
-        body: ty::TyCodeBlock { contents: vec![] },
+        body: ty::TyCodeBlock::default(),
         parameters: vec![
             ty::TyFunctionParameter {
                 name: Ident::dummy(),

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

@@ -81,7 +81,7 @@ impl ty::TyTraitFn {
         ty::TyFunctionDecl {
             purity: self.purity,
             name: self.name.clone(),
-            body: ty::TyCodeBlock { contents: vec![] },
+            body: ty::TyCodeBlock::default(),
             parameters: self.parameters.clone(),
             implementing_type: match abi_mode.clone() {
                 AbiMode::ImplAbiFn(abi_name, abi_decl_id) => {

sway-core/src/semantic_analysis/ast_node/expression/match_expression/typed/instantiate.rs

@@ -105,6 +105,7 @@ impl Instantiate {
     pub(super) fn code_block_with_implicit_return_u64(&self, value: u64) -> ty::TyExpression {
         ty::TyExpression {
             expression: ty::TyExpressionVariant::CodeBlock(ty::TyCodeBlock {
+                whole_block_span: self.dummy_span(),
                 contents: vec![ty::TyAstNode {
                     content: ty::TyAstNodeContent::ImplicitReturnExpression(ty::TyExpression {
                         expression: ty::TyExpressionVariant::Literal(Literal::U64(value)),
@@ -127,6 +128,7 @@ impl Instantiate {
     ) -> ty::TyExpression {
         ty::TyExpression {
             expression: ty::TyExpressionVariant::CodeBlock(ty::TyCodeBlock {
+                whole_block_span: self.dummy_span(),
                 contents: vec![ty::TyAstNode {
                     content: ty::TyAstNodeContent::ImplicitReturnExpression(ty::TyExpression {
                         expression: ty::TyExpressionVariant::IntrinsicFunction(

sway-core/src/semantic_analysis/ast_node/expression/match_expression/typed/typed_match_branch.rs

@@ -150,6 +150,7 @@ impl ty::TyMatchBranch {
         let new_result = ty::TyExpression {
             expression: ty::TyExpressionVariant::CodeBlock(ty::TyCodeBlock {
                 contents: code_block_contents,
+                whole_block_span: sway_types::Span::dummy(),
             }),
             return_type: typed_result.return_type,
             span: typed_result_span,
@@ -650,6 +651,7 @@ fn instantiate_branch_condition_result_var_declarations_and_matched_or_variant_i
                         condition: Box::new(condition),
                         then: Box::new(ty::TyExpression {
                             expression: ty::TyExpressionVariant::CodeBlock(ty::TyCodeBlock {
+                                whole_block_span: instantiate.dummy_span(),
                                 contents: code_block_contents,
                             }),
                             return_type: tuple_type,

sway-core/src/semantic_analysis/ast_node/expression/match_expression/typed/typed_match_expression.rs

@@ -156,6 +156,7 @@ impl ty::TyMatchExpression {
 
                     Some(ty::TyExpression {
                         expression: ty::TyExpressionVariant::CodeBlock(ty::TyCodeBlock {
+                            whole_block_span: Span::dummy(),
                             contents: code_block_contents,
                         }),
                         return_type: self.return_type_id,