extern/nushell
1
0
Fork 1
mirror of https://github.com/nushell/nushell.git synced 2025-05-29 14:21:45 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor(lsp): flat_map with mutable accumulator (#15567)

Browse Source 
# Description

Mainly performance improvement of lsp operations involving flat_map on
AST nodes.
Previous flat_map traversing is functional, which is a nice property to
have, but the heavy cost of vector collection on each tree node makes it
undesirable.

This PR mitigates the problem with a mutable accumulator.

# User-Facing Changes

Should be none.

# Tests + Formatting

# After Submitting
This commit is contained in:
zc he 2025-04-15 20:21:23 +08:00 committed by GitHub
parent 8c4d3eaa7e
commit e5f589ccdd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 188 additions and 194 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
crates/nu-lsp/src/ast.rs
View File

@ -46,6 +46,10 @@ fn command_name_span_from_call_head(
head_span: Span,
) -> Span {
let name = working_set.get_decl(decl_id).name();
// shortcut for most cases
if name.len() == head_span.end.saturating_sub(head_span.start) {
return head_span;
}
let head_content = working_set.get_span_contents(head_span);
let mut head_words = head_content.split(|c| *c == b' ').collect::<Vec<_>>();
let mut name_words = name.split(' ').collect::<Vec<_>>();
@ -500,38 +504,34 @@ fn find_reference_by_id_in_expr(
expr: &Expression,
working_set: &StateWorkingSet,
id: &Id,
) -> Option<Vec<Span>> {
let closure = |e| find_reference_by_id_in_expr(e, working_set, id);
) -> Vec<Span> {
match (&expr.expr, id) {
(Expr::Var(vid1), Id::Variable(vid2, _)) if *vid1 == *vid2 => Some(vec![Span::new(
(Expr::Var(vid1), Id::Variable(vid2, _)) if *vid1 == *vid2 => vec![Span::new(
// we want to exclude the `$` sign for renaming
expr.span.start.saturating_add(1),
expr.span.end,
)]),
(Expr::VarDecl(vid1), Id::Variable(vid2, _)) if *vid1 == *vid2 => Some(vec![expr.span]),
)],
(Expr::VarDecl(vid1), Id::Variable(vid2, _)) if *vid1 == *vid2 => vec![expr.span],
// also interested in `var_id` in call.arguments of `use` command
// and `module_id` in `module` command
(Expr::Call(call), _) => {
let mut occurs: Vec<Span> = call
.arguments
.iter()
.filter_map(|arg| arg.expr())
.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(command_name_span_from_call_head(
vec![command_name_span_from_call_head(
working_set,
call.decl_id,
call.head,
));
return Some(occurs);
)]
}
if let Some((_, span_found)) = try_find_id_in_misc(call, working_set, None, Some(id)) {
occurs.push(span_found);
// Check for misc matches (use, module, etc.)
else if let Some((_, span_found)) =
try_find_id_in_misc(call, working_set, None, Some(id))
{
vec![span_found]
} else {
vec![]
}
Some(occurs)
}
_ => None,
_ => vec![],
}
}
@ -540,7 +540,8 @@ pub(crate) fn find_reference_by_id(
working_set: &StateWorkingSet,
id: &Id,
) -> Vec<Span> {
ast.flat_map(working_set, &|e| {
find_reference_by_id_in_expr(e, working_set, id)
})
let mut results = Vec::new();
let closure = |e| find_reference_by_id_in_expr(e, working_set, id);
ast.flat_map(working_set, &closure, &mut results);
results
}

68
crates/nu-lsp/src/hints.rs
View File

@ -26,14 +26,9 @@ fn extract_inlay_hints_from_expression(
working_set: &StateWorkingSet,
offset: &usize,
file: &FullTextDocument,
) -> Option<Vec<InlayHint>> {
let closure = |e| extract_inlay_hints_from_expression(e, working_set, offset, file);
) -> Vec<InlayHint> {
match &expr.expr {
Expr::BinaryOp(lhs, op, rhs) => {
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);
@ -43,7 +38,7 @@ fn extract_inlay_hints_from_expression(
(_, "any") => type_lhs,
_ => type_lhs,
};
hints.push(InlayHint {
vec![InlayHint {
kind: Some(InlayHintKind::TYPE),
label: InlayHintLabel::String(format!(": {}", type_string)),
position,
@ -52,9 +47,10 @@ fn extract_inlay_hints_from_expression(
data: None,
padding_left: None,
padding_right: None,
})
}]
} else {
vec![]
}
Some(hints)
}
Expr::VarDecl(var_id) => {
let position = span_to_range(&expr.span, file, *offset).end;
@ -69,27 +65,30 @@ fn extract_inlay_hints_from_expression(
}))
.contains(':')
{
return Some(Vec::new());
return vec![];
}
let var = working_set.get_variable(*var_id);
let type_string = type_short_name(&var.ty);
Some(vec![
(InlayHint {
kind: Some(InlayHintKind::TYPE),
label: InlayHintLabel::String(format!(": {}", type_string)),
position,
text_edits: None,
tooltip: None,
data: None,
padding_left: None,
padding_right: None,
}),
])
vec![InlayHint {
kind: Some(InlayHintKind::TYPE),
label: InlayHintLabel::String(format!(": {}", type_string)),
position,
text_edits: None,
tooltip: None,
data: None,
padding_left: None,
padding_right: None,
}]
}
Expr::Call(call) => {
let decl = working_set.get_decl(call.decl_id);
// skip those defined outside of the project
working_set.get_block(decl.block_id()?).span?;
let Some(block_id) = decl.block_id() else {
return vec![];
};
if working_set.get_block(block_id).span.is_none() {
return vec![];
};
let signatures = decl.signature();
let signatures = [
signatures.required_positional,
@ -102,17 +101,11 @@ fn extract_inlay_hints_from_expression(
for arg in arguments {
match arg {
// skip the rest when spread/unknown arguments encountered
Argument::Spread(expr) | Argument::Unknown(expr) => {
hints.extend(expr.flat_map(working_set, &closure));
Argument::Spread(_) | Argument::Unknown(_) => {
sig_idx = signatures.len();
continue;
}
// skip current for flags
Argument::Named((_, _, Some(expr))) => {
hints.extend(expr.flat_map(working_set, &closure));
continue;
}
Argument::Positional(expr) => {
Argument::Positional(_) => {
if let Some(sig) = signatures.get(sig_idx) {
sig_idx += 1;
let position = span_to_range(&arg.span(), file, *offset).start;
@ -130,16 +123,16 @@ fn extract_inlay_hints_from_expression(
padding_right: None,
});
}
hints.extend(expr.flat_map(working_set, &closure));
}
// skip current for flags
_ => {
continue;
}
}
}
Some(hints)
hints
}
_ => None,
_ => vec![],
}
}
@ -154,9 +147,10 @@ impl LanguageServer {
offset: usize,
file: &FullTextDocument,
) -> Vec<InlayHint> {
block.flat_map(working_set, &|e| {
extract_inlay_hints_from_expression(e, working_set, &offset, file)
})
let closure = |e| extract_inlay_hints_from_expression(e, working_set, &offset, file);
let mut results = Vec::new();
block.flat_map(working_set, &closure, &mut results);
results
}
}

