@@ -192,13 +192,15 @@ pub fn indent_level_for_line(line: RopeSlice, tab_width: usize) -> usize {
192192
193193/// Computes for node and all ancestors whether they are the first node on their line.
194194/// The first entry in the return value represents the root node, the last one the node itself
195−fn get_first_in_line(mut node: Node, byte_pos: usize, new_line: bool) -> Vec<bool> {
195+fn get_first_in_line(mut node: Node, new_line_byte_pos: Option<usize>) -> Vec<bool> {
196196 let mut first_in_line = Vec::new();
197197 loop {
198198 if let Some(prev) = node.prev_sibling() {
199199 // If we insert a new line, the first node at/after the cursor is considered to be the first in its line
200200 let first = prev.end_position().row != node.start_position().row
201− || (new_line && node.start_byte() >= byte_pos && prev.start_byte() < byte_pos);
201+ || new_line_byte_pos.map_or(false, |byte_pos| {
202+ node.start_byte() >= byte_pos && prev.start_byte() < byte_pos
203+ });
202204 first_in_line.push(Some(first));
203205 } else {
204206 // Nodes that have no previous siblings are first in their line if and only if their parent is
@@ -298,8 +300,21 @@ enum IndentScope {
298300 Tail,
299301}
300302
301−/// Execute the indent query.
302−/// Returns for each node (identified by its id) a list of indent captures for that node.
303+/// A capture from the indent query which does not define an indent but extends
304+/// the range of a node. This is used before the indent is calculated.
305+enum ExtendCapture {
306+ ExtendIndented,
307+ StopExtend,
308+}
309+
310+/// The result of running a tree-sitter indent query. This stores for
311+/// each node (identified by its ID) the relevant captures (already filtered
312+/// by predicates).
313+struct IndentQueryResult {
314+ indent_captures: HashMap<usize, Vec<IndentCapture>>,
315+ extend_captures: HashMap<usize, Vec<ExtendCapture>>,
316+}
317+
303318fn query_indents(
304319 query: &Query,
305320 syntax: &Syntax,
@@ -309,8 +324,9 @@ fn query_indents(
309324 // Position of the (optional) newly inserted line break.
310325 // Given as (line, byte_pos)
311326 new_line_break: Option<(usize, usize)>,
312−) -> HashMap<usize, Vec<IndentCapture>> {
327+) -> IndentQueryResult {
313328 let mut indent_captures: HashMap<usize, Vec<IndentCapture>> = HashMap::new();
329+ let mut extend_captures: HashMap<usize, Vec<ExtendCapture>> = HashMap::new();
314330 cursor.set_byte_range(range);
315331 // Iterate over all captures from the query
316332 for m in cursor.matches(query, syntax.tree().root_node(), RopeProvider(text)) {
@@ -374,10 +390,24 @@ fn query_indents(
374390 continue;
375391 }
376392 for capture in m.captures {
377− let capture_type = query.capture_names()[capture.index as usize].as_str();
378− let capture_type = match capture_type {
393+ let capture_name = query.capture_names()[capture.index as usize].as_str();
394+ let capture_type = match capture_name {
379395 "indent" => IndentCaptureType::Indent,
380396 "outdent" => IndentCaptureType::Outdent,
397+ "extend-indented" => {
398+ extend_captures
399+ .entry(capture.node.id())
400+ .or_insert_with(|| Vec::with_capacity(1))
401+ .push(ExtendCapture::ExtendIndented);
402+ continue;
403+ }
404+ "stop-extend" => {
405+ extend_captures
406+ .entry(capture.node.id())
407+ .or_insert_with(|| Vec::with_capacity(1))
408+ .push(ExtendCapture::StopExtend);
409+ continue;
410+ }
381411 _ => {
382412 // Ignore any unknown captures (these may be needed for predicates such as #match?)
383413 continue;
@@ -420,7 +450,10 @@ fn query_indents(
420450 .push(indent_capture);
421451 }
422452 }
423− indent_captures
453+ IndentQueryResult {
454+ indent_captures,
455+ extend_captures,
456+ }
424457}
425458
426459/// Use the syntax tree to determine the indentation for a given position.
@@ -463,36 +496,114 @@ pub fn treesitter_indent_for_pos(
463496 query: &Query,
464497 syntax: &Syntax,
465498 indent_style: &IndentStyle,
499+ tab_width: usize,
466500 text: RopeSlice,
467501 line: usize,
468502 pos: usize,
469503 new_line: bool,
470504) -> Option<String> {
471505 let byte_pos = text.char_to_byte(pos);
506+ // The innermost tree-sitter node which is considered for the indent
507+ // computation. It may change if some predeceding node is extended
472508 let mut node = syntax
473509 .tree()
474510 .root_node()
475511 .descendant_for_byte_range(byte_pos, byte_pos)?;
476− let mut first_in_line = get_first_in_line(node, byte_pos, new_line);
477− let new_line_break = if new_line {
478− Some((line, byte_pos))
479− } else {
480− None
481− };
482− let query_result = crate::syntax::PARSER.with(|ts_parser| {
512+ let (query_result, prev_child) = crate::syntax::PARSER.with(|ts_parser| {
483513 let mut ts_parser = ts_parser.borrow_mut();
484514 let mut cursor = ts_parser.cursors.pop().unwrap_or_else(QueryCursor::new);
515+ // The query range should intersect with all nodes directly preceding
516+ // the cursor in case one of them is extended.
517+ // prev_child is the deepest such node.
518+ let (query_range, prev_child) = {
519+ // TODO Is there some way we can reuse this cursor?
520+ let mut tree_cursor = node.walk();
521+ let mut prev_child = None;
522+ for child in node.children(&mut tree_cursor) {
523+ if child.byte_range().end <= byte_pos {
524+ prev_child = Some(child);
525+ }
526+ }
527+ match prev_child {
528+ Some(mut prev_child) => {
529+ // Get the deepest directly preceding node
530+ while prev_child.child_count() > 0 {
531+ prev_child = prev_child.child(prev_child.child_count() - 1).unwrap();
532+ }
533+ (
534+ prev_child.byte_range().end - 1..byte_pos + 1,
535+ Some(prev_child),
536+ )
537+ }
538+ None => (byte_pos..byte_pos + 1, None),
539+ }
540+ };
485541 let query_result = query_indents(
486542 query,
487543 syntax,
488544 &mut cursor,
489545 text,
490− byte_pos..byte_pos + 1,
491− new_line_break,
546+ query_range,
547+ new_line.then(|| (line, byte_pos)),
492548 );
493549 ts_parser.cursors.push(cursor);
494− query_result
550+ (query_result, prev_child)
495551 });
552+ let indent_captures = query_result.indent_captures;
553+ let extend_captures = query_result.extend_captures;
554+
555+ // Check for extend captures (starting with the deepest
556+ // candidate node and then going up the syntax tree).
557+ if let Some(mut prev_child) = prev_child {
558+ let mut stop_extend = false;
559+ while prev_child != node {
560+ let mut extend_node = false;
561+ if let Some(captures) = extend_captures.get(&prev_child.id()) {
562+ for capture in captures {
563+ match capture {
564+ ExtendCapture::StopExtend => {
565+ stop_extend = true;
566+ }
567+ ExtendCapture::ExtendIndented => {
568+ // We extend the node if
569+ // - the cursor is on the same line as the end of the node OR
570+ // - the line that the cursor is on is more indented than the
571+ // first line of the node
572+ if prev_child.end_position().row == line {
573+ extend_node = true;
574+ } else {
575+ let cursor_indent =
576+ indent_level_for_line(text.line(line), tab_width);
577+ let node_indent = indent_level_for_line(
578+ text.line(prev_child.start_position().row),
579+ tab_width,
580+ );
581+ if cursor_indent > node_indent {
582+ extend_node = true;
583+ }
584+ }
585+ }
586+ }
587+ }
588+ }
589+ // If we encountered some `StopExtend` capture before, we don't
590+ // extend the node even if we otherwise would
591+ match (extend_node, stop_extend) {
592+ (true, true) => {
593+ stop_extend = false;
594+ }
595+ (true, false) => {
596+ node = prev_child;
597+ break;
598+ }
599+ _ => {}
600+ };
601+ // This parent always exists since node is an ancestor of prev_child
602+ prev_child = prev_child.parent().unwrap();
603+ }
604+ }
605+
606+ let mut first_in_line = get_first_in_line(node, new_line.then(|| byte_pos));
496607
497608 let mut result = Indentation::default();
498609 // We always keep track of all the indent changes on one line, in order to only indent once
@@ -504,7 +615,7 @@ pub fn treesitter_indent_for_pos(
504615 // one entry for each ancestor of the node (which is what we iterate over)
505616 let is_first = *first_in_line.last().unwrap();
506617 // Apply all indent definitions for this node
507− if let Some(definitions) = query_result.get(&node.id()) {
618+ if let Some(definitions) = indent_captures.get(&node.id()) {
508619 for definition in definitions {
509620 match definition.scope {
510621 IndentScope::All => {
@@ -579,6 +690,7 @@ pub fn indent_for_newline(
579690 query,
580691 syntax,
581692 indent_style,
693+ tab_width,
582694 text,
583695 line_before,
584696 line_before_end_pos,