Add tokens for storage reassignment (FuelLabs#4261)

## Description

Produce proper definition tokens for assignment to storage and
underlying types.

Fix FuelLabs#4166

## 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: Mohammad Fawaz <[email protected]>

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

@@ -8,6 +8,7 @@ pub struct StorageDeclaration {
     pub attributes: transform::AttributesMap,
     pub fields: Vec<StorageField>,
     pub span: Span,
+    pub storage_keyword: Ident,
 }
 
 /// An individual field in a storage declaration.

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

@@ -223,7 +223,7 @@ pub enum ExpressionKind {
 #[derive(Debug, Clone)]
 pub enum ReassignmentTarget {
     VariableExpression(Box<Expression>),
-    StorageField(Vec<Ident>),
+    StorageField(Span, Vec<Ident>),
 }
 
 #[derive(Debug, Clone)]

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

@@ -13,12 +13,12 @@ pub struct TyStorageDeclaration {
     pub fields: Vec<TyStorageField>,
     pub span: Span,
     pub attributes: transform::AttributesMap,
-    name: Ident,
+    pub storage_keyword: Ident,
 }
 
 impl Named for TyStorageDeclaration {
     fn name(&self) -> &Ident {
-        &self.name
+        &self.storage_keyword
     }
 }
 
@@ -37,7 +37,7 @@ impl HashWithEngines for TyStorageDeclaration {
             // reliable source of obj v. obj distinction
             span: _,
             attributes: _,
-            name: _,
+            storage_keyword: _,
         } = self;
         fields.hash(state, engines);
     }
@@ -50,19 +50,6 @@ impl Spanned for TyStorageDeclaration {
 }
 
 impl TyStorageDeclaration {
-    pub fn new(
-        fields: Vec<TyStorageField>,
-        span: Span,
-        attributes: transform::AttributesMap,
-    ) -> Self {
-        TyStorageDeclaration {
-            name: Ident::new_with_override("storage".to_string(), span.clone()),
-            fields,
-            span,
-            attributes,
-        }
-    }
-
     /// Given a field, find its type information in the declaration and return it. If the field has not
     /// been declared as a part of storage, return an error.
     pub fn apply_storage_load(

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

@@ -164,6 +164,7 @@ pub struct TyStorageReassignment {
     pub fields: Vec<TyStorageReassignDescriptor>,
     pub(crate) ix: StateIndex,
     pub rhs: TyExpression,
+    pub storage_keyword_span: Span,
 }
 
 impl EqWithEngines for TyStorageReassignment {}
@@ -177,7 +178,14 @@ impl PartialEqWithEngines for TyStorageReassignment {
 
 impl HashWithEngines for TyStorageReassignment {
     fn hash<H: Hasher>(&self, state: &mut H, engines: Engines<'_>) {
-        let TyStorageReassignment { fields, ix, rhs } = self;
+        let TyStorageReassignment {
+            fields,
+            ix,
+            rhs,
+            // these fields are not hashed because they aren't relevant/a
+            // reliable source of obj v. obj distinction
+            storage_keyword_span: _,
+        } = self;
         fields.hash(state, engines);
         ix.hash(state);
         rhs.hash(state, engines);

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

@@ -411,7 +411,7 @@ impl ty::TyDeclaration {
                 span,
                 fields,
                 attributes,
-                ..
+                storage_keyword,
             }) => {
                 let mut fields_buf = Vec::with_capacity(fields.len());
                 for parsed::StorageField {
@@ -446,7 +446,12 @@ impl ty::TyDeclaration {
                         attributes,
                     });
                 }
-                let decl = ty::TyStorageDeclaration::new(fields_buf, span, attributes);
+                let decl = ty::TyStorageDeclaration {
+                    fields: fields_buf,
+                    span,
+                    attributes,
+                    storage_keyword,
+                };
                 let decl_ref = decl_engine.insert(decl);
                 // insert the storage declaration into the symbols
                 // if there already was one, return an error that duplicate storage

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

@@ -1733,12 +1733,12 @@ impl ty::TyExpression {
                     errors,
                 )
             }
-            ReassignmentTarget::StorageField(fields) => {
+            ReassignmentTarget::StorageField(storage_keyword_span, fields) => {
                 let ctx = ctx
                     .with_type_annotation(type_engine.insert(decl_engine, TypeInfo::Unknown))
                     .with_help_text("");
                 let reassignment = check!(
-                    reassign_storage_subfield(ctx, fields, rhs, span.clone()),
+                    reassign_storage_subfield(ctx, fields, rhs, span.clone(), storage_keyword_span),
                     return err(warnings, errors),
                     warnings,
                     errors,

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

@@ -120,6 +120,7 @@ pub(crate) fn reassign_storage_subfield(
     fields: Vec<Ident>,
     rhs: Expression,
     span: Span,
+    storage_keyword_span: Span,
 ) -> CompileResult<ty::TyStorageReassignment> {
     let mut errors = vec![];
     let mut warnings = vec![];
@@ -219,6 +220,7 @@ pub(crate) fn reassign_storage_subfield(
 
     ok(
         ty::TyStorageReassignment {
+            storage_keyword_span,
             fields: type_checked_buf,
             ix,
             rhs,

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

@@ -1,5 +1,10 @@
 use crate::{
-    decl_engine::*, engine_threading::Engines, error::*, language::ty, namespace::*, type_system::*,
+    decl_engine::*,
+    engine_threading::Engines,
+    error::*,
+    language::ty::{self, TyStorageDeclaration},
+    namespace::*,
+    type_system::*,
 };
 
 use super::TraitMap;
@@ -169,17 +174,20 @@ impl Items {
         self.declared_storage.is_some()
     }
 
+    pub fn get_declared_storage(&self, decl_engine: &DeclEngine) -> Option<TyStorageDeclaration> {
+        self.declared_storage
+            .as_ref()
+            .map(|decl_ref| decl_engine.get_storage(decl_ref))
+    }
+
     pub(crate) fn get_storage_field_descriptors(
         &self,
         decl_engine: &DeclEngine,
     ) -> CompileResult<Vec<ty::TyStorageField>> {
         let warnings = vec![];
         let mut errors = vec![];
-        match self.declared_storage {
-            Some(ref decl_ref) => {
-                let storage = decl_engine.get_storage(decl_ref);
-                ok(storage.fields, warnings, errors)
-            }
+        match self.get_declared_storage(decl_engine) {
+            Some(storage) => ok(storage.fields, warnings, errors),
             None => {
                 let msg = "unknown source location";
                 let span = Span::new(Arc::from(msg), 0, msg.len(), None).unwrap();

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

@@ -830,6 +830,7 @@ fn item_storage_to_storage_declaration(
         attributes,
         span,
         fields,
+        storage_keyword: item_storage.storage_token.into(),
     };
     Ok(storage_declaration)
 }
@@ -3529,7 +3530,7 @@ fn assignable_to_reassignment_target(
             Assignable::Var(name) => {
                 if name.as_str() == "storage" {
                     let idents = idents.into_iter().rev().cloned().collect();
-                    return Ok(ReassignmentTarget::StorageField(idents));
+                    return Ok(ReassignmentTarget::StorageField(name.span(), idents));
                 }
                 break;
             }

sway-lsp/src/capabilities/document_symbol.rs

@@ -24,7 +24,7 @@ pub(crate) fn symbol_kind(symbol_kind: &SymbolKind) -> lsp_types::SymbolKind {
         SymbolKind::Function | SymbolKind::DeriveHelper => lsp_types::SymbolKind::FUNCTION,
         SymbolKind::Const => lsp_types::SymbolKind::CONSTANT,
         SymbolKind::Struct => lsp_types::SymbolKind::STRUCT,
-        SymbolKind::Trait => lsp_types::SymbolKind::INTERFACE,
+        SymbolKind::Trait | SymbolKind::Storage => lsp_types::SymbolKind::INTERFACE,
         SymbolKind::Module => lsp_types::SymbolKind::MODULE,
         SymbolKind::Enum => lsp_types::SymbolKind::ENUM,
         SymbolKind::Variant => lsp_types::SymbolKind::ENUM_MEMBER,

sway-lsp/src/capabilities/semantic_tokens.rs

@@ -156,7 +156,7 @@ fn semantic_token_type(kind: &SymbolKind) -> SemanticTokenType {
         SymbolKind::Struct => SemanticTokenType::STRUCT,
         SymbolKind::Enum => SemanticTokenType::ENUM,
         SymbolKind::Variant => SemanticTokenType::ENUM_MEMBER,
-        SymbolKind::Trait => SemanticTokenType::INTERFACE,
+        SymbolKind::Trait | SymbolKind::Storage => SemanticTokenType::INTERFACE,
         SymbolKind::TypeParameter => SemanticTokenType::TYPE_PARAMETER,
         SymbolKind::Module => SemanticTokenType::NAMESPACE,
         SymbolKind::StringLiteral => SemanticTokenType::STRING,

sway-lsp/src/core/token.rs

@@ -68,6 +68,7 @@ pub enum TypedAstToken {
     TypedTraitFn(ty::TyTraitFn),
     TypedSupertrait(Supertrait),
     TypedStorageField(ty::TyStorageField),
+    TyStorageResassignment(Box<ty::TyStorageReassignment>),
     TyStorageAccessDescriptor(ty::TyStorageAccessDescriptor),
     TypeCheckedStorageReassignDescriptor(ty::TyStorageReassignDescriptor),
     TypedReassignment(ty::TyReassignment),
@@ -88,6 +89,7 @@ pub enum SymbolKind {
     Const,
     Struct,
     Trait,
+    Storage,
     Enum,
     Variant,
     BoolLiteral,

sway-lsp/src/traverse/parsed_tree.rs

@@ -320,7 +320,13 @@ impl Parse for ReassignmentExpression {
             ReassignmentTarget::VariableExpression(exp) => {
                 exp.parse(ctx);
             }
-            ReassignmentTarget::StorageField(idents) => {
+            ReassignmentTarget::StorageField(storage_keyword_span, idents) => {
+                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::Storage),
+                );
+
                 for ident in idents {
                     ctx.tokens.insert(
                         to_ident_key(ident),

sway-lsp/src/traverse/typed_tree.rs

@@ -798,18 +798,91 @@ impl<'a> TypedTree<'a> {
                 }
             }
             ty::TyExpressionVariant::StorageReassignment(storage_reassignment) => {
-                for field in &storage_reassignment.fields {
+                let type_engine = self.ctx.engines.te();
+                let decl_engine = self.ctx.engines.de();
+
+                // collect storage keyword
+                if let Some(mut token) = self
+                    .ctx
+                    .tokens
+                    .try_get_mut(&to_ident_key(&Ident::new(
+                        storage_reassignment.storage_keyword_span.clone(),
+                    )))
+                    .try_unwrap()
+                {
+                    token.typed = Some(TypedAstToken::TyStorageResassignment(
+                        storage_reassignment.clone(),
+                    ));
+                    if let Some(storage) = self.namespace.get_declared_storage(decl_engine) {
+                        token.type_def = Some(TypeDefinition::Ident(storage.storage_keyword));
+                    }
+                }
+
+                if let Some((head_field, tail_fields)) = storage_reassignment.fields.split_first() {
+                    // collect the first ident as a field of the storage definition
                     if let Some(mut token) = self
                         .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::TypeCheckedStorageReassignDescriptor(
-                            field.clone(),
+                            head_field.clone(),
                         ));
+
+                        if let Some(storage_field) = self
+                            .namespace
+                            .get_declared_storage(decl_engine)
+                            .and_then(|storage| {
+                                // find the corresponding field in the storage declaration
+                                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_reassignment.fields.iter().map(|f| f.type_id))
+                    {
+                        if let Some(mut token) = self
+                            .ctx
+                            .tokens
+                            .try_get_mut(&to_ident_key(&field.name))
+                            .try_unwrap()
+                        {
+                            token.typed = Some(
+                                TypedAstToken::TypeCheckedStorageReassignDescriptor(field.clone()),
+                            );
+
+                            match type_engine.get(container_type_id) {
+                                TypeInfo::Struct(decl_ref) => {
+                                    if let Some(field_name) = decl_engine
+                                        .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));
+                                }
+                            }
+                        }
                     }
                 }
+
                 self.handle_expression(&storage_reassignment.rhs);
             }
             ty::TyExpressionVariant::Return(exp) => self.handle_expression(exp),

sway-lsp/tests/fixtures/tokens/storage/.gitignore

@@ -0,0 +1,2 @@
+out
+target

sway-lsp/tests/fixtures/tokens/storage/Forc.toml

@@ -0,0 +1,8 @@
+[project]
+authors = ["Fuel Labs <[email protected]>"]
+entry = "main.sw"
+license = "Apache-2.0"
+name = "storage"
+
+[dependencies]
+std = { path = "../../../../../sway-lib-std" }

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

@@ -0,0 +1,29 @@
+contract;
+
+struct Type1 {
+    x: u64,
+    y: bool,
+    z: Type2,
+}
+
+struct Type2 {
+    x: u64,
+}
+
+storage {
+    var1: Type1 = Type1 { x:0, y: false, z: Type2 { x:0 } },
+}
+
+abi StorageExample {
+    #[storage(write)]
+    fn store_something();
+}
+
+impl StorageExample  for Contract {
+    #[storage(write)]
+    fn store_something() {
+        storage.var1.x = 42;
+        storage.var1.y = true;
+        storage.var1.z.x = 1337;
+    }
+}