33
crates/nu-lsp/src/semantic_tokens.rs
View File

@ -19,32 +19,21 @@ use crate::{span_to_range, LanguageServer};
fn extract_semantic_tokens_from_expression(
expr: &Expression,
working_set: &StateWorkingSet,
) -> Option<Vec<Span>> {
let closure = |e| extract_semantic_tokens_from_expression(e, working_set);
) -> Vec<Span> {
match &expr.expr {
Expr::Call(call) => {
let command_name_bytes = working_set.get_span_contents(call.head);
let head_span = if command_name_bytes.contains(&b' ')
let command_name = working_set.get_decl(call.decl_id).name();
if command_name.contains(' ')
// Some keywords that are already highlighted properly, e.g. by tree-sitter-nu
&& !command_name_bytes.starts_with(b"export")
&& !command_name_bytes.starts_with(b"overlay")
&& !command_name.starts_with("export")
&& !command_name.starts_with("overlay")
{
vec![call.head]
} else {
vec![]
};
let spans = head_span
.into_iter()
.chain(
call.arguments
.iter()
.filter_map(|arg| arg.expr())
.flat_map(|e| e.flat_map(working_set, &closure)),
)
.collect();
Some(spans)
}
}
_ => None,
_ => vec![],
}
}
@ -67,14 +56,14 @@ impl LanguageServer {
offset: usize,
file: &FullTextDocument,
) -> Vec<SemanticToken> {
let spans = block.flat_map(working_set, &|e| {
extract_semantic_tokens_from_expression(e, working_set)
});
let mut results = Vec::new();
let closure = |e| extract_semantic_tokens_from_expression(e, working_set);
block.flat_map(working_set, &closure, &mut results);
let mut last_token_line = 0;
let mut last_token_char = 0;
let mut last_span = Span::unknown();
let mut tokens = vec![];
for sp in spans {
for sp in results {
let range = span_to_range(&sp, file, offset);
// shouldn't happen
if sp < last_span {

238
crates/nu-protocol/src/ast/traverse.rs
View File

@ -15,17 +15,19 @@ pub enum FindMapResult<T> {
/// 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
/// Generic function that do flat_map on an AST node.
/// Concatenates all recursive results on sub-expressions
/// into the `results` accumulator.
///
/// # Arguments
/// * `f` - function that overrides the default behavior
fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
/// * `f` - function that generates leaf elements
/// * `results` - accumulator
fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
where
F: Fn(&'a Expression) -> Option<Vec<T>>;
F: Fn(&'a Expression) -> Vec<T>;
/// Generic function that do find_map on an AST node
/// return the first Some
/// Generic function that do find_map on an AST node.
/// Return the first result found by applying `f` on sub-expressions.
///
/// # Arguments
/// * `f` - function that overrides the default behavior
@ -35,24 +37,18 @@ pub trait Traverse {
}
impl Traverse for Block {
fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
F: Fn(&'a Expression) -> 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()
for pipeline in self.pipelines.iter() {
for element in pipeline.elements.iter() {
element.expr.flat_map(working_set, f, results);
if let Some(redir) = &element.redirection {
redir.flat_map(working_set, f, results);
};
}
}
}
fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
@ -71,21 +67,19 @@ impl Traverse for Block {
}
impl Traverse for PipelineRedirection {
fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
F: Fn(&'a Expression) -> Vec<T>,
{
let recur = |expr: &'a Expression| expr.flat_map(working_set, f);
let mut recur = |expr: &'a Expression| expr.flat_map(working_set, f, results);
match self {
PipelineRedirection::Single { target, .. } => {
target.expr().map(recur).unwrap_or_default()
PipelineRedirection::Single { target, .. } => target.expr().map(recur),
PipelineRedirection::Separate { out, err } => {
out.expr().map(&mut recur);
err.expr().map(&mut recur)
}
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>
@ -94,9 +88,7 @@ impl Traverse for PipelineRedirection {
{
let recur = |expr: &'a Expression| expr.find_map(working_set, f);
match self {
PipelineRedirection::Single { target, .. } => {
target.expr().map(recur).unwrap_or_default()
}
PipelineRedirection::Single { target, .. } => target.expr().and_then(recur),
PipelineRedirection::Separate { out, err } => {
[out, err].iter().filter_map(|t| t.expr()).find_map(recur)
}
@ -105,87 +97,97 @@ impl Traverse for PipelineRedirection {
}
impl Traverse for Expression {
fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
F: Fn(&'a Expression) -> 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);
// leaf elements generated by `f` for this expression
results.extend(f(self));
let mut recur = |expr: &'a Expression| expr.flat_map(working_set, f, results);
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)
block.flat_map(working_set, f, results)
}
Expr::Range(range) => {
for sub_expr in [&range.from, &range.next, &range.to].into_iter().flatten() {
recur(sub_expr);
}
}
Expr::Call(call) => {
for arg in &call.arguments {
if let Some(sub_expr) = arg.expr() {
recur(sub_expr);
}
}
}
Expr::ExternalCall(head, args) => {
recur(head.as_ref());
for arg in args {
recur(arg.expr());
}
}
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::BinaryOp(lhs, op, rhs) => {
recur(lhs);
recur(op);
recur(rhs);
}
Expr::MatchBlock(matches) => {
for (pattern, expr) in matches {
pattern.flat_map(working_set, f, results);
expr.flat_map(working_set, f, results);
}
}
Expr::List(items) => {
for item in items {
match item {
ListItem::Item(expr) | ListItem::Spread(_, expr) => recur(expr),
}
}
}
Expr::Record(items) => {
for item in items {
match item {
RecordItem::Spread(_, expr) => recur(expr),
RecordItem::Pair(key, val) => {
recur(key);
recur(val);
}
}
}
}
Expr::Table(table) => {
for column in &table.columns {
recur(column);
}
for row in &table.rows {
for item in row {
recur(item);
}
}
}
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()
for item in vec {
recur(item);
}
}
Expr::AttributeBlock(ab) => {
for attr in &ab.attributes {
recur(&attr.expr);
}
recur(&ab.item);
}
Expr::AttributeBlock(ab) => ab
.attributes
.iter()
.flat_map(|attr| recur(&attr.expr))
.chain(recur(&ab.item))
.collect(),
_ => Vec::new(),
}
_ => (),
};
}
fn find_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Option<T>
@ -203,7 +205,9 @@ impl Traverse for Expression {
| Expr::Subexpression(block_id)
| Expr::Block(block_id)
| Expr::Closure(block_id) => {
let block = working_set.get_block(block_id.to_owned());
// Clone the block_id to create an owned value
let block_id = block_id.to_owned();
let block = working_set.get_block(block_id);
block.find_map(working_set, f)
}
Expr::Range(range) => [&range.from, &range.next, &range.to]
@ -253,25 +257,31 @@ impl Traverse for Expression {
}
impl Traverse for MatchPattern {
fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F) -> Vec<T>
fn flat_map<'a, T, F>(&'a self, working_set: &'a StateWorkingSet, f: &F, results: &mut Vec<T>)
where
F: Fn(&'a Expression) -> Option<Vec<T>>,
F: Fn(&'a Expression) -> 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);
let mut recur_pattern =
|pattern: &'a MatchPattern| pattern.flat_map(working_set, f, results);
match &self.pattern {
Pattern::Expression(expr) => recur(expr),
Pattern::Expression(expr) => expr.flat_map(working_set, f, results),
Pattern::List(patterns) | Pattern::Or(patterns) => {
patterns.iter().flat_map(recur_pattern).collect()
for pattern in patterns {
recur_pattern(pattern);
}
}
Pattern::Record(entries) => {
entries.iter().flat_map(|(_, p)| recur_pattern(p)).collect()
for (_, p) in entries {
recur_pattern(p);
}
}
_ => Vec::new(),
_ => (),
};
if let Some(g) = self.guard.as_ref() {
g.flat_map(working_set, f, results);
}
.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>