LSP: generate doc comments for trait functions (FuelLabs#4665)

## Description

Closes FuelLabs#4645

This PR lets the user generate doc comments for ABI methods. Previously
it was working for functions and function implementations of ABI/trait,
but not on the ABI/trait definition itself.

Since the output is exactly the same for FunctionDecl to TraitFn, I
wanted to do this in a way where they could share the code. I ended up
writing a trait called `FunctionSignature` that both types implement
that exposes `parameters()` and `return_type()`. This allows me to use
the same code for code actions and I suspect it will come in handy in
the future.

One problem was that `FunctionDecl`'s return_type was `TypeArgument`
while `TraitFn` had its return type as `TypeId` with an additional
`return_type_span` field. I changed it so that `TraitFn`'s return_type
is now `TypeArgument` and removed `return_type_span`.

I also needed to add a `span` field to `TraitFn` since previously its
`span()` was returning only the span of the function name, not the whole
function. From what I could see, this doesn't impact the spans of any
existing errors.

### Testing

I added an integration test to the LSP for this scenario.

### Example

![Jun-13-2023
14-40-21](https://github.com/FuelLabs/sway/assets/47993817/5ba750b6-9f9d-4cb9-96a8-c791bf93b77d)

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [ ] 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.

forc-plugins/forc-doc/src/render/item/context.rs

@@ -155,7 +155,7 @@ impl Renderable for Context {
                             )?;
                         }
                     }
-                    write!(fn_sig, ") -> {}", method.return_type_span.as_str())?;
+                    write!(fn_sig, ") -> {}", method.return_type.span.as_str())?;
                     let multiline = fn_sig.chars().count() >= 60;
                     let fn_sig = format!("fn {}(", method.name);
                     let method_id = format!("tymethod.{}", method.name.as_str());
@@ -215,9 +215,9 @@ impl Renderable for Context {
                                     }
                                     : ")";
                                 }
-                                @ if !method.return_type_span.as_str().contains(&fn_sig) {
+                                @ if !method.return_type.span.as_str().contains(&fn_sig) {
                                     : " -> ";
-                                    : method.return_type_span.as_str();
+                                    : method.return_type.span.as_str();
                                 }
                             }
                         }

sway-core/src/control_flow_analysis/dead_code_analysis.rs

