Store the variable decl Ident in VariableExpression and pass on the t…

…ype_ascription_span (FuelLabs#2544) This PR does 2 things: 1. Pass on the type_ascription_span : Option<Span> from VariableDeclaration to TypedVariableDeclaration. This enables the language server to collect any type ascriptions if they are provided. 2. Store the variable declaration Ident in VariableExpression along with the local Span. This now allows the language server to provide go to definition support for any plane a VariableExpression is used in the parse tree.
Kristal1523 · Aug 16, 2022 · 84badb6 · 84badb6
1 parent 58eeb6b
commit 84badb6
Show file tree

Hide file tree

Showing 10 changed files with 52 additions and 14 deletions.
diff --git a/sway-core/src/ir_generation/const_eval.rs b/sway-core/src/ir_generation/const_eval.rs
@@ -203,7 +203,7 @@ fn const_eval_typed_expr(
             }
             res
         }
-        TypedExpressionVariant::VariableExpression { name } => match known_consts.get(name) {
+        TypedExpressionVariant::VariableExpression { name, .. } => match known_consts.get(name) {
             // 1. Check if name is in known_consts.
             Some(cvs) => Some(cvs.clone()),
             None => {

diff --git a/sway-core/src/ir_generation/function.rs b/sway-core/src/ir_generation/function.rs
@@ -279,7 +279,7 @@ impl FnCompiler {
             TypedExpressionVariant::LazyOperator { op, lhs, rhs } => {
                 self.compile_lazy_op(context, md_mgr, op, *lhs, *rhs, span_md_idx)
             }
-            TypedExpressionVariant::VariableExpression { name } => {
+            TypedExpressionVariant::VariableExpression { name, .. } => {
                 self.compile_var_expr(context, name.as_str(), span_md_idx)
             }
             TypedExpressionVariant::Array { contents } => {

diff --git a/sway-core/src/semantic_analysis/ast_node/declaration/function/function_parameter.rs b/sway-core/src/semantic_analysis/ast_node/declaration/function/function_parameter.rs
@@ -78,6 +78,7 @@ impl TypedFunctionParameter {
                 },
                 mutability,
                 type_ascription: type_id,
+                type_ascription_span: None,
             }),
         );
         let parameter = TypedFunctionParameter {

diff --git a/sway-core/src/semantic_analysis/ast_node/declaration/variable.rs b/sway-core/src/semantic_analysis/ast_node/declaration/variable.rs
@@ -1,4 +1,5 @@
 use crate::{semantic_analysis::*, type_system::*, Ident, Visibility};
+use sway_types::Span;
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum VariableMutability {
@@ -55,6 +56,7 @@ pub struct TypedVariableDeclaration {
     pub body: TypedExpression,
     pub mutability: VariableMutability,
     pub type_ascription: TypeId,
+    pub type_ascription_span: Option<Span>,
 }
 
 // NOTE: Hash and PartialEq must uphold the invariant:

diff --git a/...re/src/semantic_analysis/ast_node/expression/match_expression/typed/typed_match_branch.rs b/...re/src/semantic_analysis/ast_node/expression/match_expression/typed/typed_match_branch.rs
@@ -68,6 +68,7 @@ impl TypedMatchBranch {
                 body: right_decl,
                 mutability: VariableMutability::Immutable,
                 type_ascription,
+                type_ascription_span: None,
             });
             ctx.namespace.insert_symbol(left_decl, var_decl.clone());
             code_block_contents.push(TypedAstNode {

diff --git a/sway-core/src/semantic_analysis/ast_node/expression/typed_expression.rs b/sway-core/src/semantic_analysis/ast_node/expression/typed_expression.rs
@@ -613,21 +613,31 @@ impl TypedExpression {
         let mut errors = vec![];
         let exp = match namespace.resolve_symbol(&name).value {
             Some(TypedDeclaration::VariableDeclaration(TypedVariableDeclaration {
-                body, ..
+                name: decl_name,
+                body,
+                ..
             })) => TypedExpression {
                 return_type: body.return_type,
                 is_constant: body.is_constant,
-                expression: TypedExpressionVariant::VariableExpression { name: name.clone() },
+                expression: TypedExpressionVariant::VariableExpression {
+                    name: decl_name.clone(),
+                    span: name.span(),
+                },
                 span,
             },
             Some(TypedDeclaration::ConstantDeclaration(TypedConstantDeclaration {
-                value, ..
+                name: decl_name,
+                value,
+                ..
             })) => TypedExpression {
                 return_type: value.return_type,
                 is_constant: IsConstant::Yes,
                 // Although this isn't strictly a 'variable' expression we can treat it as one for
                 // this context.
-                expression: TypedExpressionVariant::VariableExpression { name: name.clone() },
+                expression: TypedExpressionVariant::VariableExpression {
+                    name: decl_name.clone(),
+                    span: name.span(),
+                },
                 span,
             },
             Some(TypedDeclaration::AbiDeclaration(decl)) => TypedExpression {

diff --git a/sway-core/src/semantic_analysis/ast_node/expression/typed_expression_variant.rs b/sway-core/src/semantic_analysis/ast_node/expression/typed_expression_variant.rs
@@ -35,6 +35,7 @@ pub enum TypedExpressionVariant {
     },
     VariableExpression {
         name: Ident,
+        span: Span,
     },
     Tuple {
         fields: Vec<TypedExpression>,
@@ -156,9 +157,15 @@ impl PartialEq for TypedExpressionVariant {
                 },
             ) => l_op == r_op && (**l_lhs) == (**r_lhs) && (**l_rhs) == (**r_rhs),
             (
-                Self::VariableExpression { name: l_name },
-                Self::VariableExpression { name: r_name },
-            ) => l_name == r_name,
+                Self::VariableExpression {
+                    name: l_name,
+                    span: l_span,
+                },
+                Self::VariableExpression {
+                    name: r_name,
+                    span: r_span,
+                },
+            ) => l_name == r_name && l_span == r_span,
             (Self::Tuple { fields: l_fields }, Self::Tuple { fields: r_fields }) => {
                 l_fields == r_fields
             }

diff --git a/sway-core/src/semantic_analysis/ast_node/mod.rs b/sway-core/src/semantic_analysis/ast_node/mod.rs
@@ -226,12 +226,10 @@ impl TypedAstNode {
                             body,
                             is_mutable,
                         }) => {
-                            let type_ascription_span =
-                                type_ascription_span.unwrap_or_else(|| name.span());
                             let type_ascription = check!(
                                 ctx.resolve_type_with_self(
                                     insert_type(type_ascription),
-                                    &type_ascription_span,
+                                    &type_ascription_span.clone().unwrap_or_else(|| name.span()),
                                     EnforceTypeArguments::Yes,
                                     None
                                 ),
@@ -252,6 +250,7 @@ impl TypedAstNode {
                                     body,
                                     mutability: convert_to_variable_immutability(false, is_mutable),
                                     type_ascription,
+                                    type_ascription_span,
                                 });
                             ctx.namespace.insert_symbol(name, typed_var_decl.clone());
                             typed_var_decl
@@ -725,6 +724,7 @@ fn type_check_trait_methods(
                         },
                         mutability: convert_to_variable_immutability(is_reference, is_mutable),
                         type_ascription: r#type,
+                        type_ascription_span: None,
                     }),
                 );
             },

diff --git a/sway-lsp/src/core/traverse_parse_tree.rs b/sway-lsp/src/core/traverse_parse_tree.rs
@@ -76,6 +76,13 @@ fn handle_declaration(declaration: &Declaration, tokens: &TokenMap) {
                     to_ident_key(&variable.name),
                     Token::from_parsed(AstToken::Declaration(declaration.clone())),
                 );
+
+                if let Some(type_ascription_span) = &variable.type_ascription_span {
+                    tokens.insert(
+                        to_ident_key(&Ident::new(type_ascription_span.clone())),
+                        Token::from_parsed(AstToken::Declaration(declaration.clone())),
+                    );
+                }
             }
             handle_expression(&variable.body, tokens);
         }

diff --git a/sway-lsp/src/core/traverse_typed_tree.rs b/sway-lsp/src/core/traverse_typed_tree.rs
@@ -34,6 +34,15 @@ fn handle_declaration(declaration: &TypedDeclaration, tokens: &TokenMap) {
             if let Some(mut token) = tokens.get_mut(&to_ident_key(&variable.name)) {
                 token.typed = Some(TypedAstToken::TypedDeclaration(declaration.clone()));
             }
+            if let Some(type_ascription_span) = &variable.type_ascription_span {
+                if let Some(mut token) =
+                    tokens.get_mut(&to_ident_key(&Ident::new(type_ascription_span.clone())))
+                {
+                    token.typed = Some(TypedAstToken::TypedDeclaration(declaration.clone()));
+                    token.type_def = Some(TypeDefinition::TypeId(variable.type_ascription));
+                }
+            }
+
             handle_expression(&variable.body, tokens);
         }
         TypedDeclaration::ConstantDeclaration(const_decl) => {
@@ -269,9 +278,10 @@ fn handle_expression(expression: &TypedExpression, tokens: &TokenMap) {
             handle_expression(lhs, tokens);
             handle_expression(rhs, tokens);
         }
-        TypedExpressionVariant::VariableExpression { ref name } => {
-            if let Some(mut token) = tokens.get_mut(&to_ident_key(name)) {
+        TypedExpressionVariant::VariableExpression { ref name, ref span } => {
+            if let Some(mut token) = tokens.get_mut(&to_ident_key(&Ident::new(span.clone()))) {
                 token.typed = Some(TypedAstToken::TypedExpression(expression.clone()));
+                token.type_def = Some(TypeDefinition::Ident(name.clone()));
             }
         }
         TypedExpressionVariant::Tuple { fields } => {