From dc443487d4be520efb667f9f93b6c3cbd22e81fd Mon Sep 17 00:00:00 2001
From: Daniel Ebert <daniel.ebert@protonmail.com>
Date: Sat, 17 Sep 2022 15:03:39 +0200
Subject: [PATCH] Slightly change the behavior of the `@stop-extend` capture.
 This improves the behavior in case of multiple nested extensions.

---
 book/src/guides/indent.md |  9 +++++----
 helix-core/src/indent.rs  | 21 +++++++++++----------
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/book/src/guides/indent.md b/book/src/guides/indent.md
index c9306a101..13550846f 100644
--- a/book/src/guides/indent.md
+++ b/book/src/guides/indent.md
@@ -54,10 +54,11 @@ for languages like Python, where for the purpose of indentation some nodes
 
 - `@stop-extend`:
 Prevents the first extension of an ancestor of this node. For example, in Python
-a return expression always ends the block that it is in. Note that this only prevents
-the next extension of one ancestor: If multiple ancestors can be extended (for example
-multiple nested conditional blocks in python), only the extension of the innermost
-ancestor is prevented.
+a return expression always ends the block that it is in. Note that this only stops the
+extension of the next `@extend-indented` capture. If multiple ancestors are captured,
+only the extension of the innermost one is prevented. All other ancestors are unaffected
+(regardless of whether the innermost ancestor would actually have been extended).
+
 
 ## Predicates
 
diff --git a/helix-core/src/indent.rs b/helix-core/src/indent.rs
index 7556c1699..1ea8f18b8 100644
--- a/helix-core/src/indent.rs
+++ b/helix-core/src/indent.rs
@@ -471,6 +471,10 @@ fn extend_nodes<'a>(
         let mut stop_extend = false;
         while deepest_preceding != *node {
             let mut extend_node = false;
+            // This will be set to true if this node is captured, regardless of whether
+            // it actually will be extended (e.g. because the cursor isn't indented
+            // more than the node).
+            let mut node_captured = false;
             if let Some(captures) = extend_captures.get(&deepest_preceding.id()) {
                 for capture in captures {
                     match capture {
@@ -478,6 +482,7 @@ fn extend_nodes<'a>(
                             stop_extend = true;
                         }
                         ExtendCapture::ExtendIndented => {
+                            node_captured = true;
                             // We extend the node if
                             // - the cursor is on the same line as the end of the node OR
                             // - the line that the cursor is on is more indented than the
@@ -501,16 +506,12 @@ fn extend_nodes<'a>(
             }
             // If we encountered some `StopExtend` capture before, we don't
             // extend the node even if we otherwise would
-            match (extend_node, stop_extend) {
-                (true, true) => {
-                    stop_extend = false;
-                }
-                (true, false) => {
-                    *node = deepest_preceding;
-                    break;
-                }
-                _ => {}
-            };
+            if node_captured && stop_extend {
+                stop_extend = false;
+            } else if extend_node && !stop_extend {
+                *node = deepest_preceding;
+                break;
+            }
             // This parent always exists since node is an ancestor of deepest_preceding
             deepest_preceding = deepest_preceding.parent().unwrap();
         }