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
This commit is contained in:
zc he 2025-02-06 20:49:13 +08:00 committed by GitHub
parent 0b2d1327d2
commit 164a089656
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 619 additions and 466 deletions

View File

@ -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<EngineState>,
@ -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<String> = 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<String> = 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 = &current_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 = &current_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;
}
}
};
}
}

View File

@ -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<Block>,
working_set: &'a StateWorkingSet,
f_special: &F,
) -> Vec<T>
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
{
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<T>
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
{
// 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<T>
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
{
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<T>
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
{
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<Vec<(Id, Span)>> {
) -> 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<PathMember> = 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<Vec<Span>> {
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<Span> {
ast_flat_map(ast, working_set, &|e| {
ast.flat_map(working_set, &|e| {
find_reference_by_id_in_expr(e, working_set, id)
})
}

View File

@ -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<Vec<InlayHint>> {
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<InlayHint> =
[lhs, op, rhs].into_iter().flat_map(|e| recur(e)).collect();
let mut hints: Vec<InlayHint> = [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<InlayHint> {
ast_flat_map(block, working_set, &|e| {
block.flat_map(working_set, &|e| {
extract_inlay_hints_from_expression(e, working_set, &offset, file)
})
}

View File

@ -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`].

View File

@ -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::*;

View File

@ -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<T> {
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<T>
where
F: Fn(&'a Expression) -> Option<Vec<T>>;
/// 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<T>
where
F: Fn(&'a Expression) -> FindMapResult<T>;
}
impl Traverse for Block {
fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
{
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<T>
where
F: Fn(&'a Expression) -> FindMapResult<T>,
{
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<T>
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
{
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<T>
where
F: Fn(&'a Expression) -> FindMapResult<T>,
{
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<T>
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
{
// 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<T>
where
F: Fn(&'a Expression) -> FindMapResult<T>,
{
// 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<T>
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
{
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<T>
where
F: Fn(&'a Expression) -> FindMapResult<T>,
{
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)))
}
}