From dfed81b4b39b2f783d6e81a78ee27fba7032e01c Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Fri, 13 Sep 2024 12:11:13 -0300 Subject: [PATCH] feat: let LSP suggest trait impl methods as you are typing them (#6029) # Description ## Problem Part of #1577 ## Summary Now that we have the "Implement missing members" code action, it's a small step to have autocompletion suggest a trait method as you are typing its name (like in Rust Analyzer). The only tricky part was that `fn foo` produced a parse error so I made that parse to a function without arguments and without body (but still reporting a parser error). A nice side-effect of this is that if you type that and you think "Hm, what should the parameters be" the rest of the LSP features don't stop working (like, inlay hints don't disappear, etc.) Even though "Implement missing members" might seem to be better than offering completions one by one, this completion has two advantages: 1. It offers default methods 2. It teaches the user that LSP can suggest trait methods, so they could think "hm, I wonder if there's a code action that suggests all of them" (at least this was my experience for discovering the code action) ![lsp-suggest-trait-impl-method](https://github.com/user-attachments/assets/00d1a6f8-597b-4686-ab20-78cc145f22f4) ## Additional Context ## Documentation Check one: - [x] No documentation needed. - [ ] Documentation included in this PR. - [ ] **[For Experimental Features]** Documentation to be submitted in a separate PR. # PR Checklist - [x] I have tested the changes locally. - [x] I have formatted the changes with [Prettier](https://prettier.io/) and/or `cargo fmt` on default settings. --- compiler/noirc_frontend/src/parser/errors.rs | 2 + .../src/parser/parser/function.rs | 78 ++- tooling/lsp/src/lib.rs | 1 + .../code_action/implement_missing_members.rs | 440 +---------------- tooling/lsp/src/requests/completion.rs | 102 +++- .../requests/completion/completion_items.rs | 9 +- tooling/lsp/src/requests/completion/tests.rs | 47 ++ .../src/trait_impl_method_stub_generator.rs | 449 ++++++++++++++++++ 8 files changed, 667 insertions(+), 461 deletions(-) create mode 100644 tooling/lsp/src/trait_impl_method_stub_generator.rs diff --git a/compiler/noirc_frontend/src/parser/errors.rs b/compiler/noirc_frontend/src/parser/errors.rs index 6ba4cb68500..aa07972ae5d 100644 --- a/compiler/noirc_frontend/src/parser/errors.rs +++ b/compiler/noirc_frontend/src/parser/errors.rs @@ -31,6 +31,8 @@ pub enum ParserErrorReason { ExpectedLeftBracketOrWhereOrLeftBraceOrArrowAfterImplType, #[error("expected <, where or {{ after trait impl for type")] ExpectedLeftBracketOrWhereOrLeftBraceOrArrowAfterTraitImplForType, + #[error("expected ( or < after function name")] + ExpectedLeftParenOrLeftBracketAfterFunctionName, #[error("Expected a ; separating these two statements")] MissingSeparatingSemi, #[error("constrain keyword is deprecated")] diff --git a/compiler/noirc_frontend/src/parser/parser/function.rs b/compiler/noirc_frontend/src/parser/parser/function.rs index 05138bfffd9..dd52f640622 100644 --- a/compiler/noirc_frontend/src/parser/parser/function.rs +++ b/compiler/noirc_frontend/src/parser/parser/function.rs @@ -48,30 +48,66 @@ pub(super) fn function_definition(allow_self: bool) -> impl NoirParser CodeActionFinder<'a> { for (name, func_id) in method_ids { let func_meta = self.interner.function_meta(func_id); - let mut generator = MethodStubGenerator::new( + let mut generator = TraitImplMethodStubGenerator::new( name, func_meta, trait_, @@ -120,6 +112,8 @@ impl<'a> CodeActionFinder<'a> { self.module_id, indent + 4, ); + generator.set_body(format!("panic(f\"Implement {}\")", name)); + let stub = generator.generate(); stubs.push(stub); } @@ -136,430 +130,6 @@ impl<'a> CodeActionFinder<'a> { } } -struct MethodStubGenerator<'a> { - name: &'a str, - func_meta: &'a FuncMeta, - trait_: &'a Trait, - noir_trait_impl: &'a NoirTraitImpl, - interner: &'a NodeInterner, - def_maps: &'a BTreeMap, - module_id: ModuleId, - indent: usize, - string: String, -} - -impl<'a> MethodStubGenerator<'a> { - #[allow(clippy::too_many_arguments)] - fn new( - name: &'a str, - func_meta: &'a FuncMeta, - trait_: &'a Trait, - noir_trait_impl: &'a NoirTraitImpl, - interner: &'a NodeInterner, - def_maps: &'a BTreeMap, - module_id: ModuleId, - indent: usize, - ) -> Self { - Self { - name, - func_meta, - trait_, - noir_trait_impl, - interner, - def_maps, - module_id, - indent, - string: String::new(), - } - } - - fn generate(&mut self) -> String { - let indent_string = " ".repeat(self.indent); - - self.string.push_str(&indent_string); - self.string.push_str("fn "); - self.string.push_str(self.name); - self.append_resolved_generics(&self.func_meta.direct_generics); - self.string.push('('); - for (index, (pattern, typ, _visibility)) in self.func_meta.parameters.iter().enumerate() { - if index > 0 { - self.string.push_str(", "); - } - if self.append_pattern(pattern) { - self.string.push_str(": "); - self.append_type(typ); - } - } - self.string.push(')'); - - let return_type = self.func_meta.return_type(); - if return_type != &Type::Unit { - self.string.push_str(" -> "); - self.append_type(return_type); - } - - if !self.func_meta.trait_constraints.is_empty() { - self.string.push_str(" where "); - for (index, constraint) in self.func_meta.trait_constraints.iter().enumerate() { - if index > 0 { - self.string.push_str(", "); - } - self.append_type(&constraint.typ); - self.string.push_str(": "); - let trait_ = self.interner.get_trait(constraint.trait_id); - self.string.push_str(&trait_.name.0.contents); - self.append_trait_generics(&constraint.trait_generics); - } - } - - self.string.push_str(" {\n"); - - let body_indent_string = " ".repeat(self.indent + 4); - self.string.push_str(&body_indent_string); - self.string.push_str("panic(f\"Implement "); - self.string.push_str(self.name); - self.string.push_str("\")\n"); - self.string.push_str(&indent_string); - self.string.push_str("}\n"); - std::mem::take(&mut self.string) - } - - /// Appends a pattern and returns true if this was not the self type - fn append_pattern(&mut self, pattern: &HirPattern) -> bool { - match pattern { - HirPattern::Identifier(hir_ident) => { - let definition = self.interner.definition(hir_ident.id); - self.string.push_str(&definition.name); - &definition.name != "self" - } - HirPattern::Mutable(pattern, _) => { - self.string.push_str("mut "); - self.append_pattern(pattern) - } - HirPattern::Tuple(patterns, _) => { - self.string.push('('); - for (index, pattern) in patterns.iter().enumerate() { - if index > 0 { - self.string.push_str(", "); - } - self.append_pattern(pattern); - } - self.string.push(')'); - true - } - HirPattern::Struct(typ, patterns, _) => { - self.append_type(typ); - self.string.push_str(" { "); - for (index, (name, _pattern)) in patterns.iter().enumerate() { - if index > 0 { - self.string.push_str(", "); - } - self.string.push_str(&name.0.contents); - } - self.string.push_str(" }"); - true - } - } - } - - fn append_type(&mut self, typ: &Type) { - match typ { - Type::FieldElement => self.string.push_str("Field"), - Type::Array(n, e) => { - self.string.push('['); - self.append_type(e); - self.string.push_str("; "); - self.append_type(n); - self.string.push(']'); - } - Type::Slice(typ) => { - self.string.push('['); - self.append_type(typ); - self.string.push(']'); - } - Type::Tuple(types) => { - self.string.push('('); - for (index, typ) in types.iter().enumerate() { - if index > 0 { - self.string.push_str(", "); - } - self.append_type(typ); - } - self.string.push(')'); - } - Type::Struct(struct_type, generics) => { - let struct_type = struct_type.borrow(); - - let current_module_data = - &self.def_maps[&self.module_id.krate].modules()[self.module_id.local_id.0]; - - // Check if the struct type is already imported/visible in this module - let per_ns = current_module_data.find_name(&struct_type.name); - if let Some((module_def_id, _, _)) = per_ns.types { - if module_def_id == ModuleDefId::TypeId(struct_type.id) { - self.string.push_str(&struct_type.name.0.contents); - self.append_generics(generics); - return; - } - } - - let module_id = struct_type.id.module_id(); - let module_data = &self.def_maps[&module_id.krate].modules()[module_id.local_id.0]; - let parent_module_local_id = module_data.parent.unwrap(); - let parent_module_id = - ModuleId { krate: module_id.krate, local_id: parent_module_local_id }; - - let current_module_parent_id = current_module_data - .parent - .map(|parent| ModuleId { krate: self.module_id.krate, local_id: parent }); - - let relative_path = relative_module_id_path( - parent_module_id, - &self.module_id, - current_module_parent_id, - self.interner, - ); - - if !relative_path.is_empty() { - self.string.push_str(&relative_path); - self.string.push_str("::"); - } - self.string.push_str(&struct_type.name.0.contents); - self.append_generics(generics); - } - Type::Alias(type_alias, generics) => { - let type_alias = type_alias.borrow(); - - let current_module_data = - &self.def_maps[&self.module_id.krate].modules()[self.module_id.local_id.0]; - - // Check if the alias type is already imported/visible in this module - let per_ns = current_module_data.find_name(&type_alias.name); - if let Some((module_def_id, _, _)) = per_ns.types { - if module_def_id == ModuleDefId::TypeAliasId(type_alias.id) { - self.string.push_str(&type_alias.name.0.contents); - self.append_generics(generics); - return; - } - } - - let parent_module_id = - self.interner.reference_module(ReferenceId::Alias(type_alias.id)).unwrap(); - - let current_module_parent_id = current_module_data - .parent - .map(|parent| ModuleId { krate: self.module_id.krate, local_id: parent }); - - let relative_path = relative_module_id_path( - *parent_module_id, - &self.module_id, - current_module_parent_id, - self.interner, - ); - - if !relative_path.is_empty() { - self.string.push_str(&relative_path); - self.string.push_str("::"); - } - self.string.push_str(&type_alias.name.0.contents); - self.append_generics(generics); - } - Type::TraitAsType(trait_id, _, trait_generics) => { - let trait_ = self.interner.get_trait(*trait_id); - - let current_module_data = - &self.def_maps[&self.module_id.krate].modules()[self.module_id.local_id.0]; - - // Check if the trait type is already imported/visible in this module - let per_ns = current_module_data.find_name(&trait_.name); - if let Some((module_def_id, _, _)) = per_ns.types { - if module_def_id == ModuleDefId::TraitId(*trait_id) { - self.string.push_str(&trait_.name.0.contents); - self.append_trait_generics(trait_generics); - return; - } - } - - let parent_module_id = - self.interner.reference_module(ReferenceId::Trait(*trait_id)).unwrap(); - - let current_module_parent_id = current_module_data - .parent - .map(|parent| ModuleId { krate: self.module_id.krate, local_id: parent }); - - let relative_path = relative_module_id_path( - *parent_module_id, - &self.module_id, - current_module_parent_id, - self.interner, - ); - - if !relative_path.is_empty() { - self.string.push_str(&relative_path); - self.string.push_str("::"); - } - self.string.push_str(&trait_.name.0.contents); - self.append_trait_generics(trait_generics); - } - Type::TypeVariable(typevar, _) => { - if typevar.id() == self.trait_.self_type_typevar.id() { - self.string.push_str("Self"); - return; - } - - let generics = &self.trait_.generics; - if let Some(index) = - generics.iter().position(|generic| generic.type_var.id() == typevar.id()) - { - if let Some(typ) = self.noir_trait_impl.trait_generics.ordered_args.get(index) { - self.string.push_str(&typ.to_string()); - return; - } - } - - for associated_type in &self.trait_.associated_types { - if typevar.id() == associated_type.type_var.id() { - self.string.push_str("Self::"); - self.string.push_str(&associated_type.name); - return; - } - } - - for generic in &self.func_meta.direct_generics { - if typevar.id() == generic.type_var.id() { - self.string.push_str(&generic.name); - return; - } - } - - self.string.push_str("error"); - } - Type::NamedGeneric(typevar, _name, _kind) => { - self.append_type(&Type::TypeVariable(typevar.clone(), TypeVariableKind::Normal)); - } - Type::Function(args, ret, env, unconstrained) => { - if *unconstrained { - self.string.push_str("unconstrained "); - } - self.string.push_str("fn"); - - if let Type::Unit = **env { - } else { - self.string.push('['); - self.append_type(env); - self.string.push(']'); - } - - self.string.push('('); - for (index, arg) in args.iter().enumerate() { - if index > 0 { - self.string.push_str(", "); - } - self.append_type(arg); - } - self.string.push(')'); - - if let Type::Unit = **ret { - } else { - self.string.push_str(" -> "); - self.append_type(ret); - } - } - Type::MutableReference(typ) => { - self.string.push_str("&mut "); - self.append_type(typ); - } - Type::Forall(_, _) => { - panic!("Shouldn't get a Type::Forall"); - } - Type::InfixExpr(left, op, right) => { - self.append_type(left); - self.string.push(' '); - self.string.push_str(&op.to_string()); - self.string.push(' '); - self.append_type(right); - } - Type::Constant(_) - | Type::Integer(_, _) - | Type::Bool - | Type::String(_) - | Type::FmtString(_, _) - | Type::Unit - | Type::Quoted(_) - | Type::Error => self.string.push_str(&typ.to_string()), - } - } - - fn append_generics(&mut self, generics: &[Type]) { - if generics.is_empty() { - return; - } - - self.string.push('<'); - for (index, typ) in generics.iter().enumerate() { - if index > 0 { - self.string.push_str(", "); - } - self.append_type(typ); - } - self.string.push('>'); - } - - fn append_trait_generics(&mut self, generics: &TraitGenerics) { - if generics.named.is_empty() && generics.ordered.is_empty() { - return; - } - - let mut index = 0; - - self.string.push('<'); - for generic in &generics.ordered { - if index > 0 { - self.string.push_str(", "); - } - self.append_type(generic); - index += 1; - } - for named_type in &generics.named { - if index > 0 { - self.string.push_str(", "); - } - self.string.push_str(&named_type.name.0.contents); - self.string.push_str(" = "); - self.append_type(&named_type.typ); - index += 1; - } - self.string.push('>'); - } - - fn append_resolved_generics(&mut self, generics: &[ResolvedGeneric]) { - if generics.is_empty() { - return; - } - - self.string.push('<'); - for (index, generic) in self.func_meta.direct_generics.iter().enumerate() { - if index > 0 { - self.string.push_str(", "); - } - self.append_resolved_generic(generic); - } - self.string.push('>'); - } - - fn append_resolved_generic(&mut self, generic: &ResolvedGeneric) { - match &generic.kind { - Kind::Normal => self.string.push_str(&generic.name), - Kind::Numeric(typ) => { - self.string.push_str("let "); - self.string.push_str(&generic.name); - self.string.push_str(": "); - self.append_type(typ); - } - } - } -} - #[cfg(test)] mod tests { use tokio::test; diff --git a/tooling/lsp/src/requests/completion.rs b/tooling/lsp/src/requests/completion.rs index d07ad826094..18e406860ab 100644 --- a/tooling/lsp/src/requests/completion.rs +++ b/tooling/lsp/src/requests/completion.rs @@ -4,7 +4,10 @@ use std::{ }; use async_lsp::ResponseError; -use completion_items::{field_completion_item, simple_completion_item, snippet_completion_item}; +use completion_items::{ + field_completion_item, simple_completion_item, snippet_completion_item, + trait_impl_method_completion_item, +}; use convert_case::{Case, Casing}; use fm::{FileId, FileMap, PathString}; use kinds::{FunctionCompletionKind, FunctionKind, RequestedItems}; @@ -15,8 +18,9 @@ use noirc_frontend::{ AsTraitPath, AttributeTarget, BlockExpression, CallExpression, ConstructorExpression, Expression, ExpressionKind, ForLoopStatement, GenericTypeArgs, Ident, IfExpression, ItemVisibility, Lambda, LetStatement, MemberAccessExpression, MethodCallExpression, - NoirFunction, NoirStruct, NoirTraitImpl, Path, PathKind, Pattern, Statement, TypeImpl, - UnresolvedGeneric, UnresolvedGenerics, UnresolvedType, UseTree, UseTreeKind, Visitor, + NoirFunction, NoirStruct, NoirTraitImpl, Path, PathKind, Pattern, Statement, + TraitImplItemKind, TypeImpl, UnresolvedGeneric, UnresolvedGenerics, UnresolvedType, + UseTree, UseTreeKind, Visitor, }, graph::{CrateId, Dependency}, hir::def_map::{CrateDefMap, LocalModuleId, ModuleId}, @@ -29,7 +33,10 @@ use noirc_frontend::{ }; use sort_text::underscore_sort_text; -use crate::{requests::to_lsp_location, utils, visibility::is_visible, LspState}; +use crate::{ + requests::to_lsp_location, trait_impl_method_stub_generator::TraitImplMethodStubGenerator, + utils, visibility::is_visible, LspState, +}; use super::process_request; @@ -81,6 +88,7 @@ pub(crate) fn on_completion_request( struct NodeFinder<'a> { files: &'a FileMap, file: FileId, + source: &'a str, lines: Vec<&'a str>, byte_index: usize, byte: Option, @@ -133,6 +141,7 @@ impl<'a> NodeFinder<'a> { Self { files, file, + source, lines: source.lines().collect(), byte_index, byte, @@ -843,6 +852,73 @@ impl<'a> NodeFinder<'a> { } } + fn suggest_trait_impl_function( + &mut self, + noir_trait_impl: &NoirTraitImpl, + noir_function: &NoirFunction, + ) { + // First find the trait + let location = Location::new(noir_trait_impl.trait_name.span(), self.file); + let Some(ReferenceId::Trait(trait_id)) = self.interner.find_referenced(location) else { + return; + }; + + let trait_ = self.interner.get_trait(trait_id); + + // Get all methods + let mut method_ids = trait_.method_ids.clone(); + + // Remove the ones that already are implemented + for item in &noir_trait_impl.items { + if let TraitImplItemKind::Function(noir_function) = &item.item.kind { + method_ids.remove(noir_function.name()); + } + } + + let indent = 0; + + // Suggest the ones that match the name + let prefix = noir_function.name(); + for (name, func_id) in method_ids { + if !name_matches(&name, prefix) { + continue; + } + + let func_meta = self.interner.function_meta(&func_id); + + let mut generator = TraitImplMethodStubGenerator::new( + &name, + func_meta, + trait_, + noir_trait_impl, + self.interner, + self.def_maps, + self.module_id, + indent, + ); + generator.set_body("${1}".to_string()); + + let stub = generator.generate(); + + // We don't need the initial indent nor the final newlines + let stub = stub.trim(); + // We also don't need the leading "fn " as that's already in the code; + let stub = stub.strip_prefix("fn ").unwrap(); + + let label = if func_meta.parameters.is_empty() { + format!("fn {}()", &name) + } else { + format!("fn {}(..)", &name) + }; + + let completion_item = trait_impl_method_completion_item(label, stub); + let completion_item = self + .completion_item_with_doc_comments(ReferenceId::Function(func_id), completion_item); + + self.completion_items.push(completion_item); + } + } + fn try_set_self_type(&mut self, pattern: &Pattern) { match pattern { Pattern::Identifier(ident) => { @@ -949,6 +1025,24 @@ impl<'a> Visitor for NodeFinder<'a> { self.collect_type_parameters_in_generics(&noir_trait_impl.impl_generics); for item in &noir_trait_impl.items { + if let TraitImplItemKind::Function(noir_function) = &item.item.kind { + // Check if it's `fn foo>|<` and neither `(` nor `<` follow + if noir_function.name_ident().span().end() as usize == self.byte_index + && noir_function.parameters().is_empty() + { + let bytes = self.source.as_bytes(); + let mut cursor = self.byte_index; + while cursor < bytes.len() && bytes[cursor].is_ascii_whitespace() { + cursor += 1; + } + let char = bytes[cursor] as char; + if char != '(' && char != '<' { + self.suggest_trait_impl_function(noir_trait_impl, noir_function); + return false; + } + } + } + item.item.accept(self); } diff --git a/tooling/lsp/src/requests/completion/completion_items.rs b/tooling/lsp/src/requests/completion/completion_items.rs index 163a4e15d00..56b1776b228 100644 --- a/tooling/lsp/src/requests/completion/completion_items.rs +++ b/tooling/lsp/src/requests/completion/completion_items.rs @@ -320,7 +320,7 @@ impl<'a> NodeFinder<'a> { text } - fn completion_item_with_doc_comments( + pub(super) fn completion_item_with_doc_comments( &self, id: ReferenceId, completion_item: CompletionItem, @@ -368,6 +368,13 @@ pub(super) fn module_completion_item(name: impl Into) -> CompletionItem ) } +pub(super) fn trait_impl_method_completion_item( + label: impl Into, + insert_text: impl Into, +) -> CompletionItem { + snippet_completion_item(label, CompletionItemKind::METHOD, insert_text, None) +} + fn func_meta_type_to_string(func_meta: &FuncMeta, has_self_type: bool) -> String { let mut typ = &func_meta.typ; if let Type::Forall(_, typ_) = typ { diff --git a/tooling/lsp/src/requests/completion/tests.rs b/tooling/lsp/src/requests/completion/tests.rs index 6809e24e645..f019d9980ab 100644 --- a/tooling/lsp/src/requests/completion/tests.rs +++ b/tooling/lsp/src/requests/completion/tests.rs @@ -8,6 +8,7 @@ mod completion_tests { completion_item_with_detail, completion_item_with_sort_text, completion_item_with_trigger_parameter_hints_command, module_completion_item, simple_completion_item, snippet_completion_item, + trait_impl_method_completion_item, }, sort_text::{auto_import_sort_text, self_mismatch_sort_text}, }, @@ -1968,4 +1969,50 @@ mod completion_tests { ) .await; } + + #[test] + async fn test_suggests_trait_impl_function() { + let src = r#" + trait Trait { + fn foo(x: i32) -> i32; + } + + struct Foo {} + + impl Trait for Foo { + fn f>|< + }"#; + + assert_completion( + src, + vec![trait_impl_method_completion_item( + "fn foo(..)", + "foo(x: i32) -> i32 {\n ${1}\n}", + )], + ) + .await; + } + + #[test] + async fn test_suggests_trait_impl_default_function() { + let src = r#" + trait Trait { + fn foo(x: i32) -> i32 { 1 } + } + + struct Foo {} + + impl Trait for Foo { + fn f>|< + }"#; + + assert_completion( + src, + vec![trait_impl_method_completion_item( + "fn foo(..)", + "foo(x: i32) -> i32 {\n ${1}\n}", + )], + ) + .await; + } } diff --git a/tooling/lsp/src/trait_impl_method_stub_generator.rs b/tooling/lsp/src/trait_impl_method_stub_generator.rs new file mode 100644 index 00000000000..e4c22f1790c --- /dev/null +++ b/tooling/lsp/src/trait_impl_method_stub_generator.rs @@ -0,0 +1,449 @@ +use std::collections::BTreeMap; + +use noirc_frontend::{ + ast::NoirTraitImpl, + graph::CrateId, + hir::{ + def_map::{CrateDefMap, ModuleId}, + type_check::generics::TraitGenerics, + }, + hir_def::{function::FuncMeta, stmt::HirPattern, traits::Trait}, + macros_api::{ModuleDefId, NodeInterner}, + node_interner::ReferenceId, + Kind, ResolvedGeneric, Type, TypeVariableKind, +}; + +use crate::modules::relative_module_id_path; + +pub(crate) struct TraitImplMethodStubGenerator<'a> { + name: &'a str, + func_meta: &'a FuncMeta, + trait_: &'a Trait, + noir_trait_impl: &'a NoirTraitImpl, + interner: &'a NodeInterner, + def_maps: &'a BTreeMap, + module_id: ModuleId, + indent: usize, + body: Option, + string: String, +} + +impl<'a> TraitImplMethodStubGenerator<'a> { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + name: &'a str, + func_meta: &'a FuncMeta, + trait_: &'a Trait, + noir_trait_impl: &'a NoirTraitImpl, + interner: &'a NodeInterner, + def_maps: &'a BTreeMap, + module_id: ModuleId, + indent: usize, + ) -> Self { + Self { + name, + func_meta, + trait_, + noir_trait_impl, + interner, + def_maps, + module_id, + indent, + body: None, + string: String::new(), + } + } + + /// Sets the body to include in the stub method. By default an empty body will be generated. + pub(crate) fn set_body(&mut self, body: String) { + self.body = Some(body); + } + + pub(crate) fn generate(&mut self) -> String { + let indent_string = " ".repeat(self.indent); + + self.string.push_str(&indent_string); + self.string.push_str("fn "); + self.string.push_str(self.name); + self.append_resolved_generics(&self.func_meta.direct_generics); + self.string.push('('); + for (index, (pattern, typ, _visibility)) in self.func_meta.parameters.iter().enumerate() { + if index > 0 { + self.string.push_str(", "); + } + if self.append_pattern(pattern) { + self.string.push_str(": "); + self.append_type(typ); + } + } + self.string.push(')'); + + let return_type = self.func_meta.return_type(); + if return_type != &Type::Unit { + self.string.push_str(" -> "); + self.append_type(return_type); + } + + if !self.func_meta.trait_constraints.is_empty() { + self.string.push_str(" where "); + for (index, constraint) in self.func_meta.trait_constraints.iter().enumerate() { + if index > 0 { + self.string.push_str(", "); + } + self.append_type(&constraint.typ); + self.string.push_str(": "); + let trait_ = self.interner.get_trait(constraint.trait_id); + self.string.push_str(&trait_.name.0.contents); + self.append_trait_generics(&constraint.trait_generics); + } + } + + self.string.push_str(" {\n"); + + if let Some(body) = &self.body { + let body_indent_string = " ".repeat(self.indent + 4); + self.string.push_str(&body_indent_string); + self.string.push_str(body); + self.string.push('\n'); + self.string.push_str(&indent_string); + } + + self.string.push_str("}\n"); + std::mem::take(&mut self.string) + } + + /// Appends a pattern and returns true if this was not the self type + fn append_pattern(&mut self, pattern: &HirPattern) -> bool { + match pattern { + HirPattern::Identifier(hir_ident) => { + let definition = self.interner.definition(hir_ident.id); + self.string.push_str(&definition.name); + &definition.name != "self" + } + HirPattern::Mutable(pattern, _) => { + self.string.push_str("mut "); + self.append_pattern(pattern) + } + HirPattern::Tuple(patterns, _) => { + self.string.push('('); + for (index, pattern) in patterns.iter().enumerate() { + if index > 0 { + self.string.push_str(", "); + } + self.append_pattern(pattern); + } + self.string.push(')'); + true + } + HirPattern::Struct(typ, patterns, _) => { + self.append_type(typ); + self.string.push_str(" { "); + for (index, (name, _pattern)) in patterns.iter().enumerate() { + if index > 0 { + self.string.push_str(", "); + } + self.string.push_str(&name.0.contents); + } + self.string.push_str(" }"); + true + } + } + } + + fn append_type(&mut self, typ: &Type) { + match typ { + Type::FieldElement => self.string.push_str("Field"), + Type::Array(n, e) => { + self.string.push('['); + self.append_type(e); + self.string.push_str("; "); + self.append_type(n); + self.string.push(']'); + } + Type::Slice(typ) => { + self.string.push('['); + self.append_type(typ); + self.string.push(']'); + } + Type::Tuple(types) => { + self.string.push('('); + for (index, typ) in types.iter().enumerate() { + if index > 0 { + self.string.push_str(", "); + } + self.append_type(typ); + } + self.string.push(')'); + } + Type::Struct(struct_type, generics) => { + let struct_type = struct_type.borrow(); + + let current_module_data = + &self.def_maps[&self.module_id.krate].modules()[self.module_id.local_id.0]; + + // Check if the struct type is already imported/visible in this module + let per_ns = current_module_data.find_name(&struct_type.name); + if let Some((module_def_id, _, _)) = per_ns.types { + if module_def_id == ModuleDefId::TypeId(struct_type.id) { + self.string.push_str(&struct_type.name.0.contents); + self.append_generics(generics); + return; + } + } + + let module_id = struct_type.id.module_id(); + let module_data = &self.def_maps[&module_id.krate].modules()[module_id.local_id.0]; + let parent_module_local_id = module_data.parent.unwrap(); + let parent_module_id = + ModuleId { krate: module_id.krate, local_id: parent_module_local_id }; + + let current_module_parent_id = current_module_data + .parent + .map(|parent| ModuleId { krate: self.module_id.krate, local_id: parent }); + + let relative_path = relative_module_id_path( + parent_module_id, + &self.module_id, + current_module_parent_id, + self.interner, + ); + + if !relative_path.is_empty() { + self.string.push_str(&relative_path); + self.string.push_str("::"); + } + self.string.push_str(&struct_type.name.0.contents); + self.append_generics(generics); + } + Type::Alias(type_alias, generics) => { + let type_alias = type_alias.borrow(); + + let current_module_data = + &self.def_maps[&self.module_id.krate].modules()[self.module_id.local_id.0]; + + // Check if the alias type is already imported/visible in this module + let per_ns = current_module_data.find_name(&type_alias.name); + if let Some((module_def_id, _, _)) = per_ns.types { + if module_def_id == ModuleDefId::TypeAliasId(type_alias.id) { + self.string.push_str(&type_alias.name.0.contents); + self.append_generics(generics); + return; + } + } + + let parent_module_id = + self.interner.reference_module(ReferenceId::Alias(type_alias.id)).unwrap(); + + let current_module_parent_id = current_module_data + .parent + .map(|parent| ModuleId { krate: self.module_id.krate, local_id: parent }); + + let relative_path = relative_module_id_path( + *parent_module_id, + &self.module_id, + current_module_parent_id, + self.interner, + ); + + if !relative_path.is_empty() { + self.string.push_str(&relative_path); + self.string.push_str("::"); + } + self.string.push_str(&type_alias.name.0.contents); + self.append_generics(generics); + } + Type::TraitAsType(trait_id, _, trait_generics) => { + let trait_ = self.interner.get_trait(*trait_id); + + let current_module_data = + &self.def_maps[&self.module_id.krate].modules()[self.module_id.local_id.0]; + + // Check if the trait type is already imported/visible in this module + let per_ns = current_module_data.find_name(&trait_.name); + if let Some((module_def_id, _, _)) = per_ns.types { + if module_def_id == ModuleDefId::TraitId(*trait_id) { + self.string.push_str(&trait_.name.0.contents); + self.append_trait_generics(trait_generics); + return; + } + } + + let parent_module_id = + self.interner.reference_module(ReferenceId::Trait(*trait_id)).unwrap(); + + let current_module_parent_id = current_module_data + .parent + .map(|parent| ModuleId { krate: self.module_id.krate, local_id: parent }); + + let relative_path = relative_module_id_path( + *parent_module_id, + &self.module_id, + current_module_parent_id, + self.interner, + ); + + if !relative_path.is_empty() { + self.string.push_str(&relative_path); + self.string.push_str("::"); + } + self.string.push_str(&trait_.name.0.contents); + self.append_trait_generics(trait_generics); + } + Type::TypeVariable(typevar, _) => { + if typevar.id() == self.trait_.self_type_typevar.id() { + self.string.push_str("Self"); + return; + } + + let generics = &self.trait_.generics; + if let Some(index) = + generics.iter().position(|generic| generic.type_var.id() == typevar.id()) + { + if let Some(typ) = self.noir_trait_impl.trait_generics.ordered_args.get(index) { + self.string.push_str(&typ.to_string()); + return; + } + } + + for associated_type in &self.trait_.associated_types { + if typevar.id() == associated_type.type_var.id() { + self.string.push_str("Self::"); + self.string.push_str(&associated_type.name); + return; + } + } + + for generic in &self.func_meta.direct_generics { + if typevar.id() == generic.type_var.id() { + self.string.push_str(&generic.name); + return; + } + } + + self.string.push_str("error"); + } + Type::NamedGeneric(typevar, _name, _kind) => { + self.append_type(&Type::TypeVariable(typevar.clone(), TypeVariableKind::Normal)); + } + Type::Function(args, ret, env, unconstrained) => { + if *unconstrained { + self.string.push_str("unconstrained "); + } + self.string.push_str("fn"); + + if let Type::Unit = **env { + } else { + self.string.push('['); + self.append_type(env); + self.string.push(']'); + } + + self.string.push('('); + for (index, arg) in args.iter().enumerate() { + if index > 0 { + self.string.push_str(", "); + } + self.append_type(arg); + } + self.string.push(')'); + + if let Type::Unit = **ret { + } else { + self.string.push_str(" -> "); + self.append_type(ret); + } + } + Type::MutableReference(typ) => { + self.string.push_str("&mut "); + self.append_type(typ); + } + Type::Forall(_, _) => { + panic!("Shouldn't get a Type::Forall"); + } + Type::InfixExpr(left, op, right) => { + self.append_type(left); + self.string.push(' '); + self.string.push_str(&op.to_string()); + self.string.push(' '); + self.append_type(right); + } + Type::Constant(_) + | Type::Integer(_, _) + | Type::Bool + | Type::String(_) + | Type::FmtString(_, _) + | Type::Unit + | Type::Quoted(_) + | Type::Error => self.string.push_str(&typ.to_string()), + } + } + + fn append_generics(&mut self, generics: &[Type]) { + if generics.is_empty() { + return; + } + + self.string.push('<'); + for (index, typ) in generics.iter().enumerate() { + if index > 0 { + self.string.push_str(", "); + } + self.append_type(typ); + } + self.string.push('>'); + } + + fn append_trait_generics(&mut self, generics: &TraitGenerics) { + if generics.named.is_empty() && generics.ordered.is_empty() { + return; + } + + let mut index = 0; + + self.string.push('<'); + for generic in &generics.ordered { + if index > 0 { + self.string.push_str(", "); + } + self.append_type(generic); + index += 1; + } + for named_type in &generics.named { + if index > 0 { + self.string.push_str(", "); + } + self.string.push_str(&named_type.name.0.contents); + self.string.push_str(" = "); + self.append_type(&named_type.typ); + index += 1; + } + self.string.push('>'); + } + + fn append_resolved_generics(&mut self, generics: &[ResolvedGeneric]) { + if generics.is_empty() { + return; + } + + self.string.push('<'); + for (index, generic) in self.func_meta.direct_generics.iter().enumerate() { + if index > 0 { + self.string.push_str(", "); + } + self.append_resolved_generic(generic); + } + self.string.push('>'); + } + + fn append_resolved_generic(&mut self, generic: &ResolvedGeneric) { + match &generic.kind { + Kind::Normal => self.string.push_str(&generic.name), + Kind::Numeric(typ) => { + self.string.push_str("let "); + self.string.push_str(&generic.name); + self.string.push_str(": "); + self.append_type(typ); + } + } + } +}