Add tokens for generic custom type fields (FuelLabs#4066)

## Description
Collect LSP tokens for fields of enums, structs and storage field
annotations by refactoring the `TypeId` handling to use `TypeArgument`
and collect name spans.

Fix FuelLabs#4052
Fix FuelLabs#4051 

## 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.

Co-authored-by: Joshua Batty <[email protected]>

forc-plugins/forc-doc/src/render.rs

@@ -583,7 +583,7 @@ impl Renderable for TyStructField {
                 code {
                     : format!("{}: ", self.name.as_str());
                     // TODO: Add links to types based on visibility
-                    : self.type_span.as_str();
+                    : self.type_argument.span.as_str();
                 }
             }
             @ if !self.attributes.is_empty() {
@@ -603,7 +603,7 @@ impl Renderable for TyStorageField {
                 code {
                     : format!("{}: ", self.name.as_str());
                     // TODO: Add links to types based on visibility
-                    : self.type_span.as_str();
+                    : self.type_argument.span.as_str();
                 }
             }
             @ if !self.attributes.is_empty() {
@@ -623,7 +623,7 @@ impl Renderable for TyEnumVariant {
                 a(class="anchor field", href=format!("{IDENTITY}{enum_variant_id}"));
                 code {
                     : format!("{}: ", self.name.as_str());
-                    : self.type_span.as_str();
+                    : self.type_argument.span.as_str();
                 }
             }
             @ if !self.attributes.is_empty() {

sway-core/src/abi_generation/evm_json_abi.rs

@@ -141,7 +141,9 @@ pub fn json_abi_param_type(type_info: &TypeInfo, type_engine: &TypeEngine) -> et
         Struct { fields, .. } => ethabi::ParamType::Tuple(
             fields
                 .iter()
-                .map(|f| json_abi_param_type(&type_engine.get(f.type_id), type_engine))
+                .map(|f| {
+                    json_abi_param_type(&type_engine.get(f.type_argument.type_id), type_engine)
+                })
                 .collect::<Vec<ethabi::ParamType>>(),
         ),
         Array(elem_ty, ..) => ethabi::ParamType::Array(Box::new(json_abi_param_type(

sway-core/src/abi_generation/fuel_json_abi.rs

@@ -265,17 +265,20 @@ impl TypeId {
                 let variants = variant_types
                     .iter()
                     .map(|x| program_abi::TypeDeclaration {
-                        type_id: x.initial_type_id.index(),
-                        type_field: x.initial_type_id.get_json_type_str(type_engine, x.type_id),
-                        components: x.initial_type_id.get_json_type_components(
+                        type_id: x.type_argument.initial_type_id.index(),
+                        type_field: x
+                            .type_argument
+                            .initial_type_id
+                            .get_json_type_str(type_engine, x.type_argument.type_id),
+                        components: x.type_argument.initial_type_id.get_json_type_components(
                             type_engine,
                             types,
-                            x.type_id,
+                            x.type_argument.type_id,
                         ),
-                        type_parameters: x.initial_type_id.get_json_type_parameters(
+                        type_parameters: x.type_argument.initial_type_id.get_json_type_parameters(
                             type_engine,
                             types,
-                            x.type_id,
+                            x.type_argument.type_id,
                         ),
                     })
                     .collect::<Vec<_>>();
@@ -288,12 +291,15 @@ impl TypeId {
                         .iter()
                         .map(|x| program_abi::TypeApplication {
                             name: x.name.to_string(),
-                            type_id: x.initial_type_id.index(),
-                            type_arguments: x.initial_type_id.get_json_type_arguments(
-                                type_engine,
-                                types,
-                                x.type_id,
-                            ),
+                            type_id: x.type_argument.initial_type_id.index(),
+                            type_arguments: x
+                                .type_argument
+                                .initial_type_id
+                                .get_json_type_arguments(
+                                    type_engine,
+                                    types,
+                                    x.type_argument.type_id,
+                                ),
                         })
                         .collect(),
                 )
@@ -303,17 +309,20 @@ impl TypeId {
                 let field_types = fields
                     .iter()
                     .map(|x| program_abi::TypeDeclaration {
-                        type_id: x.initial_type_id.index(),
-                        type_field: x.initial_type_id.get_json_type_str(type_engine, x.type_id),
-                        components: x.initial_type_id.get_json_type_components(
+                        type_id: x.type_argument.initial_type_id.index(),
+                        type_field: x
+                            .type_argument
+                            .initial_type_id
+                            .get_json_type_str(type_engine, x.type_argument.type_id),
+                        components: x.type_argument.initial_type_id.get_json_type_components(
                             type_engine,
                             types,
-                            x.type_id,
+                            x.type_argument.type_id,
                         ),
-                        type_parameters: x.initial_type_id.get_json_type_parameters(
+                        type_parameters: x.type_argument.initial_type_id.get_json_type_parameters(
                             type_engine,
                             types,
-                            x.type_id,
+                            x.type_argument.type_id,
                         ),
                     })
                     .collect::<Vec<_>>();
@@ -326,12 +335,15 @@ impl TypeId {
                         .iter()
                         .map(|x| program_abi::TypeApplication {
                             name: x.name.to_string(),
-                            type_id: x.initial_type_id.index(),
-                            type_arguments: x.initial_type_id.get_json_type_arguments(
-                                type_engine,
-                                types,
-                                x.type_id,
-                            ),
+                            type_id: x.type_argument.initial_type_id.index(),
+                            type_arguments: x
+                                .type_argument
+                                .initial_type_id
+                                .get_json_type_arguments(
+                                    type_engine,
+                                    types,
+                                    x.type_argument.type_id,
+                                ),
                         })
                         .collect(),
                 )

sway-core/src/control_flow_analysis/dead_code_analysis.rs

@@ -570,7 +570,7 @@ fn get_struct_type_info_from_type_id(
             }
             for var in variant_types.iter() {
                 if let Ok(Some(type_info)) =
-                    get_struct_type_info_from_type_id(type_engine, var.type_id)
+                    get_struct_type_info_from_type_id(type_engine, var.type_argument.type_id)
                 {
                     return Ok(Some(type_info));
                 }

sway-core/src/ir_generation/convert.rs

@@ -102,7 +102,7 @@ fn convert_resolved_type(
             context,
             fields
                 .iter()
-                .map(|field| field.type_id)
+                .map(|field| field.type_argument.type_id)
                 .collect::<Vec<_>>()
                 .as_slice(),
         )?,

sway-core/src/ir_generation/types.rs

@@ -19,7 +19,9 @@ pub(super) fn create_enum_aggregate(
     // getting one here anyway.  They don't need to be a tagged union either.
     let field_types: Vec<_> = variants
         .iter()
-        .map(|tev| convert_resolved_typeid_no_span(type_engine, context, &tev.type_id))
+        .map(|tev| {
+            convert_resolved_typeid_no_span(type_engine, context, &tev.type_argument.type_id)
+        })
         .collect::<Result<Vec<_>, CompileError>>()?;
 
     // Enums where all the variants are unit types don't really need the union. Only a tag is
@@ -89,7 +91,7 @@ pub(super) fn get_struct_name_field_index_and_type(
                 .iter()
                 .enumerate()
                 .find(|(_, field)| field.name == field_name)
-                .map(|(idx, field)| (idx as u64, field.type_id)),
+                .map(|(idx, field)| (idx as u64, field.type_argument.type_id)),
         )),
         _otherwise => None,
     }
@@ -161,7 +163,7 @@ pub(super) fn get_indices_for_struct_access(
                             .enumerate()
                             .find(|(_, field)| field.name == *field_name);
                         let (field_idx, field_type) = match field_idx_and_type_opt {
-                            Some((idx, field)) => (idx as u64, field.type_id),
+                            Some((idx, field)) => (idx as u64, field.type_argument.type_id),
                             None => {
                                 return Err(CompileError::InternalOwned(
                                     format!(

sway-core/src/language/parsed/declaration/enum.rs

@@ -15,8 +15,7 @@ pub struct EnumDeclaration {
 pub struct EnumVariant {
     pub name: Ident,
     pub attributes: transform::AttributesMap,
-    pub type_info: TypeInfo,
-    pub type_span: Span,
+    pub type_argument: TypeArgument,
     pub(crate) tag: usize,
     pub(crate) span: Span,
 }

sway-core/src/language/parsed/declaration/storage.rs

@@ -18,8 +18,7 @@ pub struct StorageDeclaration {
 pub struct StorageField {
     pub name: Ident,
     pub attributes: transform::AttributesMap,
-    pub type_info: TypeInfo,
-    pub type_info_span: Span,
+    pub type_argument: TypeArgument,
     pub span: Span,
     pub initializer: Expression,
 }

sway-core/src/language/parsed/declaration/struct.rs

@@ -1,8 +1,4 @@
-use crate::{
-    language::Visibility,
-    transform,
-    type_system::{TypeInfo, TypeParameter},
-};
+use crate::{language::Visibility, transform, type_system::TypeParameter, TypeArgument};
 use sway_types::{ident::Ident, span::Span};
 
 #[derive(Debug, Clone)]
@@ -19,7 +15,6 @@ pub struct StructDeclaration {
 pub struct StructField {
     pub name: Ident,
     pub attributes: transform::AttributesMap,
-    pub type_info: TypeInfo,
     pub(crate) span: Span,
-    pub type_span: Span,
+    pub type_argument: TypeArgument,
 }

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

@@ -131,43 +131,37 @@ impl TyEnumDeclaration {
 #[derive(Debug, Clone)]
 pub struct TyEnumVariant {
     pub name: Ident,
-    pub type_id: TypeId,
-    pub initial_type_id: TypeId,
-    pub type_span: Span,
+    pub type_argument: TypeArgument,
     pub(crate) tag: usize,
     pub span: Span,
     pub attributes: transform::AttributesMap,
 }
 
 impl HashWithEngines for TyEnumVariant {
     fn hash<H: Hasher>(&self, state: &mut H, engines: Engines<'_>) {
-        let type_engine = engines.te();
         self.name.hash(state);
-        type_engine.get(self.type_id).hash(state, engines);
+        self.type_argument.hash(state, engines);
         self.tag.hash(state);
     }
 }
 
 impl EqWithEngines for TyEnumVariant {}
 impl PartialEqWithEngines for TyEnumVariant {
     fn eq(&self, other: &Self, engines: Engines<'_>) -> bool {
-        let type_engine = engines.te();
         self.name == other.name
-            && type_engine
-                .get(self.type_id)
-                .eq(&type_engine.get(other.type_id), engines)
+            && self.type_argument.eq(&other.type_argument, engines)
             && self.tag == other.tag
     }
 }
 
 impl SubstTypes for TyEnumVariant {
     fn subst_inner(&mut self, type_mapping: &TypeSubstMap, engines: Engines<'_>) {
-        self.type_id.subst(type_mapping, engines);
+        self.type_argument.subst_inner(type_mapping, engines);
     }
 }
 
 impl ReplaceSelfType for TyEnumVariant {
     fn replace_self_type(&mut self, engines: Engines<'_>, self_type: TypeId) {
-        self.type_id.replace_self_type(engines, self_type);
+        self.type_argument.replace_self_type(engines, self_type);
     }
 }

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

@@ -71,12 +71,9 @@ impl TyStorageDeclaration {
             .enumerate()
             .find(|(_, TyStorageField { name, .. })| name == &first_field)
         {
-            Some((
-                ix,
-                TyStorageField {
-                    type_id: r#type, ..
-                },
-            )) => (StateIndex::new(ix), r#type),
+            Some((ix, TyStorageField { type_argument, .. })) => {
+                (StateIndex::new(ix), type_argument.type_id)
+            }
             None => {
                 errors.push(CompileError::StorageFieldDoesNotExist {
                     name: first_field.clone(),
@@ -88,7 +85,7 @@ impl TyStorageDeclaration {
 
         type_checked_buf.push(TyStorageAccessDescriptor {
             name: first_field.clone(),
-            type_id: *initial_field_type,
+            type_id: initial_field_type,
             span: first_field.span(),
         });
 
@@ -99,7 +96,7 @@ impl TyStorageDeclaration {
 
         // if the previously iterated type was a struct, put its fields here so we know that,
         // in the case of a subfield, we can type check the that the subfield exists and its type.
-        let mut available_struct_fields = update_available_struct_fields(*initial_field_type);
+        let mut available_struct_fields = update_available_struct_fields(initial_field_type);
 
         // get the initial field's type
         // make sure the next field exists in that type
@@ -111,10 +108,11 @@ impl TyStorageDeclaration {
                 Some(struct_field) => {
                     type_checked_buf.push(TyStorageAccessDescriptor {
                         name: field.clone(),
-                        type_id: struct_field.type_id,
+                        type_id: struct_field.type_argument.type_id,
                         span: field.span().clone(),
                     });
-                    available_struct_fields = update_available_struct_fields(struct_field.type_id);
+                    available_struct_fields =
+                        update_available_struct_fields(struct_field.type_argument.type_id);
                 }
                 None => {
                     let available_fields = available_struct_fields
@@ -153,17 +151,14 @@ impl TyStorageDeclaration {
             .map(
                 |TyStorageField {
                      ref name,
-                     type_id: ref r#type,
+                     ref type_argument,
                      ref span,
-                     ref initializer,
                      ref attributes,
                      ..
                  }| TyStructField {
                     name: name.clone(),
-                    type_id: *r#type,
-                    initial_type_id: *r#type,
                     span: span.clone(),
-                    type_span: initializer.span.clone(),
+                    type_argument: type_argument.clone(),
                     attributes: attributes.clone(),
                 },
             )
@@ -174,8 +169,7 @@ impl TyStorageDeclaration {
 #[derive(Clone, Debug)]
 pub struct TyStorageField {
     pub name: Ident,
-    pub type_id: TypeId,
-    pub type_span: Span,
+    pub type_argument: TypeArgument,
     pub initializer: TyExpression,
     pub(crate) span: Span,
     pub attributes: transform::AttributesMap,
@@ -184,11 +178,8 @@ pub struct TyStorageField {
 impl EqWithEngines for TyStorageField {}
 impl PartialEqWithEngines for TyStorageField {
     fn eq(&self, other: &Self, engines: Engines<'_>) -> bool {
-        let type_engine = engines.te();
         self.name == other.name
-            && type_engine
-                .get(self.type_id)
-                .eq(&type_engine.get(other.type_id), engines)
+            && self.type_argument.eq(&other.type_argument, engines)
             && self.initializer.eq(&other.initializer, engines)
     }
 }
@@ -197,17 +188,15 @@ impl HashWithEngines for TyStorageField {
     fn hash<H: Hasher>(&self, state: &mut H, engines: Engines<'_>) {
         let TyStorageField {
             name,
-            type_id,
+            type_argument,
             initializer,
             // these fields are not hashed because they aren't relevant/a
             // reliable source of obj v. obj distinction
             span: _,
             attributes: _,
-            type_span: _,
         } = self;
-        let type_engine = engines.te();
         name.hash(state);
-        type_engine.get(*type_id).hash(state, engines);
+        type_argument.hash(state, engines);
         initializer.hash(state, engines);
     }
 }