Fix bug with folding If with multiple statements

src/server/state.rs

@@ -304,15 +304,16 @@ impl ServerState {
             let if_statement = m.captures[0].node;
             let location = if_statement.range().to_location(url);
 
-            let if_body = if_statement.child_by_field_name("body");
-            if let Some(if_body) = if_body {
-                let if_body_location = if_body.range().to_location(url);
-
+            let if_bottom_location = if_statement
+                .children_by_field_name("body", &mut if_statement.walk())
+                .map(|c| c.range().to_location(url))
+                .max_by_key(|l| l.range.end.line);
+            if let Some(if_bottom_location) = if_bottom_location {
                 // If statement body
                 folding_ranges.push(FoldingRange {
                     start_line: location.range.start.line,
                     start_character: None,
-                    end_line: if_body_location.range.end.line,
+                    end_line: if_bottom_location.range.end.line,
                     end_character: None,
                     kind: Some(FoldingRangeKind::Region),
                 });

tree-sitter-quickbms/test/corpus/if.bms

@@ -23,6 +23,23 @@ If 3 == 3
     Endif
 Endif
 
+If 3 == 3
+    print "1"
+    print "2"
+    print "3"
+    print "4"
+    print "5"
+Elif 2 == 2
+    print "1"
+    print "2"
+    print "3"
+Else
+    print "1"
+    print "2"
+    print "3"
+    print "4"
+Endif
+
 --------------------------------------------------------------------------------
 
 (source_file
@@ -77,4 +94,53 @@ Endif
         (print)
         (string_literal))
       (endif))
+    (endif))
+  (if_statement
+    (if)
+    (integer_literal)
+    (comparison)
+    (integer_literal)
+    (print_statement
+      (print)
+      (string_literal))
+    (print_statement
+      (print)
+      (string_literal))
+    (print_statement
+      (print)
+      (string_literal))
+    (print_statement
+      (print)
+      (string_literal))
+    (print_statement
+      (print)
+      (string_literal))
+    (elif_statement
+      (elif)
+      (integer_literal)
+      (comparison)
+      (integer_literal)
+      (print_statement
+        (print)
+        (string_literal))
+      (print_statement
+        (print)
+        (string_literal))
+      (print_statement
+        (print)
+        (string_literal)))
+    (else_statement
+      (else)
+      (print_statement
+        (print)
+        (string_literal))
+      (print_statement
+        (print)
+        (string_literal))
+      (print_statement
+        (print)
+        (string_literal))
+      (print_statement
+        (print)
+        (string_literal)))
     (endif)))