@@ -456,6 +456,67 @@ fn query_indents(
456456 }
457457}
458458
459+/// Handle extend queries. deepest_preceding is the deepest descendant of node that directly precedes the cursor position.
460+/// Any ancestor of deepest_preceding which is also a descendant of node may be "extended". In that case, node will be updated,
461+/// so that the indent computation starts with the correct syntax node.
462+fn extend_nodes<'a>(
463+ node: &mut Node<'a>,
464+ deepest_preceding: Option<Node<'a>>,
465+ extend_captures: &HashMap<usize, Vec<ExtendCapture>>,
466+ text: RopeSlice,
467+ line: usize,
468+ tab_width: usize,
469+) {
470+ if let Some(mut deepest_preceding) = deepest_preceding {
471+ let mut stop_extend = false;
472+ while deepest_preceding != *node {
473+ let mut extend_node = false;
474+ if let Some(captures) = extend_captures.get(&deepest_preceding.id()) {
475+ for capture in captures {
476+ match capture {
477+ ExtendCapture::StopExtend => {
478+ stop_extend = true;
479+ }
480+ ExtendCapture::ExtendIndented => {
481+ // We extend the node if
482+ // - the cursor is on the same line as the end of the node OR
483+ // - the line that the cursor is on is more indented than the
484+ // first line of the node
485+ if deepest_preceding.end_position().row == line {
486+ extend_node = true;
487+ } else {
488+ let cursor_indent =
489+ indent_level_for_line(text.line(line), tab_width);
490+ let node_indent = indent_level_for_line(
491+ text.line(deepest_preceding.start_position().row),
492+ tab_width,
493+ );
494+ if cursor_indent > node_indent {
495+ extend_node = true;
496+ }
497+ }
498+ }
499+ }
500+ }
501+ }
502+ // If we encountered some `StopExtend` capture before, we don't
503+ // extend the node even if we otherwise would
504+ match (extend_node, stop_extend) {
505+ (true, true) => {
506+ stop_extend = false;
507+ }
508+ (true, false) => {
509+ *node = deepest_preceding;
510+ break;
511+ }
512+ _ => {}
513+ };
514+ // This parent always exists since node is an ancestor of deepest_preceding
515+ deepest_preceding = deepest_preceding.parent().unwrap();
516+ }
517+ }
518+}
519+
459520/// Use the syntax tree to determine the indentation for a given position.
460521/// This can be used in 2 ways:
461522///
@@ -510,100 +571,54 @@ pub fn treesitter_indent_for_pos(
510571 .tree()
511572 .root_node()
512573 .descendant_for_byte_range(byte_pos, byte_pos)?;
513− let (query_result, prev_child) = crate::syntax::PARSER.with(|ts_parser| {
514− let mut ts_parser = ts_parser.borrow_mut();
515− let mut cursor = ts_parser.cursors.pop().unwrap_or_else(QueryCursor::new);
574+ let (query_result, deepest_preceding) = {
516575 // The query range should intersect with all nodes directly preceding
517576 // the cursor in case one of them is extended.
518− // prev_child is the deepest such node.
519− let (query_range, prev_child) = {
520− // TODO Is there some way we can reuse this cursor?
521− let mut tree_cursor = node.walk();
522− let mut prev_child = None;
523− for child in node.children(&mut tree_cursor) {
524− if child.byte_range().end <= byte_pos {
525− prev_child = Some(child);
526− }
577+ let mut deepest_preceding = None; // The deepest node preceding the cursor
578+ let mut tree_cursor = node.walk();
579+ for child in node.children(&mut tree_cursor) {
580+ if child.byte_range().end <= byte_pos {
581+ deepest_preceding = Some(child);
527582 }
528− match prev_child {
529− Some(mut prev_child) => {
530− // Get the deepest directly preceding node
531− while prev_child.child_count() > 0 {
532− prev_child = prev_child.child(prev_child.child_count() - 1).unwrap();
533− }
534− (
535− prev_child.byte_range().end - 1..byte_pos + 1,
536− Some(prev_child),
537− )
538− }
539− None => (byte_pos..byte_pos + 1, None),
583+ }
584+ deepest_preceding = deepest_preceding.map(|mut prec| {
585+ // Get the deepest directly preceding node
586+ while prec.child_count() > 0 {
587+ prec = prec.child(prec.child_count() - 1).unwrap();
540588 }
541− };
542− let query_result = query_indents(
543− query,
544− syntax,
545− &mut cursor,
546− text,
547− query_range,
548− new_line.then(|| (line, byte_pos)),
549− );
550− ts_parser.cursors.push(cursor);
551− (query_result, prev_child)
552− });
589+ prec
590+ });
591+ let query_range = deepest_preceding
592+ .map(|prec| prec.byte_range().end - 1..byte_pos + 1)
593+ .unwrap_or(byte_pos..byte_pos + 1);
594+
595+ crate::syntax::PARSER.with(|ts_parser| {
596+ let mut ts_parser = ts_parser.borrow_mut();
597+ let mut cursor = ts_parser.cursors.pop().unwrap_or_else(QueryCursor::new);
598+ let query_result = query_indents(
599+ query,
600+ syntax,
601+ &mut cursor,
602+ text,
603+ query_range,
604+ new_line.then(|| (line, byte_pos)),
605+ );
606+ ts_parser.cursors.push(cursor);
607+ (query_result, deepest_preceding)
608+ })
609+ };
553610 let indent_captures = query_result.indent_captures;
554611 let extend_captures = query_result.extend_captures;
555612
556− // Check for extend captures (starting with the deepest
557− // candidate node and then going up the syntax tree).
558− if let Some(mut prev_child) = prev_child {
559− let mut stop_extend = false;
560− while prev_child != node {
561− let mut extend_node = false;
562− if let Some(captures) = extend_captures.get(&prev_child.id()) {
563− for capture in captures {
564− match capture {
565− ExtendCapture::StopExtend => {
566− stop_extend = true;
567− }
568− ExtendCapture::ExtendIndented => {
569− // We extend the node if
570− // - the cursor is on the same line as the end of the node OR
571− // - the line that the cursor is on is more indented than the
572− // first line of the node
573− if prev_child.end_position().row == line {
574− extend_node = true;
575− } else {
576− let cursor_indent =
577− indent_level_for_line(text.line(line), tab_width);
578− let node_indent = indent_level_for_line(
579− text.line(prev_child.start_position().row),
580− tab_width,
581− );
582− if cursor_indent > node_indent {
583− extend_node = true;
584− }
585− }
586− }
587− }
588− }
589− }
590− // If we encountered some `StopExtend` capture before, we don't
591− // extend the node even if we otherwise would
592− match (extend_node, stop_extend) {
593− (true, true) => {
594− stop_extend = false;
595− }
596− (true, false) => {
597− node = prev_child;
598− break;
599− }
600− _ => {}
601− };
602− // This parent always exists since node is an ancestor of prev_child
603− prev_child = prev_child.parent().unwrap();
604− }
605− }
606−
613+ // Check for extend captures, potentially changing the node that the indent calculation starts with
614+ extend_nodes(
615+ &mut node,
616+ deepest_preceding,
617+ &extend_captures,
618+ text,
619+ line,
620+ tab_width,
621+ );
607622 let mut first_in_line = get_first_in_line(node, new_line.then(|| byte_pos));
608623
609624 let mut result = Indentation::default();