@@ -791,7 +791,7 @@ fn connect_abi_declaration(
                 if let Some(TypeInfo::Struct(decl_ref)) = get_struct_type_info_from_type_id(
                     type_engine,
                     decl_engine,
-                    fn_decl.return_type,
+                    fn_decl.return_type.type_id,
                 )? {
                     let decl = decl_engine.get_struct(&decl_ref);
                     if let Some(ns) = graph.namespace.get_struct(&decl.call_path.suffix).cloned() {

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

@@ -63,9 +63,9 @@ impl HashWithEngines for Supertrait {
 #[derive(Debug, Clone)]
 pub struct TraitFn {
     pub name: Ident,
+    pub span: Span,
     pub attributes: transform::AttributesMap,
     pub purity: Purity,
     pub parameters: Vec<FunctionParameter>,
-    pub return_type: TypeInfo,
-    pub return_type_span: Span,
+    pub return_type: TypeArgument,
 }

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

@@ -43,6 +43,16 @@ impl Named for TyFunctionDecl {
     }
 }
 
+impl declaration::FunctionSignature for TyFunctionDecl {
+    fn parameters(&self) -> &Vec<TyFunctionParameter> {
+        &self.parameters
+    }
+
+    fn return_type(&self) -> &TypeArgument {
+        &self.return_type
+    }
+}
+
 impl EqWithEngines for TyFunctionDecl {}
 impl PartialEqWithEngines for TyFunctionDecl {
     fn eq(&self, other: &Self, engines: &Engines) -> bool {

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

@@ -24,3 +24,10 @@ pub use storage::*;
 pub use trait_fn::*;
 pub use type_alias::*;
 pub use variable::*;
+
+use crate::TypeArgument;
+
+pub trait FunctionSignature {
+    fn parameters(&self) -> &Vec<TyFunctionParameter>;
+    fn return_type(&self) -> &TypeArgument;
+}

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

@@ -12,10 +12,10 @@ use crate::{
 #[derive(Clone, Debug)]
 pub struct TyTraitFn {
     pub name: Ident,
+    pub(crate) span: Span,
     pub(crate) purity: Purity,
     pub parameters: Vec<TyFunctionParameter>,
-    pub return_type: TypeId,
-    pub return_type_span: Span,
+    pub return_type: TypeArgument,
     pub attributes: transform::AttributesMap,
 }
 
@@ -27,7 +27,17 @@ impl Named for TyTraitFn {
 
 impl Spanned for TyTraitFn {
     fn span(&self) -> Span {
-        self.name.span()
+        self.span.clone()
+    }
+}
+
+impl declaration::FunctionSignature for TyTraitFn {
+    fn parameters(&self) -> &Vec<TyFunctionParameter> {
+        &self.parameters
+    }
+
+    fn return_type(&self) -> &TypeArgument {
+        &self.return_type
     }
 }
 
@@ -39,8 +49,8 @@ impl PartialEqWithEngines for TyTraitFn {
             && self.purity == other.purity
             && self.parameters.eq(&other.parameters, engines)
             && type_engine
-                .get(self.return_type)
-                .eq(&type_engine.get(other.return_type), engines)
+                .get(self.return_type.type_id)
+                .eq(&type_engine.get(other.return_type.type_id), engines)
             && self.attributes == other.attributes
     }
 }
@@ -54,13 +64,13 @@ impl HashWithEngines for TyTraitFn {
             return_type,
             // these fields are not hashed because they aren't relevant/a
             // reliable source of obj v. obj distinction
-            return_type_span: _,
+            span: _,
             attributes: _,
         } = self;
         let type_engine = engines.te();
         name.hash(state);
         parameters.hash(state, engines);
-        type_engine.get(*return_type).hash(state, engines);
+        type_engine.get(return_type.type_id).hash(state, engines);
         purity.hash(state);
     }
 }

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

@@ -809,10 +809,10 @@ fn type_check_impl_method(
         (true, true) | (false, false) => (), // no payability mismatch
     }
 
-    if !type_engine
-        .get(impl_method.return_type.type_id)
-        .eq(&type_engine.get(impl_method_signature.return_type), engines)
-    {
+    if !type_engine.get(impl_method.return_type.type_id).eq(
+        &type_engine.get(impl_method_signature.return_type.type_id),
+        engines,
+    ) {
         errors.push(CompileError::MismatchedTypeInInterfaceSurface {
             interface_name: interface_name(),
             span: impl_method.return_type.span.clone(),

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

@@ -17,10 +17,10 @@ impl ty::TyTraitFn {
 
         let parsed::TraitFn {
             name,
+            span,
             purity,
             parameters,
-            return_type,
-            return_type_span,
+            mut return_type,
             attributes,
         } = trait_fn;
 
@@ -45,10 +45,10 @@ impl ty::TyTraitFn {
         }
 
         // Type check the return type.
-        let return_type = check!(
+        return_type.type_id = check!(
             ctx.resolve_type_with_self(
-                type_engine.insert(engines, return_type),
-                &return_type_span,
+                return_type.type_id,
+                &return_type.span,
                 EnforceTypeArguments::Yes,
                 None
             ),
@@ -59,9 +59,9 @@ impl ty::TyTraitFn {
 
         let trait_fn = ty::TyTraitFn {
             name,
+            span,
             parameters: typed_parameters,
             return_type,
-            return_type_span,
             purity,
             attributes,
         };
@@ -81,12 +81,7 @@ impl ty::TyTraitFn {
             implementing_type: None,
             span: self.name.span(),
             attributes: self.attributes.clone(),
-            return_type: TypeArgument {
-                type_id: self.return_type,
-                initial_type_id: self.return_type,
-                span: self.return_type_span.clone(),
-                call_path_tree: None,
-            },
+            return_type: self.return_type.clone(),
             visibility: Visibility::Public,
             type_parameters: vec![],
             is_contract_call: mode == Mode::ImplAbiFn,

sway-core/src/semantic_analysis/node_dependencies.rs

@@ -339,7 +339,7 @@ impl Dependencies {
                         .gather_from_iter(sig.parameters.iter(), |deps, param| {
                             deps.gather_from_type_argument(engines, &param.type_argument)
                         })
-                        .gather_from_typeinfo(engines, &sig.return_type),
+                        .gather_from_type_argument(engines, &sig.return_type),
                     TraitItem::Constant(const_decl) => {
                         deps.gather_from_constant_decl(engines, const_decl)
                     }
@@ -389,7 +389,7 @@ impl Dependencies {
                         .gather_from_iter(sig.parameters.iter(), |deps, param| {
                             deps.gather_from_type_argument(engines, &param.type_argument)
                         })
-                        .gather_from_typeinfo(engines, &sig.return_type),
+                        .gather_from_type_argument(engines, &sig.return_type),
                     TraitItem::Constant(const_decl) => {
                         deps.gather_from_constant_decl(engines, const_decl)
                     }

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

@@ -1374,12 +1374,23 @@ fn fn_signature_to_trait_fn(
     fn_signature: FnSignature,
     attributes: AttributesMap,
 ) -> Result<TraitFn, ErrorEmitted> {
-    let return_type_span = match &fn_signature.return_type_opt {
-        Some((_right_arrow_token, ty)) => ty.span(),
-        None => fn_signature.span(),
+    let return_type = match &fn_signature.return_type_opt {
+        Some((_right_arrow, ty)) => ty_to_type_argument(context, handler, engines, ty.clone())?,
+        None => {
+            let type_id = engines.te().insert(engines, TypeInfo::Tuple(Vec::new()));
+            TypeArgument {
+                type_id,
+                initial_type_id: type_id,
+                // TODO: Fix as part of https://github.com/FuelLabs/sway/issues/3635
+                span: fn_signature.span(),
+                call_path_tree: None,
+            }
+        }
     };
+
     let trait_fn = TraitFn {
-        name: fn_signature.name,
+        name: fn_signature.name.clone(),
+        span: fn_signature.span(),
         purity: get_attributed_purity(context, handler, &attributes)?,
         attributes,
         parameters: fn_args_to_function_parameters(
@@ -1388,11 +1399,7 @@ fn fn_signature_to_trait_fn(
             engines,
             fn_signature.arguments.into_inner(),
         )?,
-        return_type: match fn_signature.return_type_opt {
-            Some((_right_arrow_token, ty)) => ty_to_type_info(context, handler, engines, ty)?,
-            None => TypeInfo::Tuple(Vec::new()),
-        },
-        return_type_span,
+        return_type,
     };
     Ok(trait_fn)
 }

sway-lsp/src/capabilities/code_actions/abi_decl/abi_impl.rs

@@ -59,10 +59,10 @@ impl AbiImplCodeAction<'_> {
     fn return_type_string(&self, function_decl: &TyTraitFn) -> String {
         let type_engine = self.engines.te();
         // Unit is the implicit return type for ABI functions.
-        if type_engine.get(function_decl.return_type).is_unit() {
+        if type_engine.get(function_decl.return_type.type_id).is_unit() {
             String::from("")
         } else {
-            format!(" -> {}", function_decl.return_type_span.as_str())
+            format!(" -> {}", function_decl.return_type.span.as_str())
         }
     }
 

sway-lsp/src/capabilities/code_actions/common/basic_doc_comment.rs

@@ -0,0 +1,38 @@
+use crate::capabilities::code_actions::{CodeAction, CodeActionContext, CODE_ACTION_DOC_TITLE};
+use lsp_types::{Range, Url};
+use sway_types::Spanned;
+
+use super::generate_doc::GenerateDocCodeAction;
+
+pub struct BasicDocCommentCodeAction<'a, T: Spanned> {
+    decl: &'a T,
+    uri: &'a Url,
+}
+
+impl<'a, T: Spanned> GenerateDocCodeAction<'a, T> for BasicDocCommentCodeAction<'a, T> {}
+
+impl<'a, T: Spanned> CodeAction<'a, T> for BasicDocCommentCodeAction<'a, T> {
+    fn new(ctx: CodeActionContext<'a>, decl: &'a T) -> Self {
+        Self { decl, uri: ctx.uri }
+    }
+
+    fn new_text(&self) -> String {
+        self.default_template()
+    }
+
+    fn range(&self) -> Range {
+        self.range_before()
+    }
+
+    fn title(&self) -> String {
+        CODE_ACTION_DOC_TITLE.to_string()
+    }
+
+    fn decl(&self) -> &T {
+        self.decl
+    }
+
+    fn uri(&self) -> &Url {
+        self.uri
+    }
+}