1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

use erg_common::error::Location;
use erg_common::traits::Locational;
use erg_compiler::artifact::BuildRunnable;
use erg_compiler::erg_parser::ast::Expr;
use erg_compiler::erg_parser::parse::Parsable;

use lsp_types::{FoldingRange, FoldingRangeKind, FoldingRangeParams};

use crate::_log;
use crate::server::{ELSResult, RedirectableStdout, Server};
use crate::util::NormalizedUrl;

fn imports_range(start: &Location, end: &Location) -> Option<FoldingRange> {
    Some(FoldingRange {
        start_line: start.ln_begin()?.saturating_sub(1),
        start_character: start.col_begin(),
        end_line: end.ln_end()?.saturating_sub(1),
        end_character: end.col_end(),
        kind: Some(FoldingRangeKind::Imports),
    })
}

impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
    pub(crate) fn handle_folding_range(
        &mut self,
        params: FoldingRangeParams,
    ) -> ELSResult<Option<Vec<FoldingRange>>> {
        _log!(self, "folding range requested: {params:?}");
        let uri = NormalizedUrl::new(params.text_document.uri);
        let mut res = vec![];
        res.extend(self.fold_imports(&uri));
        Ok(Some(res))
    }

    fn fold_imports(&self, uri: &NormalizedUrl) -> Vec<FoldingRange> {
        let mut res = vec![];
        if let Ok(module) = self.build_ast(uri) {
            let mut ranges = vec![];
            for chunk in module.into_iter() {
                match chunk {
                    Expr::Def(def) if def.def_kind().is_import() => {
                        ranges.push(def.loc());
                    }
                    _ => {
                        if !ranges.is_empty() {
                            let start = ranges.first().unwrap();
                            let end = ranges.last().unwrap();
                            res.extend(imports_range(start, end));
                            ranges.clear();
                        }
                    }
                }
            }
            if !ranges.is_empty() {
                let start = ranges.first().unwrap();
                let end = ranges.last().unwrap();
                res.extend(imports_range(start, end));
            }
        }
        res
    }
}