Simplify stored types during token collection (FuelLabs#4202)

## Description
This PR simplifies the token types that are collected during traversal.
We were being a bit verbose previously which was making it hard to
reason about the context. I had previously started to implement this in
FuelLabs#3722 but decided it would be nicer to have this be a separate PR, then
I can focus on the new Parse trait in a follow-up PR.

Also fixes a bug where inlay hints weren't showing and variables were
having incorrect semantic definitions applied.

works towards FuelLabs#3799

sway-lsp/src/capabilities/rename.rs

@@ -1,7 +1,4 @@
-use crate::core::{
-    session::Session,
-    token::{get_range_from_span, AstToken},
-};
+use crate::core::{session::Session, token::get_range_from_span};
 use std::collections::HashMap;
 use std::sync::Arc;
 use sway_types::Spanned;
@@ -37,12 +34,9 @@ pub fn prepare_rename(
     url: Url,
     position: Position,
 ) -> Option<PrepareRenameResponse> {
-    let (ident, token) = session.token_map().token_at_position(&url, position)?;
-    match token.parsed {
-        AstToken::Reassignment(_) => None,
-        _ => Some(PrepareRenameResponse::RangeWithPlaceholder {
-            range: get_range_from_span(&ident.span()),
-            placeholder: ident.as_str().to_string(),
-        }),
-    }
+    let (ident, ..) = session.token_map().token_at_position(&url, position)?;
+    Some(PrepareRenameResponse::RangeWithPlaceholder {
+        range: get_range_from_span(&ident.span()),
+        placeholder: ident.as_str().to_string(),
+    })
 }

sway-lsp/src/core/token.rs

@@ -3,8 +3,8 @@ use sway_core::{
     language::{
         parsed::{
             ConstantDeclaration, Declaration, EnumVariant, Expression, FunctionDeclaration,
-            FunctionParameter, ReassignmentExpression, Scrutinee, StorageField,
-            StructExpressionField, StructField, Supertrait, TraitFn, TreeType, UseStatement,
+            FunctionParameter, Scrutinee, StorageField, StructExpressionField, StructField,
+            StructScrutineeField, Supertrait, TraitFn, UseStatement,
         },
         ty,
     },
@@ -24,21 +24,26 @@ pub enum AstToken {
     Declaration(Declaration),
     Expression(Expression),
     StructExpressionField(StructExpressionField),
+    StructScrutineeField(StructScrutineeField),
     FunctionDeclaration(FunctionDeclaration),
     FunctionParameter(FunctionParameter),
     StructField(StructField),
     EnumVariant(EnumVariant),
     TraitFn(TraitFn),
+    TraitConstraint(TraitConstraint),
     ConstantDeclaration(ConstantDeclaration),
-    Reassignment(ReassignmentExpression),
     StorageField(StorageField),
     Scrutinee(Scrutinee),
     Keyword(Ident),
     Intrinsic(Intrinsic),
     Attribute(Attribute),
-    TreeType(TreeType),
+    LibraryName(Ident),
     IncludeStatement,
     UseStatement(UseStatement),
+    TypeArgument(TypeArgument),
+    TypeParameter(TypeParameter),
+    Supertrait(Supertrait),
+    Ident(Ident),
 }
 
 /// The `TypedAstToken` holds the types produced by the [sway_core::language::ty::TyProgram].
@@ -55,6 +60,7 @@ pub enum TypedAstToken {
     TypedTraitFn(ty::TyTraitFn),
     TypedSupertrait(Supertrait),
     TypedStorageField(ty::TyStorageField),
+    TyStorageAccessDescriptor(ty::TyStorageAccessDescriptor),
     TypeCheckedStorageReassignDescriptor(ty::TyStorageReassignDescriptor),
     TypedReassignment(ty::TyReassignment),
     TypedArgument(TypeArgument),
@@ -64,6 +70,7 @@ pub enum TypedAstToken {
     TypedLibraryName(Ident),
     TypedIncludeStatement,
     TypedUseStatement(ty::TyUseStatement),
+    Ident(Ident),
 }
 
 /// These variants are used to represent the semantic type of the [Token].

sway-lsp/src/traverse/lexed_tree.rs

@@ -10,8 +10,9 @@ use sway_ast::{
     ty::TyTupleDescriptor, Assignable, CodeBlockContents, ConfigurableField, Expr,
     ExprArrayDescriptor, ExprStructField, ExprTupleDescriptor, FnArg, FnArgs, FnSignature,
     IfCondition, IfExpr, ItemAbi, ItemConfigurable, ItemConst, ItemEnum, ItemFn, ItemImpl,
-    ItemImplItem, ItemKind, ItemStorage, ItemStruct, ItemTrait, ItemUse, MatchBranchKind, Pattern,
-    PatternStructField, Statement, StatementLet, StorageField, Ty, TypeField, UseTree,
+    ItemImplItem, ItemKind, ItemStorage, ItemStruct, ItemTrait, ItemUse, MatchBranchKind,
+    ModuleKind, Pattern, PatternStructField, Statement, StatementLet, StorageField, Ty, TypeField,
+    UseTree,
 };
 use sway_core::language::lexed::LexedProgram;
 use sway_types::{Ident, Span, Spanned};
@@ -26,20 +27,37 @@ impl<'a> LexedTree<'a> {
     }
 
     pub fn parse(&self, lexed_program: &LexedProgram) {
-        insert_keyword(self.tokens, lexed_program.root.tree.kind.span());
+        insert_module_kind(self.tokens, &lexed_program.root.tree.kind);
         for item in &lexed_program.root.tree.items {
             item.value.parse(self.tokens);
         }
 
         for (.., dep) in &lexed_program.root.submodules {
-            insert_keyword(self.tokens, dep.module.tree.kind.span());
+            insert_module_kind(self.tokens, &dep.module.tree.kind);
             for item in &dep.module.tree.items {
                 item.value.parse(self.tokens);
             }
         }
     }
 }
 
+fn insert_module_kind(tokens: &TokenMap, kind: &ModuleKind) {
+    match kind {
+        ModuleKind::Script { script_token } => {
+            insert_keyword(tokens, script_token.span());
+        }
+        ModuleKind::Contract { contract_token } => {
+            insert_keyword(tokens, contract_token.span());
+        }
+        ModuleKind::Predicate { predicate_token } => {
+            insert_keyword(tokens, predicate_token.span());
+        }
+        ModuleKind::Library { library_token, .. } => {
+            insert_keyword(tokens, library_token.span());
+        }
+    }
+}
+
 fn insert_keyword(tokens: &TokenMap, span: Span) {
     let ident = Ident::new(span);
     let token = Token::from_parsed(AstToken::Keyword(ident.clone()), SymbolKind::Keyword);