diff --git a/src/server/state.rs b/src/server/state.rs index a644b16..26c0c06 100644 --- a/src/server/state.rs +++ b/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), }); diff --git a/tree-sitter-quickbms/test/corpus/if.bms b/tree-sitter-quickbms/test/corpus/if.bms index 8d82074..fe34dcc 100644 --- a/tree-sitter-quickbms/test/corpus/if.bms +++ b/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)))