Collect storage keyword and assign type_def to storage fields (FuelLa…

…bs#4810)

## Description
The storage tests in the language server were removed in FuelLabs#4464. That PR
also removed the ability for the language server to access the `storage`
token and assign correct type definition to storage fields. This PR adds
all that functionality back in.

closes FuelLabs#4641
closes FuelLabs#4482

sway-core/src/language/parsed/expression/mod.rs

@@ -142,6 +142,7 @@ pub struct ArrayIndexExpression {
 #[derive(Debug, Clone)]
 pub struct StorageAccessExpression {
     pub field_names: Vec<Ident>,
+    pub storage_keyword_span: Span,
 }
 
 #[derive(Debug, Clone)]

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

@@ -58,6 +58,7 @@ impl TyStorageDecl {
         decl_engine: &DeclEngine,
         fields: Vec<Ident>,
         storage_fields: &[TyStorageField],
+        storage_keyword_span: Span,
     ) -> CompileResult<(TyStorageAccess, TypeId)> {
         let mut errors = vec![];
         let warnings = vec![];
@@ -137,6 +138,7 @@ impl TyStorageDecl {
                 TyStorageAccess {
                     fields: type_checked_buf,
                     ix,
+                    storage_keyword_span,
                 },
                 return_type,
             ),

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

@@ -9,6 +9,7 @@ use crate::{engine_threading::*, type_system::TypeId};
 pub struct TyStorageAccess {
     pub fields: Vec<TyStorageAccessDescriptor>,
     pub(crate) ix: StateIndex,
+    pub storage_keyword_span: Span,
 }
 
 impl EqWithEngines for TyStorageAccess {}
@@ -22,9 +23,14 @@ impl PartialEqWithEngines for TyStorageAccess {
 
 impl HashWithEngines for TyStorageAccess {
     fn hash<H: Hasher>(&self, state: &mut H, engines: &Engines) {
-        let TyStorageAccess { fields, ix } = self;
+        let TyStorageAccess {
+            fields,
+            ix,
+            storage_keyword_span,
+        } = self;
         fields.hash(state, engines);
         ix.hash(state);
+        storage_keyword_span.hash(state);
     }
 }
 
@@ -48,7 +54,7 @@ impl TyStorageAccess {
 #[derive(Clone, Debug)]
 pub struct TyStorageAccessDescriptor {
     pub name: Ident,
-    pub(crate) type_id: TypeId,
+    pub type_id: TypeId,
     pub(crate) span: Span,
 }
 

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

@@ -259,12 +259,15 @@ impl ty::TyExpression {
                     .with_help_text("");
                 Self::type_check_array_index(ctx, *prefix, *index, span)
             }
-            ExpressionKind::StorageAccess(StorageAccessExpression { field_names }) => {
+            ExpressionKind::StorageAccess(StorageAccessExpression {
+                field_names,
+                storage_keyword_span,
+            }) => {
                 let ctx = ctx
                     .by_ref()
                     .with_type_annotation(type_engine.insert(engines, TypeInfo::Unknown))
                     .with_help_text("");
-                Self::type_check_storage_access(ctx, field_names, &span)
+                Self::type_check_storage_access(ctx, field_names, storage_keyword_span, &span)
             }
             ExpressionKind::IntrinsicFunction(IntrinsicFunctionExpression {
                 kind_binding,
@@ -896,6 +899,7 @@ impl ty::TyExpression {
     fn type_check_storage_access(
         ctx: TypeCheckContext,
         checkee: Vec<Ident>,
+        storage_keyword_span: Span,
         span: &Span,
     ) -> CompileResult<Self> {
         let mut warnings = vec![];
@@ -919,8 +923,12 @@ impl ty::TyExpression {
 
         // Do all namespace checking here!
         let (storage_access, mut access_type) = check!(
-            ctx.namespace
-                .apply_storage_load(ctx.engines, checkee, &storage_fields,),
+            ctx.namespace.apply_storage_load(
+                ctx.engines,
+                checkee,
+                &storage_fields,
+                storage_keyword_span
+            ),
             return err(warnings, errors),
             warnings,
             errors

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

@@ -59,6 +59,7 @@ impl Items {
         engines: &Engines,
         fields: Vec<Ident>,
         storage_fields: &[ty::TyStorageField],
+        storage_keyword_span: Span,
     ) -> CompileResult<(ty::TyStorageAccess, TypeId)> {
         let warnings = vec![];
         let mut errors = vec![];
@@ -67,7 +68,13 @@ impl Items {
         match self.declared_storage {
             Some(ref decl_ref) => {
                 let storage = decl_engine.get_storage(&decl_ref.id().clone());
-                storage.apply_storage_load(type_engine, decl_engine, fields, storage_fields)
+                storage.apply_storage_load(
+                    type_engine,
+                    decl_engine,
+                    fields,
+                    storage_fields,
+                    storage_keyword_span,
+                )
             }
             None => {
                 errors.push(CompileError::NoDeclaredStorage {

sway-core/src/transform/to_parsed_lang/convert_parse_tree.rs

@@ -1988,7 +1988,7 @@ fn expr_to_expression(
             // For example, `storage.foo.bar` would result in `Some([foo, bar])`.
             let mut idents = vec![&name];
             let mut base = &*target;
-            let storage_access_field_names = loop {
+            let kind = loop {
                 match base {
                     // Parent is a projection itself, so check its parent.
                     Expr::FieldProjection { target, name, .. } => {
@@ -2002,22 +2002,22 @@ fn expr_to_expression(
                             && path_expr.prefix.generics_opt.is_none()
                             && path_expr.prefix.name.as_str() == "storage" =>
                     {
-                        break Some(idents)
+                        break ExpressionKind::StorageAccess(StorageAccessExpression {
+                            field_names: idents.into_iter().rev().cloned().collect(),
+                            storage_keyword_span: path_expr.prefix.name.span(),
+                        })
                     }
                     // We'll never find `storage`, so stop here.
-                    _ => break None,
+                    _ => {
+                        break ExpressionKind::Subfield(SubfieldExpression {
+                            prefix: Box::new(expr_to_expression(
+                                context, handler, engines, *target,
+                            )?),
+                            field_to_access: name,
+                        })
+                    }
                 }
             };
-
-            let kind = match storage_access_field_names {
-                Some(field_names) => ExpressionKind::StorageAccess(StorageAccessExpression {
-                    field_names: field_names.into_iter().rev().cloned().collect(),
-                }),
-                None => ExpressionKind::Subfield(SubfieldExpression {
-                    prefix: Box::new(expr_to_expression(context, handler, engines, *target)?),
-                    field_to_access: name,
-                }),
-            };
             Expression { kind, span }
         }
         Expr::TupleFieldProjection {
@@ -3768,32 +3768,33 @@ fn assignable_to_expression(
         Assignable::FieldProjection { target, name, .. } => {
             let mut idents = vec![&name];
             let mut base = &*target;
-            let storage_access_field_names_opt = loop {
+            let (storage_access_field_names_opt, storage_name_opt) = loop {
                 match base {
                     Assignable::FieldProjection { target, name, .. } => {
                         idents.push(name);
                         base = target;
                     }
                     Assignable::Var(name) => {
                         if name.as_str() == "storage" {
-                            break Some(idents);
+                            break (Some(idents), Some(name.clone()));
                         }
-                        break None;
+                        break (None, None);
                     }
-                    _ => break None,
+                    _ => break (None, None),
                 }
             };
-            match storage_access_field_names_opt {
-                Some(field_names) => {
+            match (storage_access_field_names_opt, storage_name_opt) {
+                (Some(field_names), Some(storage_name)) => {
                     let field_names = field_names.into_iter().rev().cloned().collect();
                     Expression {
                         kind: ExpressionKind::StorageAccess(StorageAccessExpression {
                             field_names,
+                            storage_keyword_span: storage_name.span(),
                         }),
                         span,
                     }
                 }
-                None => Expression {
+                _ => Expression {
                     kind: ExpressionKind::Subfield(SubfieldExpression {
                         prefix: Box::new(assignable_to_expression(
                             context, handler, engines, *target,
@@ -3840,6 +3841,7 @@ fn assignable_to_reassignment_target(
 ) -> Result<ReassignmentTarget, ErrorEmitted> {
     let mut idents = Vec::new();
     let mut base = &assignable;
+
     loop {
         match base {
             Assignable::FieldProjection { target, name, .. } => {

sway-lsp/src/core/token.rs

@@ -69,7 +69,8 @@ pub enum TypedAstToken {
     TypedTraitFn(ty::TyTraitFn),
     TypedSupertrait(Supertrait),
     TypedStorageField(ty::TyStorageField),
-    TyStorageAccessDescriptor(ty::TyStorageAccessDescriptor),
+    TypedStorageAccess(ty::TyStorageAccess),
+    TypedStorageAccessDescriptor(ty::TyStorageAccessDescriptor),
     TypedReassignment(ty::TyReassignment),
     TypedArgument(TypeArgument),
     TypedParameter(TypeParameter),

sway-lsp/src/traverse/parsed_tree.rs

@@ -292,7 +292,16 @@ impl Parse for Expression {
                 prefix.parse(ctx);
                 index.parse(ctx);
             }
-            ExpressionKind::StorageAccess(StorageAccessExpression { field_names }) => {
+            ExpressionKind::StorageAccess(StorageAccessExpression {
+                field_names,
+                storage_keyword_span,
+            }) => {
+                let storage_ident = Ident::new(storage_keyword_span.clone());
+                ctx.tokens.insert(
+                    to_ident_key(&storage_ident),
+                    Token::from_parsed(AstToken::Ident(storage_ident), SymbolKind::Unknown),
+                );
+
                 field_names.iter().for_each(|field_name| {
                     ctx.tokens.insert(
                         to_ident_key(field_name),

sway-lsp/src/traverse/typed_tree.rs

@@ -470,15 +470,77 @@ impl Parse for ty::TyExpression {
                 address.parse(ctx);
             }
             ty::TyExpressionVariant::StorageAccess(storage_access) => {
-                storage_access.fields.iter().for_each(|field| {
+                // collect storage keyword
+                if let Some(mut token) = ctx
+                    .tokens
+                    .try_get_mut(&to_ident_key(&Ident::new(
+                        storage_access.storage_keyword_span.clone(),
+                    )))
+                    .try_unwrap()
+                {
+                    token.typed = Some(TypedAstToken::TypedStorageAccess(storage_access.clone()));
+                    if let Some(storage) = ctx.namespace.get_declared_storage(ctx.engines.de()) {
+                        token.type_def = Some(TypeDefinition::Ident(storage.storage_keyword));
+                    }
+                }
+                if let Some((head_field, tail_fields)) = storage_access.fields.split_first() {
+                    // collect the first ident as a field of the storage definition
                     if let Some(mut token) = ctx
                         .tokens
-                        .try_get_mut(&to_ident_key(&field.name))
+                        .try_get_mut(&to_ident_key(&head_field.name))
                         .try_unwrap()
                     {
-                        token.typed = Some(TypedAstToken::TyStorageAccessDescriptor(field.clone()));
+                        token.typed = Some(TypedAstToken::TypedStorageAccessDescriptor(
+                            head_field.clone(),
+                        ));
+                        if let Some(storage_field) = ctx
+                            .namespace
+                            .get_declared_storage(ctx.engines.de())
+                            .and_then(|storage| {
+                                storage
+                                    .fields
+                                    .into_iter()
+                                    .find(|f| f.name.as_str() == head_field.name.as_str())
+                            })
+                        {
+                            token.type_def = Some(TypeDefinition::Ident(storage_field.name));
+                        }
                     }
-                });
+                    // collect the rest of the idents as fields of their respective types
+                    for (field, container_type_id) in tail_fields
+                        .iter()
+                        .zip(storage_access.fields.iter().map(|f| f.type_id))
+                    {
+                        if let Some(mut token) = ctx
+                            .tokens
+                            .try_get_mut(&to_ident_key(&field.name))
+                            .try_unwrap()
+                        {
+                            token.typed = Some(TypedAstToken::Ident(field.name.clone()));
+                            match ctx.engines.te().get(container_type_id) {
+                                TypeInfo::Struct(decl_ref) => {
+                                    if let Some(field_name) = ctx
+                                        .engines
+                                        .de()
+                                        .get_struct(&decl_ref)
+                                        .fields
+                                        .iter()
+                                        .find(|struct_field| {
+                                            // find the corresponding field in the containing type declaration
+                                            struct_field.name.as_str() == field.name.as_str()
+                                        })
+                                        .map(|struct_field| struct_field.name.clone())
+                                    {
+                                        token.type_def = Some(TypeDefinition::Ident(field_name));
+                                    }
+                                }
+                                _ => {
+                                    token.type_def = Some(TypeDefinition::TypeId(field.type_id));
+                                }
+                            }
+                        }
+                    }
+                }
             }
             ty::TyExpressionVariant::IntrinsicFunction(kind) => {
                 kind.parse(ctx);

sway-lsp/tests/fixtures/tokens/storage/src/main.sw

@@ -19,7 +19,7 @@ abi StorageExample {
     fn store_something();
 }
 
-impl StorageExample  for Contract {
+impl StorageExample for Contract {
     #[storage(write)]
     fn store_something() {
         storage.var1.x.write(42);