From da6e24de1751aef6a944adfcefb192b751c56f76 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Mon, 7 Nov 2022 16:52:19 -0800
Subject: [PATCH] Fix adjustment of trees' included ranges on edits

Previously, when an included range started or ended *inside* of
an edit, that range did not get updated correctly.

Co-authored-by: Mikayla Maki <mikayla@zed.dev>
---
 cli/src/tests/tree_test.rs | 65 ++++++++++++++++++++++++++++++++++++++
 lib/src/tree.c             | 26 +++++++++------
 2 files changed, 81 insertions(+), 10 deletions(-)

diff --git a/cli/src/tests/tree_test.rs b/cli/src/tests/tree_test.rs
index d5c54545..be0c4ff1 100644
--- a/cli/src/tests/tree_test.rs
+++ b/cli/src/tests/tree_test.rs
@@ -232,6 +232,71 @@ fn test_tree_edit() {
     }
 }
 
+#[test]
+fn test_tree_edit_with_included_ranges() {
+    let mut parser = Parser::new();
+    parser.set_language(get_language("html")).unwrap();
+
+    let source = "<div><% if a %><span>a</span><% else %><span>b</span><% end %></div>";
+
+    let ranges = [0..5, 15..29, 39..53, 62..68];
+
+    parser
+        .set_included_ranges(
+            &ranges
+                .iter()
+                .map(|range| Range {
+                    start_byte: range.start,
+                    end_byte: range.end,
+                    start_point: Point::new(0, range.start),
+                    end_point: Point::new(0, range.end),
+                })
+                .collect::<Vec<_>>(),
+        )
+        .unwrap();
+
+    let mut tree = parser.parse(source, None).unwrap();
+
+    tree.edit(&InputEdit {
+        start_byte: 29,
+        old_end_byte: 53,
+        new_end_byte: 29,
+        start_position: Point::new(0, 29),
+        old_end_position: Point::new(0, 53),
+        new_end_position: Point::new(0, 29),
+    });
+
+    assert_eq!(
+        tree.included_ranges(),
+        &[
+            Range {
+                start_byte: 0,
+                end_byte: 5,
+                start_point: Point::new(0, 0),
+                end_point: Point::new(0, 5),
+            },
+            Range {
+                start_byte: 15,
+                end_byte: 29,
+                start_point: Point::new(0, 15),
+                end_point: Point::new(0, 29),
+            },
+            Range {
+                start_byte: 29,
+                end_byte: 29,
+                start_point: Point::new(0, 29),
+                end_point: Point::new(0, 29),
+            },
+            Range {
+                start_byte: 38,
+                end_byte: 44,
+                start_point: Point::new(0, 38),
+                end_point: Point::new(0, 44),
+            }
+        ]
+    );
+}
+
 #[test]
 fn test_tree_cursor() {
     let mut parser = Parser::new();
diff --git a/lib/src/tree.c b/lib/src/tree.c
index 6f747dd1..f6bd2c72 100644
--- a/lib/src/tree.c
+++ b/lib/src/tree.c
@@ -66,17 +66,23 @@ void ts_tree_edit(TSTree *self, const TSInputEdit *edit) {
           range->end_point = POINT_MAX;
         }
       }
-      if (range->start_byte >= edit->old_end_byte) {
-        range->start_byte = edit->new_end_byte + (range->start_byte - edit->old_end_byte);
-        range->start_point = point_add(
-          edit->new_end_point,
-          point_sub(range->start_point, edit->old_end_point)
-        );
-        if (range->start_byte < edit->new_end_byte) {
-          range->start_byte = UINT32_MAX;
-          range->start_point = POINT_MAX;
-        }
+    } else if (range->end_byte > edit->start_byte) {
+      range->end_byte = edit->start_byte;
+      range->end_point = edit->start_point;
+    }
+    if (range->start_byte >= edit->old_end_byte) {
+      range->start_byte = edit->new_end_byte + (range->start_byte - edit->old_end_byte);
+      range->start_point = point_add(
+        edit->new_end_point,
+        point_sub(range->start_point, edit->old_end_point)
+      );
+      if (range->start_byte < edit->new_end_byte) {
+        range->start_byte = UINT32_MAX;
+        range->start_point = POINT_MAX;
       }
+    } else if (range->start_byte > edit->start_byte) {
+      range->start_byte = edit->start_byte;
+      range->start_point = edit->start_point;
     }
   }