Port TraitItem::TraitFn to be stored in the parsed declaration engine.

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

@@ -14,7 +14,7 @@ use sway_types::{ident::Ident, span::Span, Spanned};
 
 #[derive(Debug, Clone)]
 pub enum TraitItem {
-    TraitFn(TraitFn),
+    TraitFn(ParsedDeclId<TraitFn>),
     Constant(ParsedDeclId<ConstantDeclaration>),
     Type(ParsedDeclId<TraitTypeDeclaration>),
     // to handle parser recovery: Error represents an incomplete trait item

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

@@ -101,11 +101,12 @@ impl ty::TyAbiDecl {
 
                 for item in interface_surface.into_iter() {
                     let decl_name = match item {
-                        TraitItem::TraitFn(method) => {
+                        TraitItem::TraitFn(decl_id) => {
+                            let method = engines.pe().get_trait_fn(&decl_id);
                             // check that a super-trait does not define a method
                             // with the same name as the current interface method
                             error_on_shadowing_superabi_method(&method.name, &mut ctx);
-                            let method = ty::TyTraitFn::type_check(handler, ctx.by_ref(), method)?;
+                            let method = ty::TyTraitFn::type_check(handler, ctx.by_ref(), &method)?;
                             for param in &method.parameters {
                                 if param.is_reference || param.is_mutable {
                                     handler.emit_err(

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

@@ -68,7 +68,7 @@ impl ty::TyFunctionParameter {
     pub(crate) fn type_check_interface_parameter(
         handler: &Handler,
         mut ctx: TypeCheckContext,
-        parameter: FunctionParameter,
+        parameter: &FunctionParameter,
     ) -> Result<Self, ErrorEmitted> {
         let type_engine = ctx.engines.te();
         let engines = ctx.engines();
@@ -78,10 +78,11 @@ impl ty::TyFunctionParameter {
             is_reference,
             is_mutable,
             mutability_span,
-            mut type_argument,
+            type_argument,
         } = parameter;
 
-        type_argument.type_id = ctx
+        let mut new_type_argument = type_argument.clone();
+        new_type_argument.type_id = ctx
             .resolve_type(
                 handler,
                 type_argument.type_id,
@@ -92,11 +93,11 @@ impl ty::TyFunctionParameter {
             .unwrap_or_else(|err| type_engine.insert(engines, TypeInfo::ErrorRecovery(err), None));
 
         let typed_parameter = ty::TyFunctionParameter {
-            name,
-            is_reference,
-            is_mutable,
-            mutability_span,
-            type_argument,
+            name: name.clone(),
+            is_reference: *is_reference,
+            is_mutable: *is_mutable,
+            mutability_span: mutability_span.clone(),
+            type_argument: new_type_argument,
         };
 
         Ok(typed_parameter)

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

@@ -129,8 +129,9 @@ impl TyTraitDecl {
 
             for item in interface_surface.into_iter() {
                 let decl_name = match item {
-                    TraitItem::TraitFn(method) => {
-                        let method = ty::TyTraitFn::type_check(handler, ctx.by_ref(), method)?;
+                    TraitItem::TraitFn(decl_id) => {
+                        let method = engines.pe().get_trait_fn(&decl_id);
+                        let method = ty::TyTraitFn::type_check(handler, ctx.by_ref(), &method)?;
                         let decl_ref = decl_engine.insert(method.clone());
                         dummy_interface_surface.push(ty::TyImplItem::Fn(
                             decl_engine

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

@@ -16,27 +16,27 @@ impl ty::TyTraitFn {
     pub(crate) fn type_check(
         handler: &Handler,
         mut ctx: TypeCheckContext,
-        trait_fn: parsed::TraitFn,
+        trait_fn: &parsed::TraitFn,
     ) -> Result<ty::TyTraitFn, ErrorEmitted> {
         let parsed::TraitFn {
             name,
             span,
             purity,
             parameters,
-            mut return_type,
+            return_type,
             attributes,
         } = trait_fn;
 
         let type_engine = ctx.engines.te();
         let engines = ctx.engines();
 
         // Create a namespace for the trait function.
-        ctx.by_ref().with_purity(purity).scoped(|mut ctx| {
+        ctx.by_ref().with_purity(*purity).scoped(|mut ctx| {
             // TODO: when we add type parameters to trait fns, type check them here
 
             // Type check the parameters.
             let mut typed_parameters = vec![];
-            for param in parameters.into_iter() {
+            for param in parameters.iter() {
                 typed_parameters.push(
                     match ty::TyFunctionParameter::type_check_interface_parameter(
                         handler,
@@ -50,7 +50,8 @@ impl ty::TyTraitFn {
             }
 
             // Type check the return type.
-            return_type.type_id = ctx
+            let mut new_return_type = return_type.clone();
+            new_return_type.type_id = ctx
                 .resolve_type(
                     handler,
                     return_type.type_id,
@@ -63,12 +64,12 @@ impl ty::TyTraitFn {
                 });
 
             let trait_fn = ty::TyTraitFn {
-                name,
-                span,
+                name: name.clone(),
+                span: span.clone(),
                 parameters: typed_parameters,
-                return_type,
-                purity,
-                attributes,
+                return_type: new_return_type,
+                purity: *purity,
+                attributes: attributes.clone(),
             };
 
             Ok(trait_fn)

sway-core/src/semantic_analysis/node_dependencies.rs

@@ -345,11 +345,13 @@ impl Dependencies {
                 .gather_from_iter(
                     trait_decl.interface_surface.iter(),
                     |deps, item| match item {
-                        TraitItem::TraitFn(sig) => deps
-                            .gather_from_iter(sig.parameters.iter(), |deps, param| {
+                        TraitItem::TraitFn(decl_id) => {
+                            let sig = engines.pe().get_trait_fn(decl_id);
+                            deps.gather_from_iter(sig.parameters.iter(), |deps, param| {
                                 deps.gather_from_type_argument(engines, &param.type_argument)
                             })
-                            .gather_from_type_argument(engines, &sig.return_type),
+                            .gather_from_type_argument(engines, &sig.return_type)
+                        }
                         TraitItem::Constant(decl_id) => {
                             let const_decl = engines.pe().get_constant(decl_id);
                             deps.gather_from_constant_decl(engines, &const_decl)
@@ -429,11 +431,13 @@ impl Dependencies {
                     deps.gather_from_call_path(&sup.name, false, false)
                 })
                 .gather_from_iter(interface_surface.iter(), |deps, item| match item {
-                    TraitItem::TraitFn(sig) => deps
-                        .gather_from_iter(sig.parameters.iter(), |deps, param| {
+                    TraitItem::TraitFn(decl_id) => {
+                        let sig = engines.pe().get_trait_fn(decl_id);
+                        deps.gather_from_iter(sig.parameters.iter(), |deps, param| {
                             deps.gather_from_type_argument(engines, &param.type_argument)
                         })
-                        .gather_from_type_argument(engines, &sig.return_type),
+                        .gather_from_type_argument(engines, &sig.return_type)
+                    }
                     TraitItem::Constant(decl_id) => {
                         let const_decl = engines.pe().get_constant(decl_id);
                         deps.gather_from_constant_decl(engines, &const_decl)

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

@@ -844,7 +844,7 @@ fn item_abi_to_abi_declaration(
                                 context,
                                 handler,
                                 engines,
-                                &trait_fn.parameters,
+                                &engines.pe().get_trait_fn(&trait_fn).parameters,
                                 "an ABI method signature",
                             )?;
                             Ok(TraitItem::TraitFn(trait_fn))
@@ -1534,7 +1534,7 @@ fn fn_signature_to_trait_fn(
     engines: &Engines,
     fn_signature: FnSignature,
     attributes: AttributesMap,
-) -> Result<TraitFn, ErrorEmitted> {
+) -> Result<ParsedDeclId<TraitFn>, ErrorEmitted> {
     let return_type = match &fn_signature.return_type_opt {
         Some((_right_arrow, ty)) => ty_to_type_argument(context, handler, engines, ty.clone())?,
         None => {
@@ -1566,6 +1566,7 @@ fn fn_signature_to_trait_fn(
         )?,
         return_type,
     };
+    let trait_fn = engines.pe().insert(trait_fn);
     Ok(trait_fn)
 }
 

sway-lsp/src/core/token.rs

@@ -2,6 +2,7 @@ use lsp_types::{Position, Range};
 use std::path::PathBuf;
 use sway_ast::Intrinsic;
 use sway_core::{
+    decl_engine::parsed_id::ParsedDeclId,
     language::{
         parsed::{
             AbiCastExpression, AmbiguousPathExpression, Declaration, DelineatedPathExpression,
@@ -49,7 +50,7 @@ pub enum AstToken {
     StructScrutineeField(StructScrutineeField),
     Supertrait(Supertrait),
     TraitConstraint(TraitConstraint),
-    TraitFn(TraitFn),
+    TraitFn(ParsedDeclId<TraitFn>),
     TypeArgument(TypeArgument),
     TypeParameter(TypeParameter),
     UseStatement(UseStatement),

sway-lsp/src/traverse/parsed_tree.rs

@@ -987,17 +987,18 @@ impl Parse for Supertrait {
     }
 }
 
-impl Parse for TraitFn {
+impl Parse for ParsedDeclId<TraitFn> {
     fn parse(&self, ctx: &ParseContext) {
+        let trait_fn = ctx.engines.pe().get_trait_fn(self);
         ctx.tokens.insert(
-            ctx.ident(&self.name),
-            Token::from_parsed(AstToken::TraitFn(self.clone()), SymbolKind::Function),
+            ctx.ident(&trait_fn.name),
+            Token::from_parsed(AstToken::TraitFn(*self), SymbolKind::Function),
         );
-        self.parameters.par_iter().for_each(|param| {
+        trait_fn.parameters.par_iter().for_each(|param| {
             param.parse(ctx);
         });
-        self.return_type.parse(ctx);
-        self.attributes.parse(ctx);
+        trait_fn.return_type.parse(ctx);
+        trait_fn.attributes.parse(ctx);
     }
 }
 

sway-lsp/src/utils/attributes.rs

@@ -36,7 +36,10 @@ where
         },
         AstToken::StorageField(field) => callback(&field.attributes),
         AstToken::StructField(field) => callback(&field.attributes),
-        AstToken::TraitFn(trait_fn) => callback(&trait_fn.attributes),
+        AstToken::TraitFn(decl_id) => {
+            let decl = engines.pe().get_trait_fn(decl_id);
+            callback(&decl.attributes)
+        }
         AstToken::EnumVariant(variant) => callback(&variant.attributes),
         _ => {}
     }