laburnum 1.17.0

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

//! # Diagnostic Storage Conventions
//!
//! Diagnostics should be stored in Partitions using the following key
//! structure. These are conventions only - users are free to implement their
//! own strategies.
//!
//! ## Partition Key Format
//!
//! ```text
//! diagnostics
//! ```
//!
//! All diagnostics use a single partition key `diagnostics`.
//!
//! Rationale:
//! - Simplifies watcher patterns (single partition to monitor)
//! - File-level granularity provided by sort key prefix
//! - Built-in diagnostic watcher can react to all diagnostic changes
//!
//! ## Sort Key Format
//!
//! ```text
//! {source_key_u64:016x}|{severity_num:01}|{sequence:04}
//! ```
//!
//! Components:
//! - `source_key_u64`: SourceKey as hex string (16 chars, zero-padded)
//! - `severity_num`: DiagnosticSeverity as number (1=Error, 2=Warning, 3=Info,
//!   4=Hint)
//! - `sequence`: Zero-padded sequence number (4 digits)
//!
//! Examples:
//! - `00000000000a0001|1|0000` - First error in file_id=10, version=1
//! - `00000000000a0001|1|0001` - Second error in same file
//! - `00000000000a0001|2|0000` - First warning in same file
//!
//! Rationale:
//! - SourceKey prefix enables efficient file-level queries via
//!   `sort_key_begins_with`
//! - Severity ordering groups diagnostics by importance
//! - Sequence ensures stable ordering for diagnostics at same location
//!
//! ## Query Patterns
//!
//! ```rust,ignore
//! // Get all diagnostics for a file
//! let source_key_prefix = format!("{:016x}", source_key.as_u64());
//! query.sort_key_begins_with(source_key_prefix)
//!
//! // Get only errors for a file
//! let error_prefix = format!("{:016x}|1", source_key.as_u64());
//! query.sort_key_begins_with(error_prefix)
//! ```
//!
//! ## Implementation Notes
//!
//! Users should:
//! 1. Create their own diagnostic record types (e.g., `NanoRecord::ParseError`)
//! 2. Implement the `Diagnostic` trait for their types
//! 3. Store diagnostics using the partition/sort key convention above
//! 4. Implement `Partitions::get_diagnostics()` to query their diagnostic
//!    records

use {
  crate::{
    Uri,
    connect::lsp::{ClientId, notification::PublishDiagnostics},
    database::{
      PartitionWriteContextRef,
      Partitions,
    },
    protocol::lsp::{
      self,
      LanguageServer,
      PositionEncodingKind,
      PublishDiagnosticsParams,
    },
    record::LaburnumRecordRef,
    scheduler::{key_watcher::WatcherResult, task::TaskContext},
    source::SourceKey,
  },
  std::{
    collections::HashMap,
    future::Future,
    pin::Pin,
  },
};

type ClientDiagnostics = HashMap<Uri, (Option<SourceKey>, Vec<lsp::Diagnostic>)>;

pub fn diagnostic_watcher<'a, P, T>(
  ctx: &'a mut TaskContext<P, T>,
  _writer: &'a mut PartitionWriteContextRef<'a, P>,
) -> Pin<Box<dyn Future<Output = WatcherResult<P, T>> + Send + 'a>>
where
  P: Partitions,
  T: LanguageServer<P>,
{
  Box::pin(async move {
    let updated_keys = ctx.matched_keys_updated().to_vec();
    let deleted_keys = ctx.matched_keys_deleted().to_vec();

    eprintln!(
      "DEBUG: diagnostic_watcher called - updated: {}, deleted: {}",
      updated_keys.len(),
      deleted_keys.len()
    );

    let source_cache_reader = ctx.source_cache_reader();

    let mut per_client: HashMap<ClientId, ClientDiagnostics> = HashMap::new();

    let mut broadcast: ClientDiagnostics = HashMap::new();

    for key in &updated_keys {
      let results = key.get_record(ctx.query_client()).await;

      for record_meta in results.records.iter() {
        let Some(record) = results.get(record_meta) else {
          continue;
        };
        let Some(diagnostic) = record.as_dyn_diagnostic() else {
          continue;
        };
        let Some(source_key) = diagnostic.source_key() else {
          continue;
        };
        let Some(source) = source_cache_reader.get_source(source_key) else {
          continue;
        };

        let uri = source.uri().clone();
        let clients = source_cache_reader.clients_for(source_key);

        if clients.is_empty() {
          let lsp_diag = diagnostic.to_lsp_diagnostic(
            &source_cache_reader,
            &PositionEncodingKind::DEFAULT,
          );
          let entry = broadcast.entry(uri).or_default();
          entry.0 = Some(source_key);
          entry.1.push(lsp_diag);
        } else {
          for client_id in clients {
            let encoding = ctx
              .scheduler()
              .registry()
              .get(client_id)
              .map(|c| c.position_encoding().clone())
              .unwrap_or(PositionEncodingKind::DEFAULT);

            let lsp_diag =
              diagnostic.to_lsp_diagnostic(&source_cache_reader, &encoding);

            let client_map = per_client.entry(client_id).or_default();
            let entry = client_map.entry(uri.clone()).or_default();
            entry.0 = Some(source_key);
            entry.1.push(lsp_diag);
          }
        }
      }
    }

    let mut notified_uris: std::collections::HashSet<Uri> =
      std::collections::HashSet::new();

    for (client_id, uri_map) in per_client {
      for (uri, (source_key, diagnostics)) in uri_map {
        notified_uris.insert(uri.clone());
        let version =
          source_key.and_then(|sk| source_cache_reader.get_lsp_version(sk));
        let params = PublishDiagnosticsParams::new(uri, diagnostics, version);
        ctx
          .send_notification::<PublishDiagnostics>(params, Some(client_id))
          .await;
      }
    }

    for (uri, (source_key, diagnostics)) in broadcast {
      notified_uris.insert(uri.clone());
      let version =
        source_key.and_then(|sk| source_cache_reader.get_lsp_version(sk));
      let params = PublishDiagnosticsParams::new(uri, diagnostics, version);
      ctx
        .send_notification::<PublishDiagnostics>(params, None)
        .await;
    }

    let mut deleted_uris: HashMap<Uri, Option<i32>> = HashMap::new();

    for key in &deleted_keys {
      let Some(source_key) = key.source_key() else {
        continue;
      };
      let Some(source) = source_cache_reader.get_source(source_key) else {
        continue;
      };

      let uri = source.uri().clone();
      let version = source_cache_reader.get_lsp_version(source_key);
      deleted_uris.insert(uri, version);
    }

    for (uri, version) in deleted_uris {
      if !notified_uris.contains(&uri) {
        let params = PublishDiagnosticsParams::new(uri, vec![], version);
        // Broadcast cleared diagnostics to all clients
        ctx.send_notification::<PublishDiagnostics>(params, None).await;
      }
    }

    WatcherResult::empty()
  })
}