Return DeclRef from TypeBinding::type_check. (FuelLabs#4308)

## Description

The main goal of this PR is to refactor
`TypeBinding::type_check_with_ident` into a new generic trait
`TypeCheckTypeBinding<T>` with a function `type_check` that returns
`DeclRef<DeclId<T>>`. The fact that this is done with a generic trait is
an implementation detail and is a workaround to create easy generic
returns in Rust. The important bit is that this function now returns a
`DeclRef<_>` instead of a `TyDeclaration`.

In doing this I discovered a bug with the implementation of
`CollectTypesMetadata` on `TypeId`. The implementation did not account
for any generics that might appear in _nested_ positions of types (like
inside the types of enums, structs, etc). To fix this, this PR
introduces a new abstract method `extract_any` on `TypeInfo` with method
signature: `fn extract_any<F>(&self, .., filter_fn: &F) ->
HashSet<TypeId> where F: Fn(&TypeInfo) -> bool`. The `extract_any`
extracts any bugs that match the criteria for `CollectTypesMetadata`.

This PR also updates some of the previous methods on `TypeInfo` to use
the new `extract_any` method.

- Easy generic returns in Rust:
https://blog.jcoglan.com/2019/04/22/generic-returns-in-rust/

## 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/declaration/function.rs

@@ -140,7 +140,6 @@ impl UnconstrainedTypeParameters for TyFunctionDeclaration {
         type_parameter: &TypeParameter,
     ) -> bool {
         let type_engine = engines.te();
-        let decl_engine = engines.de();
         let mut all_types: HashSet<TypeId> = self
             .type_parameters
             .iter()
@@ -149,14 +148,14 @@ impl UnconstrainedTypeParameters for TyFunctionDeclaration {
         all_types.extend(self.parameters.iter().flat_map(|param| {
             let mut inner = type_engine
                 .get(param.type_argument.type_id)
-                .extract_inner_types(type_engine, decl_engine);
+                .extract_inner_types(engines);
             inner.insert(param.type_argument.type_id);
             inner
         }));
         all_types.extend(
             type_engine
                 .get(self.return_type.type_id)
-                .extract_inner_types(type_engine, decl_engine),
+                .extract_inner_types(engines),
         );
         all_types.insert(self.return_type.type_id);
         let type_parameter_info = type_engine.get(type_parameter.type_id);

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

@@ -147,8 +147,13 @@ impl CollectTypesMetadata for TyExpression {
             StructExpression {
                 fields,
                 instantiation_span,
+                struct_ref,
                 ..
             } => {
+                let struct_decl = decl_engine.get_struct(struct_ref);
+                for type_parameter in struct_decl.type_parameters {
+                    ctx.call_site_insert(type_parameter.type_id, instantiation_span.clone());
+                }
                 if let TypeInfo::Struct(decl_ref) = ctx.type_engine.get(self.return_type) {
                     let decl = decl_engine.get_struct(&decl_ref);
                     for type_parameter in decl.type_parameters {

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

@@ -189,19 +189,18 @@ impl TyProgram {
                     {
                         let storage_decl = decl_engine.get_storage(decl_id);
                         for field in storage_decl.fields.iter() {
-                            let type_info = ty_engine.get(field.type_argument.type_id);
-                            let type_info_str = engines.help_out(&type_info).to_string();
-                            let raw_ptr_type = type_info
-                                .extract_nested_types(engines, &field.span)
-                                .value
-                                .and_then(|value| {
-                                    value
-                                        .into_iter()
-                                        .find(|ty| matches!(ty, TypeInfo::RawUntypedPtr))
-                                });
-                            if raw_ptr_type.is_some() {
+                            if !field
+                                .type_argument
+                                .type_id
+                                .extract_any_including_self(engines, &|type_info| {
+                                    matches!(type_info, TypeInfo::RawUntypedPtr)
+                                })
+                                .is_empty()
+                            {
                                 errors.push(CompileError::TypeNotAllowedInContractStorage {
-                                    ty: type_info_str,
+                                    ty: engines
+                                        .help_out(&ty_engine.get(field.type_argument.type_id))
+                                        .to_string(),
                                     span: field.span.clone(),
                                 });
                             }
@@ -260,27 +259,24 @@ impl TyProgram {
                 // Directly returning a `raw_slice` is allowed, which will be just mapped to a RETD.
                 // TODO: Allow returning nested `raw_slice`s when our spec supports encoding DSTs.
                 let main_func = mains.remove(0);
-                let main_return_type_info = ty_engine.get(main_func.return_type.type_id);
-                let nested_types = check!(
-                    main_return_type_info
-                        .clone()
-                        .extract_nested_types(engines, &main_func.return_type.span),
-                    vec![],
-                    warnings,
-                    errors
-                );
-                if nested_types
-                    .iter()
-                    .any(|ty| matches!(ty, TypeInfo::RawUntypedPtr))
+                if !main_func
+                    .return_type
+                    .type_id
+                    .extract_any_including_self(engines, &|type_info| {
+                        matches!(type_info, TypeInfo::RawUntypedPtr)
+                    })
+                    .is_empty()
                 {
                     errors.push(CompileError::PointerReturnNotAllowedInMain {
                         span: main_func.return_type.span.clone(),
                     });
                 }
-                if !matches!(main_return_type_info, TypeInfo::RawUntypedSlice)
-                    && nested_types
-                        .iter()
-                        .any(|ty| matches!(ty, TypeInfo::RawUntypedSlice))
+                if !ty_engine
+                    .get(main_func.return_type.type_id)
+                    .extract_any(engines, &|type_info| {
+                        matches!(type_info, TypeInfo::RawUntypedSlice)
+                    })
+                    .is_empty()
                 {
                     errors.push(CompileError::NestedSliceReturnNotAllowedInMain {
                         span: main_func.return_type.span.clone(),

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

@@ -86,7 +86,6 @@ impl ty::TyImplTrait {
                 &new_impl_type_parameters,
                 &trait_type_arguments,
                 implementing_for.type_id,
-                &implementing_for.span
             ),
             return err(warnings, errors),
             warnings,
@@ -495,7 +494,6 @@ impl ty::TyImplTrait {
                 &new_impl_type_parameters,
                 &[],
                 implementing_for.type_id,
-                &implementing_for.span
             ),
             return err(warnings, errors),
             warnings,
@@ -1191,9 +1189,8 @@ fn check_for_unconstrained_type_parameters(
     type_parameters: &[TypeParameter],
     trait_type_arguments: &[TypeArgument],
     self_type: TypeId,
-    self_type_span: &Span,
 ) -> CompileResult<()> {
-    let mut warnings = vec![];
+    let warnings = vec![];
     let mut errors = vec![];
 
     // create a list of defined generics, with the generic and a span
@@ -1207,25 +1204,14 @@ fn check_for_unconstrained_type_parameters(
     // create a list of the generics in use in the impl signature
     let mut generics_in_use = HashSet::new();
     for type_arg in trait_type_arguments.iter() {
-        generics_in_use.extend(check!(
+        generics_in_use.extend(
             engines
                 .te()
                 .get(type_arg.type_id)
-                .extract_nested_generics(engines, &type_arg.span),
-            HashSet::new(),
-            warnings,
-            errors
-        ));
+                .extract_nested_generics(engines),
+        );
     }
-    generics_in_use.extend(check!(
-        engines
-            .te()
-            .get(self_type)
-            .extract_nested_generics(engines, self_type_span),
-        HashSet::new(),
-        warnings,
-        errors
-    ));
+    generics_in_use.extend(engines.te().get(self_type).extract_nested_generics(engines));
 
     // TODO: add a lookup in the trait constraints here and add it to
     // generics_in_use

sway-core/src/semantic_analysis/ast_node/expression/match_expression/analysis/constructor_factory.rs

@@ -18,20 +18,9 @@ pub(crate) struct ConstructorFactory {
 }
 
 impl ConstructorFactory {
-    pub(crate) fn new(engines: Engines<'_>, type_id: TypeId, span: &Span) -> CompileResult<Self> {
-        let mut warnings = vec![];
-        let mut errors = vec![];
-        let possible_types = check!(
-            engines
-                .te()
-                .get(type_id)
-                .extract_nested_types(engines, span),
-            return err(warnings, errors),
-            warnings,
-            errors
-        );
-        let factory = ConstructorFactory { possible_types };
-        ok(factory, warnings, errors)
+    pub(crate) fn new(engines: Engines<'_>, type_id: TypeId) -> Self {
+        let possible_types = engines.te().get(type_id).extract_nested_types(engines);
+        ConstructorFactory { possible_types }
     }
 
     /// Given Σ, computes a `Pattern` not present in Σ from the type of the

sway-core/src/semantic_analysis/ast_node/expression/match_expression/analysis/usefulness.rs

@@ -223,12 +223,7 @@ pub(crate) fn check_match_expression_usefulness(
         return ok((witness_report, arms_reachability), warnings, errors);
     }
 
-    let factory = check!(
-        ConstructorFactory::new(engines, type_id, &span),
-        return err(warnings, errors),
-        warnings,
-        errors
-    );
+    let factory = ConstructorFactory::new(engines, type_id);
     for scrutinee in scrutinees.into_iter() {
         let pat = check!(
             Pattern::from_scrutinee(scrutinee.clone()),

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

@@ -463,17 +463,9 @@ impl ty::TyExpression {
         let mut warnings = vec![];
         let mut errors = vec![];
 
-        // Grab the declaration.
-        let (unknown_decl, _type_id) = check!(
-            TypeBinding::type_check_with_ident(&mut call_path_binding, ctx.by_ref()),
-            return err(warnings, errors),
-            warnings,
-            errors
-        );
-
-        // Unwrap a fn ref if possible.
-        let fn_ref = check!(
-            unknown_decl.to_fn_ref(),
+        // Grab the fn declaration.
+        let (fn_ref, _): (DeclRefFunction, _) = check!(
+            TypeBinding::type_check(&mut call_path_binding, ctx.by_ref()),
             return err(warnings, errors),
             warnings,
             errors
@@ -1087,19 +1079,17 @@ impl ty::TyExpression {
         // Check if this could be a function
         let mut function_probe_warnings = Vec::new();
         let mut function_probe_errors = Vec::new();
-
-        let maybe_function = {
+        let maybe_function: Option<(DeclRefFunction, _)> = {
             let mut call_path_binding = unknown_call_path_binding.clone();
-            TypeBinding::type_check_with_ident(&mut call_path_binding, ctx.by_ref())
-                .flat_map(|(decl, _type_id)| decl.to_fn_ref())
+            TypeBinding::type_check(&mut call_path_binding, ctx.by_ref())
                 .ok(&mut function_probe_warnings, &mut function_probe_errors)
-                .map(|fn_ref| (fn_ref, call_path_binding))
+                .map(|(fn_ref, _)| (fn_ref, call_path_binding))
         };
 
         // Check if this could be an enum
         let mut enum_probe_warnings = vec![];
         let mut enum_probe_errors = vec![];
-        let maybe_enum = {
+        let maybe_enum: Option<(DeclRefEnum, _, _)> = {
             let call_path_binding = unknown_call_path_binding.clone();
             let variant_name = call_path_binding.inner.suffix.clone();
             let enum_call_path = call_path_binding.inner.rshift();
@@ -1109,10 +1099,9 @@ impl ty::TyExpression {
                 type_arguments: call_path_binding.type_arguments,
                 span: call_path_binding.span,
             };
-            TypeBinding::type_check_with_ident(&mut call_path_binding, ctx.by_ref())
-                .flat_map(|(unknown_decl, _type_id)| unknown_decl.to_enum_ref(ctx.engines()))
+            TypeBinding::type_check(&mut call_path_binding, ctx.by_ref())
                 .ok(&mut enum_probe_warnings, &mut enum_probe_errors)
-                .map(|enum_ref| (enum_ref, variant_name, call_path_binding))
+                .map(|(enum_ref, _)| (enum_ref, variant_name, call_path_binding))
         };
 
         // Check if this could be a constant
@@ -1224,10 +1213,10 @@ impl ty::TyExpression {
             call_path_binding.strip_prefixes();
         }
 
-        let const_opt = TypeBinding::type_check_with_ident(&mut call_path_binding, ctx.by_ref())
-            .flat_map(|(unknown_decl, _type_id)| unknown_decl.to_const_ref())
-            .ok(const_probe_warnings, const_probe_errors)
-            .map(|const_decl| (const_decl, call_path_binding.clone()));
+        let const_opt: Option<(DeclRefConstant, _)> =
+            TypeBinding::type_check(&mut call_path_binding, ctx.by_ref())
+                .ok(const_probe_warnings, const_probe_errors)
+                .map(|(const_ref, _)| (const_ref, call_path_binding.clone()));
         if const_opt.is_some() {
             return const_opt;
         }
@@ -1245,19 +1234,16 @@ impl ty::TyExpression {
             span: call_path_binding.span.clone(),
         };
 
-        let struct_type_id = match check!(
-            TypeBinding::type_check_with_ident(&mut const_call_path_binding, ctx.by_ref()),
+        let (_, struct_type_id): (DeclRefStruct, _) = check!(
+            TypeBinding::type_check(&mut const_call_path_binding, ctx.by_ref()),
             return None,
             const_probe_warnings,
             const_probe_errors
-        ) {
-            (ty::TyDeclaration::StructDeclaration { .. }, Some(type_id)) => type_id,
-            _ => return None,
-        };
+        );
 
         let const_decl_ref = check!(
             ctx.namespace.find_constant_for_type(
-                struct_type_id,
+                struct_type_id.unwrap(),
                 &suffix,
                 ctx.self_type(),
                 ctx.engines(),

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

@@ -2,6 +2,7 @@ use sway_error::error::CompileError;
 use sway_types::{Ident, Span, Spanned};
 
 use crate::{
+    decl_engine::DeclRefStruct,
     error::*,
     language::{parsed::*, ty, CallPath},
     semantic_analysis::TypeCheckContext,
@@ -24,7 +25,8 @@ pub(crate) fn struct_instantiation(
     // We need the call_path_binding to have types that point to proper definitions so the LSP can
     // look for them, but its types haven't been resolved yet.
     // To that end we do a dummy type check which has the side effect of resolving the types.
-    let _ = TypeBinding::type_check_with_ident(&mut call_path_binding, ctx.by_ref());
+    let _: CompileResult<(DeclRefStruct, _)> =
+        TypeBinding::type_check(&mut call_path_binding, ctx.by_ref());
 
     let TypeBinding {
         inner: CallPath {

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

@@ -477,14 +477,11 @@ impl TraitMap {
     /// with those entries for `Data<T, T>`.
     pub(crate) fn filter_by_type(&self, type_id: TypeId, engines: Engines<'_>) -> TraitMap {
         let type_engine = engines.te();
-        let decl_engine = engines.de();
         // a curried version of the decider protocol to use in the helper functions
         let decider = |type_info: &TypeInfo, map_type_info: &TypeInfo| {
             type_info.is_subset_of(map_type_info, engines)
         };
-        let mut all_types = type_engine
-            .get(type_id)
-            .extract_inner_types(type_engine, decl_engine);
+        let mut all_types = type_engine.get(type_id).extract_inner_types(engines);
         all_types.insert(type_id);
         let all_types = all_types.into_iter().collect::<Vec<_>>();
         self.filter_by_type_inner(engines, all_types, decider)
@@ -554,7 +551,6 @@ impl TraitMap {
         engines: Engines<'_>,
     ) -> TraitMap {
         let type_engine = engines.te();
-        let decl_engine = engines.de();
         // a curried version of the decider protocol to use in the helper functions
         let decider = |type_info: &TypeInfo, map_type_info: &TypeInfo| {
             type_info.is_subset_of(map_type_info, engines)
@@ -563,7 +559,7 @@ impl TraitMap {
         let mut trait_map = self.filter_by_type_inner(engines, vec![type_id], decider);
         let all_types = type_engine
             .get(type_id)
-            .extract_inner_types(type_engine, decl_engine)
+            .extract_inner_types(engines)
             .into_iter()
             .collect::<Vec<_>>();
         // a curried version of the decider protocol to use in the helper functions

sway-core/src/semantic_analysis/storage_only_types.rs

@@ -156,12 +156,7 @@ fn check_type(
         warnings,
         errors
     );
-    let nested_types = check!(
-        type_info.clone().extract_nested_types(engines, &span),
-        vec![],
-        warnings,
-        errors
-    );
+    let nested_types = type_info.clone().extract_nested_types(engines);
     for ty in nested_types {
         if ignore_self && ty.eq(&type_info, engines) {
             continue;