aster/aster-lsp/src/references.rs

use std::collections::HashMap;

use aster_core::analysis;
use aster_core::lexer::Lexer;
use lsp_server::{Connection, Message, Request};
use lsp_types::{Location, Position, Range, ReferenceParams, Uri};

use crate::completion;

pub fn handle_references(
    documents: &HashMap<Uri, String>,
    connection: &Connection,
    req: Request,
) {
    let id = req.id;
    let params: ReferenceParams = match serde_json::from_value(req.params) {
        Ok(p) => p,
        Err(e) => {
            eprintln!("Invalid references params: {}", e);
            return;
        }
    };

    let uri = params.text_document_position.text_document.uri;
    let position = params.text_document_position.position;

    let text = documents.get(&uri).map(|s| s.as_str()).unwrap_or("");
    let word = completion::get_word_at_position(text, position);

    if word.is_empty() {
        let resp = lsp_server::Response::new_ok::<Option<Vec<Location>>>(id, None);
        let _ = connection.sender.send(Message::Response(resp));
        return;
    }

    let (tokens, _) = Lexer::new(text).tokenize();
    let positions = analysis::find_all_references(&tokens, &word);

    let locations: Vec<Location> = positions
        .into_iter()
        .map(|(line, col)| {
            let l = (line as u32).saturating_sub(1);
            let c = (col as u32).saturating_sub(1);
            Location {
                uri: uri.clone(),
                range: Range {
                    start: Position {
                        line: l,
                        character: c,
                    },
                    end: Position {
                        line: l,
                        character: c + word.len() as u32,
                    },
                },
            }
        })
        .collect();

    let resp = lsp_server::Response::new_ok(id, locations);
    let _ = connection.sender.send(Message::Response(resp));
}