From 164a0896564eae74d80ecc0d04167af292a275eb Mon Sep 17 00:00:00 2001 From: zc he Date: Thu, 6 Feb 2025 20:49:13 +0800 Subject: [PATCH] refactor(completion): AST traverse to find the inner-most expression to complete (#14973) # Description As discussed [here](https://github.com/nushell/nushell/pull/14856#issuecomment-2623393017) and [here](https://github.com/nushell/nushell/discussions/14868). I feel this method is generally better. As for the new-parser, we can simply modify the implementation in `traverse.rs` to accommodate. Next, I'm gonna overhaul the `Completer` trait, so before it gets really messy, I' think this is the step to put this open for review so we can check if I'm on track. This PR closes #13897 (the `|` part) # User-Facing Changes # After Submitting --- crates/nu-cli/src/completions/completer.rs | 573 +++++++++++---------- crates/nu-lsp/src/ast.rs | 201 +------- crates/nu-lsp/src/hints.rs | 18 +- crates/nu-protocol/src/ast/call.rs | 9 + crates/nu-protocol/src/ast/mod.rs | 2 + crates/nu-protocol/src/ast/traverse.rs | 282 ++++++++++ 6 files changed, 619 insertions(+), 466 deletions(-) create mode 100644 crates/nu-protocol/src/ast/traverse.rs diff --git a/crates/nu-cli/src/completions/completer.rs b/crates/nu-cli/src/completions/completer.rs index 71125e04e9..9ab0cba566 100644 --- a/crates/nu-cli/src/completions/completer.rs +++ b/crates/nu-cli/src/completions/completer.rs @@ -5,8 +5,9 @@ use crate::completions::{ use log::debug; use nu_color_config::{color_record_to_nustyle, lookup_ansi_color_style}; use nu_engine::eval_block; -use nu_parser::{flatten_pipeline_element, parse, FlatShape}; +use nu_parser::{flatten_expression, parse, FlatShape}; use nu_protocol::{ + ast::{Expr, Expression, FindMapResult, Traverse}, debugger::WithoutDebug, engine::{Closure, EngineState, Stack, StateWorkingSet}, PipelineData, Span, Value, @@ -16,6 +17,51 @@ use std::{str, sync::Arc}; use super::base::{SemanticSuggestion, SuggestionKind}; +fn find_pipeline_element_by_position<'a>( + expr: &'a Expression, + working_set: &'a StateWorkingSet, + pos: usize, +) -> FindMapResult<&'a Expression> { + // skip the entire expression if the position is not in it + if !expr.span.contains(pos) { + return FindMapResult::Stop; + } + let closure = |expr: &'a Expression| find_pipeline_element_by_position(expr, working_set, pos); + match &expr.expr { + Expr::Call(call) => call + .arguments + .iter() + .find_map(|arg| arg.expr().and_then(|e| e.find_map(working_set, &closure))) + // if no inner call/external_call found, then this is the inner-most one + .or(Some(expr)) + .map(FindMapResult::Found) + .unwrap_or_default(), + // TODO: clear separation of internal/external completion logic + Expr::ExternalCall(head, arguments) => arguments + .iter() + .find_map(|arg| arg.expr().find_map(working_set, &closure)) + .or(head.as_ref().find_map(working_set, &closure)) + .or(Some(expr)) + .map(FindMapResult::Found) + .unwrap_or_default(), + // complete the operator + Expr::BinaryOp(lhs, _, rhs) => lhs + .find_map(working_set, &closure) + .or(rhs.find_map(working_set, &closure)) + .or(Some(expr)) + .map(FindMapResult::Found) + .unwrap_or_default(), + Expr::FullCellPath(fcp) => fcp + .head + .find_map(working_set, &closure) + .or(Some(expr)) + .map(FindMapResult::Found) + .unwrap_or_default(), + Expr::Var(_) => FindMapResult::Found(expr), + _ => FindMapResult::Continue, + } +} + #[derive(Clone)] pub struct NuCompleter { engine_state: Arc, @@ -140,139 +186,144 @@ impl NuCompleter { let config = self.engine_state.get_config(); - let outermost_block = parse(&mut working_set, Some("completer"), line.as_bytes(), false); + let block = parse(&mut working_set, Some("completer"), line.as_bytes(), false); + let Some(element_expression) = block.find_map(&working_set, &|expr: &Expression| { + find_pipeline_element_by_position(expr, &working_set, pos) + }) else { + return vec![]; + }; - // Try to get the innermost block parsed (by span) so that we consider the correct context/scope. - let target_block = working_set - .delta - .blocks - .iter() - .filter_map(|block| match block.span { - Some(span) if span.contains(pos) => Some((block, span)), - _ => None, - }) - .reduce(|prev, cur| { - // |(block, span), (block, span)| - match cur.1.start.cmp(&prev.1.start) { - core::cmp::Ordering::Greater => cur, - core::cmp::Ordering::Equal if cur.1.end < prev.1.end => cur, - _ => prev, + let flattened = flatten_expression(&working_set, element_expression); + let mut spans: Vec = vec![]; + + for (flat_idx, (span, shape)) in flattened.iter().enumerate() { + let is_passthrough_command = spans + .first() + .filter(|content| content.as_str() == "sudo" || content.as_str() == "doas") + .is_some(); + + // Read the current span to string + let current_span = working_set.get_span_contents(*span); + let current_span_str = String::from_utf8_lossy(current_span); + let is_last_span = span.contains(pos); + + // Skip the last 'a' as span item + if is_last_span { + let offset = pos - span.start; + if offset == 0 { + spans.push(String::new()) + } else { + let mut current_span_str = current_span_str.to_string(); + current_span_str.remove(offset); + spans.push(current_span_str); } - }) - .map(|(block, _)| block) - .unwrap_or(&outermost_block); + } else { + spans.push(current_span_str.to_string()); + } - for pipeline in &target_block.pipelines { - for pipeline_element in &pipeline.elements { - let flattened = flatten_pipeline_element(&working_set, pipeline_element); - let mut spans: Vec = vec![]; + // Complete based on the last span + if is_last_span { + // Context variables + let most_left_var = most_left_variable(flat_idx, &working_set, flattened.clone()); - for (flat_idx, flat) in flattened.iter().enumerate() { - let is_passthrough_command = spans - .first() - .filter(|content| content.as_str() == "sudo" || content.as_str() == "doas") - .is_some(); + // Create a new span + let new_span = Span::new(span.start, span.end - 1); - // Read the current span to string - let current_span = working_set.get_span_contents(flat.0); - let current_span_str = String::from_utf8_lossy(current_span); - let is_last_span = pos >= flat.0.start && pos < flat.0.end; + // Parses the prefix. Completion should look up to the cursor position, not after. + let index = pos - span.start; + let prefix = ¤t_span[..index]; - // Skip the last 'a' as span item - if is_last_span { - let offset = pos - flat.0.start; - if offset == 0 { - spans.push(String::new()) - } else { - let mut current_span_str = current_span_str.to_string(); - current_span_str.remove(offset); - spans.push(current_span_str); - } - } else { - spans.push(current_span_str.to_string()); + // Variables completion + if prefix.starts_with(b"$") || most_left_var.is_some() { + let mut variable_names_completer = + VariableCompletion::new(most_left_var.unwrap_or((vec![], vec![]))); + + let mut variable_completions = self.process_completion( + &mut variable_names_completer, + &working_set, + prefix, + new_span, + fake_offset, + pos, + ); + + let mut variable_operations_completer = + OperatorCompletion::new(element_expression.clone()); + + let mut variable_operations_completions = self.process_completion( + &mut variable_operations_completer, + &working_set, + prefix, + new_span, + fake_offset, + pos, + ); + + variable_completions.append(&mut variable_operations_completions); + return variable_completions; + } + + // Flags completion + if prefix.starts_with(b"-") { + // Try to complete flag internally + let mut completer = FlagCompletion::new(element_expression.clone()); + let result = self.process_completion( + &mut completer, + &working_set, + prefix, + new_span, + fake_offset, + pos, + ); + + if !result.is_empty() { + return result; } - // Complete based on the last span - if is_last_span { - // Context variables - let most_left_var = - most_left_variable(flat_idx, &working_set, flattened.clone()); - - // Create a new span - let new_span = Span::new(flat.0.start, flat.0.end - 1); - - // Parses the prefix. Completion should look up to the cursor position, not after. - let index = pos - flat.0.start; - let prefix = ¤t_span[..index]; - - // Variables completion - if prefix.starts_with(b"$") || most_left_var.is_some() { - let mut variable_names_completer = - VariableCompletion::new(most_left_var.unwrap_or((vec![], vec![]))); - - let mut variable_completions = self.process_completion( - &mut variable_names_completer, - &working_set, - prefix, - new_span, - fake_offset, - pos, - ); - - let mut variable_operations_completer = - OperatorCompletion::new(pipeline_element.expr.clone()); - - let mut variable_operations_completions = self.process_completion( - &mut variable_operations_completer, - &working_set, - prefix, - new_span, - fake_offset, - pos, - ); - - variable_completions.append(&mut variable_operations_completions); - return variable_completions; - } - - // Flags completion - if prefix.starts_with(b"-") { - // Try to complete flag internally - let mut completer = FlagCompletion::new(pipeline_element.expr.clone()); - let result = self.process_completion( - &mut completer, - &working_set, - prefix, - new_span, - fake_offset, - pos, - ); - - if !result.is_empty() { - return result; - } - - // We got no results for internal completion - // now we can check if external completer is set and use it - if let Some(closure) = config.completions.external.completer.as_ref() { - if let Some(external_result) = - self.external_completion(closure, &spans, fake_offset, new_span) - { - return external_result; - } - } - } - - // specially check if it is currently empty - always complete commands - if (is_passthrough_command && flat_idx == 1) - || (flat_idx == 0 && working_set.get_span_contents(new_span).is_empty()) + // We got no results for internal completion + // now we can check if external completer is set and use it + if let Some(closure) = config.completions.external.completer.as_ref() { + if let Some(external_result) = + self.external_completion(closure, &spans, fake_offset, new_span) { - let mut completer = CommandCompletion::new( - flattened.clone(), - // flat_idx, - FlatShape::String, - true, - ); + return external_result; + } + } + } + + // specially check if it is currently empty - always complete commands + if (is_passthrough_command && flat_idx == 1) + || (flat_idx == 0 && working_set.get_span_contents(new_span).is_empty()) + { + let mut completer = CommandCompletion::new( + flattened.clone(), + // flat_idx, + FlatShape::String, + true, + ); + return self.process_completion( + &mut completer, + &working_set, + prefix, + new_span, + fake_offset, + pos, + ); + } + + // Completions that depends on the previous expression (e.g: use, source-env) + if (is_passthrough_command && flat_idx > 1) || flat_idx > 0 { + if let Some(previous_expr) = flattened.get(flat_idx - 1) { + // Read the content for the previous expression + let prev_expr_str = working_set.get_span_contents(previous_expr.0).to_vec(); + + // Completion for .nu files + if prev_expr_str == b"use" + || prev_expr_str == b"overlay use" + || prev_expr_str == b"source-env" + { + let mut completer = DotNuCompletion::new(); + return self.process_completion( &mut completer, &working_set, @@ -281,164 +332,132 @@ impl NuCompleter { fake_offset, pos, ); - } + } else if prev_expr_str == b"ls" { + let mut completer = FileCompletion::new(); - // Completions that depends on the previous expression (e.g: use, source-env) - if (is_passthrough_command && flat_idx > 1) || flat_idx > 0 { - if let Some(previous_expr) = flattened.get(flat_idx - 1) { - // Read the content for the previous expression - let prev_expr_str = - working_set.get_span_contents(previous_expr.0).to_vec(); + return self.process_completion( + &mut completer, + &working_set, + prefix, + new_span, + fake_offset, + pos, + ); + } else if matches!( + previous_expr.1, + FlatShape::Float + | FlatShape::Int + | FlatShape::String + | FlatShape::List + | FlatShape::Bool + | FlatShape::Variable(_) + ) { + let mut completer = OperatorCompletion::new(element_expression.clone()); - // Completion for .nu files - if prev_expr_str == b"use" - || prev_expr_str == b"overlay use" - || prev_expr_str == b"source-env" - { - let mut completer = DotNuCompletion::new(); - - return self.process_completion( - &mut completer, - &working_set, - prefix, - new_span, - fake_offset, - pos, - ); - } else if prev_expr_str == b"ls" { - let mut completer = FileCompletion::new(); - - return self.process_completion( - &mut completer, - &working_set, - prefix, - new_span, - fake_offset, - pos, - ); - } else if matches!( - previous_expr.1, - FlatShape::Float - | FlatShape::Int - | FlatShape::String - | FlatShape::List - | FlatShape::Bool - | FlatShape::Variable(_) - ) { - let mut completer = - OperatorCompletion::new(pipeline_element.expr.clone()); - - let operator_suggestion = self.process_completion( - &mut completer, - &working_set, - prefix, - new_span, - fake_offset, - pos, - ); - if !operator_suggestion.is_empty() { - return operator_suggestion; - } - } + let operator_suggestion = self.process_completion( + &mut completer, + &working_set, + prefix, + new_span, + fake_offset, + pos, + ); + if !operator_suggestion.is_empty() { + return operator_suggestion; } } - - // Match other types - match &flat.1 { - FlatShape::Custom(decl_id) => { - let mut completer = CustomCompletion::new( - self.stack.clone(), - *decl_id, - initial_line, - FileCompletion::new(), - ); - - return self.process_completion( - &mut completer, - &working_set, - prefix, - new_span, - fake_offset, - pos, - ); - } - FlatShape::Directory => { - let mut completer = DirectoryCompletion::new(); - - return self.process_completion( - &mut completer, - &working_set, - prefix, - new_span, - fake_offset, - pos, - ); - } - FlatShape::Filepath | FlatShape::GlobPattern => { - let mut completer = FileCompletion::new(); - - return self.process_completion( - &mut completer, - &working_set, - prefix, - new_span, - fake_offset, - pos, - ); - } - flat_shape => { - let mut completer = CommandCompletion::new( - flattened.clone(), - // flat_idx, - flat_shape.clone(), - false, - ); - - let mut out: Vec<_> = self.process_completion( - &mut completer, - &working_set, - prefix, - new_span, - fake_offset, - pos, - ); - - if !out.is_empty() { - return out; - } - - // Try to complete using an external completer (if set) - if let Some(closure) = - config.completions.external.completer.as_ref() - { - if let Some(external_result) = self.external_completion( - closure, - &spans, - fake_offset, - new_span, - ) { - return external_result; - } - } - - // Check for file completion - let mut completer = FileCompletion::new(); - out = self.process_completion( - &mut completer, - &working_set, - prefix, - new_span, - fake_offset, - pos, - ); - - if !out.is_empty() { - return out; - } - } - }; } } + + // Match other types + match shape { + FlatShape::Custom(decl_id) => { + let mut completer = CustomCompletion::new( + self.stack.clone(), + *decl_id, + initial_line, + FileCompletion::new(), + ); + + return self.process_completion( + &mut completer, + &working_set, + prefix, + new_span, + fake_offset, + pos, + ); + } + FlatShape::Directory => { + let mut completer = DirectoryCompletion::new(); + + return self.process_completion( + &mut completer, + &working_set, + prefix, + new_span, + fake_offset, + pos, + ); + } + FlatShape::Filepath | FlatShape::GlobPattern => { + let mut completer = FileCompletion::new(); + + return self.process_completion( + &mut completer, + &working_set, + prefix, + new_span, + fake_offset, + pos, + ); + } + flat_shape => { + let mut completer = CommandCompletion::new( + flattened.clone(), + // flat_idx, + flat_shape.clone(), + false, + ); + + let mut out: Vec<_> = self.process_completion( + &mut completer, + &working_set, + prefix, + new_span, + fake_offset, + pos, + ); + + if !out.is_empty() { + return out; + } + + // Try to complete using an external completer (if set) + if let Some(closure) = config.completions.external.completer.as_ref() { + if let Some(external_result) = + self.external_completion(closure, &spans, fake_offset, new_span) + { + return external_result; + } + } + + // Check for file completion + let mut completer = FileCompletion::new(); + out = self.process_completion( + &mut completer, + &working_set, + prefix, + new_span, + fake_offset, + pos, + ); + + if !out.is_empty() { + return out; + } + } + }; } } diff --git a/crates/nu-lsp/src/ast.rs b/crates/nu-lsp/src/ast.rs index 511a25ca6b..775093c133 100644 --- a/crates/nu-lsp/src/ast.rs +++ b/crates/nu-lsp/src/ast.rs @@ -1,172 +1,11 @@ use crate::Id; use nu_protocol::{ - ast::{ - Argument, Block, Call, Expr, Expression, ExternalArgument, ListItem, MatchPattern, - PathMember, Pattern, PipelineRedirection, RecordItem, - }, + ast::{Argument, Block, Call, Expr, Expression, FindMapResult, ListItem, PathMember, Traverse}, engine::StateWorkingSet, Span, }; use std::sync::Arc; -/// similar to flatten_block, but allows extra map function -pub fn ast_flat_map<'a, T, F>( - ast: &'a Arc, - working_set: &'a StateWorkingSet, - f_special: &F, -) -> Vec -where - F: Fn(&'a Expression) -> Option>, -{ - ast.pipelines - .iter() - .flat_map(|pipeline| { - pipeline.elements.iter().flat_map(|element| { - expr_flat_map(&element.expr, working_set, f_special) - .into_iter() - .chain( - element - .redirection - .as_ref() - .map(|redir| redirect_flat_map(redir, working_set, f_special)) - .unwrap_or_default(), - ) - }) - }) - .collect() -} - -/// generic function that do flat_map on an expression -/// concats all recursive results on sub-expressions -/// -/// # Arguments -/// * `f_special` - function that overrides the default behavior -pub fn expr_flat_map<'a, T, F>( - expr: &'a Expression, - working_set: &'a StateWorkingSet, - f_special: &F, -) -> Vec -where - F: Fn(&'a Expression) -> Option>, -{ - // behavior overridden by f_special - if let Some(vec) = f_special(expr) { - return vec; - } - let recur = |expr| expr_flat_map(expr, working_set, f_special); - match &expr.expr { - Expr::RowCondition(block_id) - | Expr::Subexpression(block_id) - | Expr::Block(block_id) - | Expr::Closure(block_id) => { - let block = working_set.get_block(block_id.to_owned()); - ast_flat_map(block, working_set, f_special) - } - Expr::Range(range) => [&range.from, &range.next, &range.to] - .iter() - .filter_map(|e| e.as_ref()) - .flat_map(recur) - .collect(), - Expr::Call(call) => call - .arguments - .iter() - .filter_map(|arg| arg.expr()) - .flat_map(recur) - .collect(), - Expr::ExternalCall(head, args) => recur(head) - .into_iter() - .chain(args.iter().flat_map(|arg| match arg { - ExternalArgument::Regular(e) | ExternalArgument::Spread(e) => recur(e), - })) - .collect(), - Expr::UnaryNot(expr) | Expr::Collect(_, expr) => recur(expr), - Expr::BinaryOp(lhs, op, rhs) => recur(lhs) - .into_iter() - .chain(recur(op)) - .chain(recur(rhs)) - .collect(), - Expr::MatchBlock(matches) => matches - .iter() - .flat_map(|(pattern, expr)| { - match_pattern_flat_map(pattern, working_set, f_special) - .into_iter() - .chain(recur(expr)) - }) - .collect(), - Expr::List(items) => items - .iter() - .flat_map(|item| match item { - ListItem::Item(expr) | ListItem::Spread(_, expr) => recur(expr), - }) - .collect(), - Expr::Record(items) => items - .iter() - .flat_map(|item| match item { - RecordItem::Spread(_, expr) => recur(expr), - RecordItem::Pair(key, val) => [key, val].into_iter().flat_map(recur).collect(), - }) - .collect(), - Expr::Table(table) => table - .columns - .iter() - .flat_map(recur) - .chain(table.rows.iter().flat_map(|row| row.iter().flat_map(recur))) - .collect(), - Expr::ValueWithUnit(vu) => recur(&vu.expr), - Expr::FullCellPath(fcp) => recur(&fcp.head), - Expr::Keyword(kw) => recur(&kw.expr), - Expr::StringInterpolation(vec) | Expr::GlobInterpolation(vec, _) => { - vec.iter().flat_map(recur).collect() - } - - _ => Vec::new(), - } -} - -/// flat_map on match patterns -fn match_pattern_flat_map<'a, T, F>( - pattern: &'a MatchPattern, - working_set: &'a StateWorkingSet, - f_special: &F, -) -> Vec -where - F: Fn(&'a Expression) -> Option>, -{ - let recur = |expr| expr_flat_map(expr, working_set, f_special); - let recur_match = |p| match_pattern_flat_map(p, working_set, f_special); - match &pattern.pattern { - Pattern::Expression(expr) => recur(expr), - Pattern::List(patterns) | Pattern::Or(patterns) => { - patterns.iter().flat_map(recur_match).collect() - } - Pattern::Record(entries) => entries.iter().flat_map(|(_, p)| recur_match(p)).collect(), - _ => Vec::new(), - } - .into_iter() - .chain(pattern.guard.as_ref().map(|g| recur(g)).unwrap_or_default()) - .collect() -} - -/// flat_map on redirections -fn redirect_flat_map<'a, T, F>( - redir: &'a PipelineRedirection, - working_set: &'a StateWorkingSet, - f_special: &F, -) -> Vec -where - F: Fn(&'a Expression) -> Option>, -{ - let recur = |expr| expr_flat_map(expr, working_set, f_special); - match redir { - PipelineRedirection::Single { target, .. } => target.expr().map(recur).unwrap_or_default(), - PipelineRedirection::Separate { out, err } => [out, err] - .iter() - .filter_map(|t| t.expr()) - .flat_map(recur) - .collect(), - } -} - /// Adjust span if quoted fn strip_quotes(span: Span, working_set: &StateWorkingSet) -> Span { let text = String::from_utf8_lossy(working_set.get_span_contents(span)); @@ -436,22 +275,22 @@ fn find_id_in_expr( expr: &Expression, working_set: &StateWorkingSet, location: &usize, -) -> Option> { +) -> FindMapResult<(Id, Span)> { // skip the entire expression if the location is not in it if !expr.span.contains(*location) { - return Some(Vec::new()); + return FindMapResult::Stop; } let span = expr.span; match &expr.expr { - Expr::VarDecl(var_id) => Some(vec![(Id::Variable(*var_id), span)]), + Expr::VarDecl(var_id) => FindMapResult::Found((Id::Variable(*var_id), span)), // trim leading `$` sign - Expr::Var(var_id) => Some(vec![( + Expr::Var(var_id) => FindMapResult::Found(( Id::Variable(*var_id), Span::new(span.start.saturating_add(1), span.end), - )]), + )), Expr::Call(call) => { if call.head.contains(*location) { - Some(vec![(Id::Declaration(call.decl_id), call.head)]) + FindMapResult::Found((Id::Declaration(call.decl_id), call.head)) } else { try_find_id_in_def(call, working_set, Some(location), None) .or(try_find_id_in_mod(call, working_set, Some(location), None)) @@ -462,19 +301,20 @@ fn find_id_in_expr( Some(location), None, )) - .map(|p| vec![p]) + .map(FindMapResult::Found) + .unwrap_or_default() } } Expr::FullCellPath(fcp) => { if fcp.head.span.contains(*location) { - None + FindMapResult::Continue } else { let Expression { expr: Expr::Var(var_id), .. } = fcp.head else { - return None; + return FindMapResult::Continue; }; let tail: Vec = fcp .tail @@ -482,11 +322,13 @@ fn find_id_in_expr( .into_iter() .take_while(|pm| pm.span().start <= *location) .collect(); - let span = tail.last()?.span(); - Some(vec![(Id::CellPath(var_id, tail), span)]) + let Some(span) = tail.last().map(|pm| pm.span()) else { + return FindMapResult::Stop; + }; + FindMapResult::Found((Id::CellPath(var_id, tail), span)) } } - Expr::Overlay(Some(module_id)) => Some(vec![(Id::Module(*module_id), span)]), + Expr::Overlay(Some(module_id)) => FindMapResult::Found((Id::Module(*module_id), span)), // terminal value expressions Expr::Bool(_) | Expr::Binary(_) @@ -500,8 +342,8 @@ fn find_id_in_expr( | Expr::Nothing | Expr::RawString(_) | Expr::Signature(_) - | Expr::String(_) => Some(vec![(Id::Value(expr.ty.clone()), span)]), - _ => None, + | Expr::String(_) => FindMapResult::Found((Id::Value(expr.ty.clone()), span)), + _ => FindMapResult::Continue, } } @@ -512,7 +354,7 @@ pub(crate) fn find_id( location: &usize, ) -> Option<(Id, Span)> { let closure = |e| find_id_in_expr(e, working_set, location); - ast_flat_map(ast, working_set, &closure).first().cloned() + ast.find_map(working_set, &closure) } fn find_reference_by_id_in_expr( @@ -521,7 +363,6 @@ fn find_reference_by_id_in_expr( id: &Id, ) -> Option> { let closure = |e| find_reference_by_id_in_expr(e, working_set, id); - let recur = |expr| expr_flat_map(expr, working_set, &closure); match (&expr.expr, id) { (Expr::Var(vid1), Id::Variable(vid2)) if *vid1 == *vid2 => Some(vec![Span::new( // we want to exclude the `$` sign for renaming @@ -536,7 +377,7 @@ fn find_reference_by_id_in_expr( .arguments .iter() .filter_map(|arg| arg.expr()) - .flat_map(recur) + .flat_map(|e| e.flat_map(working_set, &closure)) .collect(); if matches!(id, Id::Declaration(decl_id) if call.decl_id == *decl_id) { occurs.push(call.head); @@ -560,7 +401,7 @@ pub(crate) fn find_reference_by_id( working_set: &StateWorkingSet, id: &Id, ) -> Vec { - ast_flat_map(ast, working_set, &|e| { + ast.flat_map(working_set, &|e| { find_reference_by_id_in_expr(e, working_set, id) }) } diff --git a/crates/nu-lsp/src/hints.rs b/crates/nu-lsp/src/hints.rs index 96c191c7db..581a2e70b0 100644 --- a/crates/nu-lsp/src/hints.rs +++ b/crates/nu-lsp/src/hints.rs @@ -1,4 +1,3 @@ -use crate::ast::{ast_flat_map, expr_flat_map}; use crate::{span_to_range, LanguageServer}; use lsp_textdocument::FullTextDocument; use lsp_types::{ @@ -6,7 +5,7 @@ use lsp_types::{ MarkupKind, Position, Range, }; use nu_protocol::{ - ast::{Argument, Block, Expr, Expression, Operator}, + ast::{Argument, Block, Expr, Expression, Operator, Traverse}, engine::StateWorkingSet, Type, }; @@ -29,11 +28,12 @@ fn extract_inlay_hints_from_expression( file: &FullTextDocument, ) -> Option> { let closure = |e| extract_inlay_hints_from_expression(e, working_set, offset, file); - let recur = |expr| expr_flat_map(expr, working_set, &closure); match &expr.expr { Expr::BinaryOp(lhs, op, rhs) => { - let mut hints: Vec = - [lhs, op, rhs].into_iter().flat_map(|e| recur(e)).collect(); + let mut hints: Vec = [lhs, op, rhs] + .into_iter() + .flat_map(|e| e.flat_map(working_set, &closure)) + .collect(); if let Expr::Operator(Operator::Assignment(_)) = op.expr { let position = span_to_range(&lhs.span, file, *offset).end; let type_rhs = type_short_name(&rhs.ty); @@ -103,13 +103,13 @@ fn extract_inlay_hints_from_expression( match arg { // skip the rest when spread/unknown arguments encountered Argument::Spread(expr) | Argument::Unknown(expr) => { - hints.extend(recur(expr)); + hints.extend(expr.flat_map(working_set, &closure)); sig_idx = signatures.len(); continue; } // skip current for flags Argument::Named((_, _, Some(expr))) => { - hints.extend(recur(expr)); + hints.extend(expr.flat_map(working_set, &closure)); continue; } Argument::Positional(expr) => { @@ -130,7 +130,7 @@ fn extract_inlay_hints_from_expression( padding_right: None, }); } - hints.extend(recur(expr)); + hints.extend(expr.flat_map(working_set, &closure)); } _ => { continue; @@ -154,7 +154,7 @@ impl LanguageServer { offset: usize, file: &FullTextDocument, ) -> Vec { - ast_flat_map(block, working_set, &|e| { + block.flat_map(working_set, &|e| { extract_inlay_hints_from_expression(e, working_set, &offset, file) }) } diff --git a/crates/nu-protocol/src/ast/call.rs b/crates/nu-protocol/src/ast/call.rs index 0c42583c5c..7a0d2b0cb4 100644 --- a/crates/nu-protocol/src/ast/call.rs +++ b/crates/nu-protocol/src/ast/call.rs @@ -76,6 +76,15 @@ pub enum ExternalArgument { Spread(Expression), } +impl ExternalArgument { + pub fn expr(&self) -> &Expression { + match self { + ExternalArgument::Regular(expr) => expr, + ExternalArgument::Spread(expr) => expr, + } + } +} + /// Parsed call of a `Command` /// /// As we also implement some internal keywords in terms of the `Command` trait, this type stores the passed arguments as [`Expression`]. diff --git a/crates/nu-protocol/src/ast/mod.rs b/crates/nu-protocol/src/ast/mod.rs index cb09168805..0138431ec6 100644 --- a/crates/nu-protocol/src/ast/mod.rs +++ b/crates/nu-protocol/src/ast/mod.rs @@ -11,6 +11,7 @@ mod operator; mod pipeline; mod range; mod table; +mod traverse; mod unit; mod value_with_unit; @@ -26,5 +27,6 @@ pub use operator::*; pub use pipeline::*; pub use range::*; pub use table::Table; +pub use traverse::*; pub use unit::*; pub use value_with_unit::*; diff --git a/crates/nu-protocol/src/ast/traverse.rs b/crates/nu-protocol/src/ast/traverse.rs new file mode 100644 index 0000000000..2c9e6dbaa1 --- /dev/null +++ b/crates/nu-protocol/src/ast/traverse.rs @@ -0,0 +1,282 @@ +use crate::engine::StateWorkingSet; + +use super::{ + Block, Expr, Expression, ListItem, MatchPattern, Pattern, PipelineRedirection, RecordItem, +}; + +/// Result of find_map closure +#[derive(Default)] +pub enum FindMapResult { + Found(T), + #[default] + Continue, + Stop, +} + +/// Trait for traversing the AST +pub trait Traverse { + /// Generic function that do flat_map on an AST node + /// concatenates all recursive results on sub-expressions + /// + /// # Arguments + /// * `f` - function that overrides the default behavior + fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec + where + F: Fn(&'a Expression) -> Option>; + + /// Generic function that do find_map on an AST node + /// return the first Some + /// + /// # Arguments + /// * `f` - function that overrides the default behavior + fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option + where + F: Fn(&'a Expression) -> FindMapResult; +} + +impl Traverse for Block { + fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec + where + F: Fn(&'a Expression) -> Option>, + { + self.pipelines + .iter() + .flat_map(|pipeline| { + pipeline.elements.iter().flat_map(|element| { + element.expr.flat_map(working_set, f).into_iter().chain( + element + .redirection + .as_ref() + .map(|redir| redir.flat_map(working_set, f)) + .unwrap_or_default(), + ) + }) + }) + .collect() + } + + fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option + where + F: Fn(&'a Expression) -> FindMapResult, + { + self.pipelines.iter().find_map(|pipeline| { + pipeline.elements.iter().find_map(|element| { + element.expr.find_map(working_set, f).or(element + .redirection + .as_ref() + .and_then(|redir| redir.find_map(working_set, f))) + }) + }) + } +} + +impl Traverse for PipelineRedirection { + fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec + where + F: Fn(&'a Expression) -> Option>, + { + let recur = |expr: &'a Expression| expr.flat_map(working_set, f); + match self { + PipelineRedirection::Single { target, .. } => { + target.expr().map(recur).unwrap_or_default() + } + PipelineRedirection::Separate { out, err } => [out, err] + .iter() + .filter_map(|t| t.expr()) + .flat_map(recur) + .collect(), + } + } + + fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option + where + F: Fn(&'a Expression) -> FindMapResult, + { + let recur = |expr: &'a Expression| expr.find_map(working_set, f); + match self { + PipelineRedirection::Single { target, .. } => { + target.expr().map(recur).unwrap_or_default() + } + PipelineRedirection::Separate { out, err } => { + [out, err].iter().filter_map(|t| t.expr()).find_map(recur) + } + } + } +} + +impl Traverse for Expression { + fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec + where + F: Fn(&'a Expression) -> Option>, + { + // behavior overridden by f + if let Some(vec) = f(self) { + return vec; + } + let recur = |expr: &'a Expression| expr.flat_map(working_set, f); + match &self.expr { + Expr::RowCondition(block_id) + | Expr::Subexpression(block_id) + | Expr::Block(block_id) + | Expr::Closure(block_id) => { + let block = working_set.get_block(block_id.to_owned()); + block.flat_map(working_set, f) + } + Expr::Range(range) => [&range.from, &range.next, &range.to] + .iter() + .filter_map(|e| e.as_ref()) + .flat_map(recur) + .collect(), + Expr::Call(call) => call + .arguments + .iter() + .filter_map(|arg| arg.expr()) + .flat_map(recur) + .collect(), + Expr::ExternalCall(head, args) => recur(head.as_ref()) + .into_iter() + .chain(args.iter().flat_map(|arg| recur(arg.expr()))) + .collect(), + Expr::UnaryNot(expr) | Expr::Collect(_, expr) => recur(expr.as_ref()), + Expr::BinaryOp(lhs, op, rhs) => recur(lhs) + .into_iter() + .chain(recur(op)) + .chain(recur(rhs)) + .collect(), + Expr::MatchBlock(matches) => matches + .iter() + .flat_map(|(pattern, expr)| { + pattern + .flat_map(working_set, f) + .into_iter() + .chain(recur(expr)) + }) + .collect(), + Expr::List(items) => items + .iter() + .flat_map(|item| match item { + ListItem::Item(expr) | ListItem::Spread(_, expr) => recur(expr), + }) + .collect(), + Expr::Record(items) => items + .iter() + .flat_map(|item| match item { + RecordItem::Spread(_, expr) => recur(expr), + RecordItem::Pair(key, val) => [key, val].into_iter().flat_map(recur).collect(), + }) + .collect(), + Expr::Table(table) => table + .columns + .iter() + .flat_map(recur) + .chain(table.rows.iter().flat_map(|row| row.iter().flat_map(recur))) + .collect(), + Expr::ValueWithUnit(vu) => recur(&vu.expr), + Expr::FullCellPath(fcp) => recur(&fcp.head), + Expr::Keyword(kw) => recur(&kw.expr), + Expr::StringInterpolation(vec) | Expr::GlobInterpolation(vec, _) => { + vec.iter().flat_map(recur).collect() + } + + _ => Vec::new(), + } + } + + fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option + where + F: Fn(&'a Expression) -> FindMapResult, + { + // behavior overridden by f + match f(self) { + FindMapResult::Found(t) => Some(t), + FindMapResult::Stop => None, + FindMapResult::Continue => { + let recur = |expr: &'a Expression| expr.find_map(working_set, f); + match &self.expr { + Expr::RowCondition(block_id) + | Expr::Subexpression(block_id) + | Expr::Block(block_id) + | Expr::Closure(block_id) => { + let block = working_set.get_block(block_id.to_owned()); + block.find_map(working_set, f) + } + Expr::Range(range) => [&range.from, &range.next, &range.to] + .iter() + .find_map(|e| e.as_ref().and_then(recur)), + Expr::Call(call) => call + .arguments + .iter() + .find_map(|arg| arg.expr().and_then(recur)), + Expr::ExternalCall(head, args) => { + recur(head.as_ref()).or(args.iter().find_map(|arg| recur(arg.expr()))) + } + Expr::UnaryNot(expr) | Expr::Collect(_, expr) => recur(expr.as_ref()), + Expr::BinaryOp(lhs, op, rhs) => recur(lhs).or(recur(op)).or(recur(rhs)), + Expr::MatchBlock(matches) => matches.iter().find_map(|(pattern, expr)| { + pattern.find_map(working_set, f).or(recur(expr)) + }), + Expr::List(items) => items.iter().find_map(|item| match item { + ListItem::Item(expr) | ListItem::Spread(_, expr) => recur(expr), + }), + Expr::Record(items) => items.iter().find_map(|item| match item { + RecordItem::Spread(_, expr) => recur(expr), + RecordItem::Pair(key, val) => [key, val].into_iter().find_map(recur), + }), + Expr::Table(table) => table + .columns + .iter() + .find_map(recur) + .or(table.rows.iter().find_map(|row| row.iter().find_map(recur))), + Expr::ValueWithUnit(vu) => recur(&vu.expr), + Expr::FullCellPath(fcp) => recur(&fcp.head), + Expr::Keyword(kw) => recur(&kw.expr), + Expr::StringInterpolation(vec) | Expr::GlobInterpolation(vec, _) => { + vec.iter().find_map(recur) + } + + _ => None, + } + } + } + } +} + +impl Traverse for MatchPattern { + fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec + where + F: Fn(&'a Expression) -> Option>, + { + let recur = |expr: &'a Expression| expr.flat_map(working_set, f); + let recur_pattern = |pattern: &'a MatchPattern| pattern.flat_map(working_set, f); + match &self.pattern { + Pattern::Expression(expr) => recur(expr), + Pattern::List(patterns) | Pattern::Or(patterns) => { + patterns.iter().flat_map(recur_pattern).collect() + } + Pattern::Record(entries) => { + entries.iter().flat_map(|(_, p)| recur_pattern(p)).collect() + } + _ => Vec::new(), + } + .into_iter() + .chain(self.guard.as_ref().map(|g| recur(g)).unwrap_or_default()) + .collect() + } + + fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option + where + F: Fn(&'a Expression) -> FindMapResult, + { + let recur = |expr: &'a Expression| expr.find_map(working_set, f); + let recur_pattern = |pattern: &'a MatchPattern| pattern.find_map(working_set, f); + match &self.pattern { + Pattern::Expression(expr) => recur(expr), + Pattern::List(patterns) | Pattern::Or(patterns) => { + patterns.iter().find_map(recur_pattern) + } + Pattern::Record(entries) => entries.iter().find_map(|(_, p)| recur_pattern(p)), + _ => None, + } + .or(self.guard.as_ref().and_then(|g| recur(g))) + } +}