use {
crate::{
TRACER,
database::{
DynPartition,
PartitionKey,
PartitionWriteContextRef,
},
partitions::{
DocumentSymbols,
TextDocumentReferences,
TextDocumentPosition,
},
protocol::{
jsonrpc,
lsp::{
DynamicRegistrationClientCapabilities,
Location,
PartialResultParams,
Position,
TextDocumentIdentifier,
TextDocumentPositionParams,
WorkDoneProgressParams,
},
},
record::LaburnumRecordRef,
scheduler::task::TaskContext,
},
opentelemetry::trace::FutureExt,
serde::{
Deserialize,
Serialize,
},
};
pub type ReferenceClientCapabilities = DynamicRegistrationClientCapabilities;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ReferenceContext {
pub include_declaration: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ReferenceParams {
#[serde(flatten)]
pub text_document_position: TextDocumentPositionParams,
#[serde(flatten)]
pub work_done_progress_params: WorkDoneProgressParams,
#[serde(flatten)]
pub partial_result_params: PartialResultParams,
pub context: ReferenceContext,
}
pub trait ReferenceService<
P: crate::database::storage::Partitions,
T: crate::protocol::lsp::LanguageServer<P>,
>: Send + Sync + 'static
{
fn references(
&self,
params: ReferenceParams,
ctx: &mut TaskContext<P, T>,
writer: &mut PartitionWriteContextRef<'_, P>,
) -> impl std::future::Future<Output = jsonrpc::Result<Option<Vec<Location>>>> + Send
{
let _ = writer;
let cx = otel::span!(^
"laburnum.lsp.references",
"document.uri" = params.text_document_position.text_document.uri.to_string(),
"position.line" = params.text_document_position.position.line as i64,
"position.character" = params.text_document_position.position.character as i64,
"include_declaration" = params.context.include_declaration
);
let encoding = ctx.position_encoding();
async move {
let uri = params.text_document_position.text_document.uri;
let position = params.text_document_position.position;
let include_declaration = params.context.include_declaration;
let source_key = {
let cache = ctx.source_cache();
let guard = cache.read();
match guard.latest_key(&uri) {
| Some(key) => key,
| None => return Ok(None),
}
};
use crate::source::line_ops::LineOps;
let source_cache = ctx.source_cache_reader();
let byte_offset = {
let source = match source_cache.get_source(source_key) {
Some(s) => s,
None => return Ok(None),
};
match source.line_col_to_byte(position.line, position.character, &encoding) {
Some(o) => o as u64,
None => return Ok(None),
}
};
let Some((ident, symbol_hash, is_ident)) = ctx
.query_client()
.span_index_get::<TextDocumentPosition>(&uri, byte_offset)
.and_then(|record| {
let pos = record.as_text_document_position()?;
Some((
pos.symbol_ident(),
pos.symbol_hash(),
pos.kind() == crate::partitions::text_document_position::PositionKind::Ident,
))
})
else {
return Ok(None);
};
let Some((source_key, symbol_kind, symbol_def_loc)) = ctx
.query_client()
.get_by_hash::<DocumentSymbols>(symbol_hash)
.and_then(|r| {
let sym = r.as_document_symbol()?;
let def_loc = if !include_declaration {
sym.definition_location(&source_cache, &encoding)
} else {
None
};
Some((sym.source_key(), sym.symbol_kind(), def_loc))
})
else {
return Ok(None);
};
let (target_ident, target_source_key, def_location) = if is_ident {
(ident, source_key, symbol_def_loc)
} else {
let symbol_sort_key = crate::partitions::document_symbols::DocumentSymbolSortKey::Symbol {
source_key,
symbol_kind,
ident,
}.to_string();
let symbol_results = ctx
.query_client()
.get_record(DocumentSymbols::KEY, symbol_sort_key)
.await;
if let Some(symbol_meta) = symbol_results.records.first() {
if let Some(symbol_record) = symbol_results.get(symbol_meta) {
if let Some(target_symbol) = symbol_record.as_document_symbol() {
let target_ident = target_symbol.get_ident();
let target_source_key = target_symbol.target_source_key().unwrap_or_else(|| target_symbol.source_key());
let def_loc = if !include_declaration {
target_symbol.definition_location(&source_cache, &encoding)
} else {
None
};
(target_ident, target_source_key, def_loc)
} else {
return Ok(None);
}
} else {
return Ok(None);
}
} else {
return Ok(None);
}
};
let references_prefix = crate::partitions::text_document_references::TextDocumentReferenceSortKey::SymbolPrefix {
source_key: target_source_key,
ident: target_ident,
};
let ref_records = ctx
.query_client()
.query_partition(TextDocumentReferences)
.sort_key_begins_with(references_prefix)
.execute()
.await;
let mut references = Vec::new();
if include_declaration {
let symbol_sort_key = crate::partitions::document_symbols::DocumentSymbolSortKey::Symbol {
source_key: target_source_key,
symbol_kind,
ident: target_ident,
}.to_string();
let symbol_results = ctx
.query_client()
.get_record(DocumentSymbols::KEY, symbol_sort_key)
.await;
if let Some(symbol_meta) = symbol_results.records.first()
&& let Some(symbol_record) = symbol_results.get(symbol_meta)
&& let Some(def_symbol) = symbol_record.as_document_symbol()
&& let Some(def_loc) = def_symbol.definition_location(&source_cache, &encoding)
{
references.push(def_loc);
}
}
for record_meta in &ref_records.records {
if let Some(record) = ref_records.get(record_meta)
&& let Some(ref_record) = record.as_text_document_reference()
&& let Some(location) = ref_record.location(&source_cache, &encoding)
{
references.push(location);
}
}
if let Some(def_loc) = def_location {
references.retain(|loc| loc != &def_loc);
}
if references.is_empty() {
Ok(None)
} else {
Ok(Some(references))
}
}.with_context(cx)
}
const REFERENCES_LANE: crate::scheduler::lanes::Lane =
crate::scheduler::lanes::DEFAULT_LANE;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_references_request_params() {
let params = ReferenceParams {
text_document_position: TextDocumentPositionParams {
text_document: TextDocumentIdentifier {
uri: "file:///test.txt".parse().unwrap(),
},
position: Position {
line: 10,
character: 5,
},
},
work_done_progress_params: WorkDoneProgressParams::default(),
partial_result_params: PartialResultParams::default(),
context: ReferenceContext {
include_declaration: true,
},
};
assert_eq!(params.text_document_position.position.line, 10);
assert_eq!(params.text_document_position.position.character, 5);
assert!(params.context.include_declaration);
}
}