laburnum 1.17.0

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

pub use notification_params::*;
use {
  super::LspVersion,
  crate::{
    Uri,
    database::PartitionWriteContextRef,
    protocol::{
      lsp::{LSPObject, PositionEncodingKind},
      macros::lsp_enum,
    },
    scheduler::task::TaskContext,
  },
  serde::{
    Deserialize,
    Serialize,
  },
};

fn deserialize_lsp_version<'de, D>(deserializer: D) -> Result<LspVersion, D::Error>
where
  D: serde::Deserializer<'de>,
{
  let value = i32::deserialize(deserializer)?;
  Ok(value.into())
}

fn serialize_lsp_version<S>(version: &LspVersion, serializer: S) -> Result<S::Ok, S::Error>
where
  S: serde::Serializer,
{
  serializer.serialize_i32(version.raw())
}

/// A notebook document.
///
/// @since 3.17.0
#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct NotebookDocument {
  /// The notebook document's URI.
  pub uri:           Uri,
  /// The type of the notebook.
  pub notebook_type: String,
  /// The version number of this document (it will increase after each
  /// change, including undo/redo).
  #[serde(deserialize_with = "deserialize_lsp_version", serialize_with = "serialize_lsp_version")]
  pub version:       LspVersion,
  /// Additional metadata stored with the notebook
  /// document.
  #[serde(skip_serializing_if = "Option::is_none")]
  pub metadata:      Option<LSPObject>,
  /// The cells of a notebook.
  pub cells:         Vec<NotebookCell>,
}

/// A notebook cell.
///
/// A cell's document URI must be unique across ALL notebook
/// cells and can therefore be used to uniquely identify a
/// notebook cell or the cell's text document.
///
/// @since 3.17.0
#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct NotebookCell {
  /// The cell's kind
  pub kind:              NotebookCellKind,
  /// The URI of the cell's text document content.
  pub document:          Uri,
  /// Additional metadata stored with the cell.
  #[serde(skip_serializing_if = "Option::is_none")]
  pub metadata:          Option<LSPObject>,
  /// Additional execution summary information
  /// if supported by the client.
  #[serde(skip_serializing_if = "Option::is_none")]
  pub execution_summary: Option<ExecutionSummary>,
}

#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ExecutionSummary {
  /// A strict monotonically increasing value
  /// indicating the execution order of a cell
  /// inside a notebook.
  pub execution_order: u32,
  /// Whether the execution was successful or
  /// not if known by the client.
  #[serde(skip_serializing_if = "Option::is_none")]
  pub success:         Option<bool>,
}

#[derive(Clone, PartialEq, Eq, Deserialize, Serialize)]
pub struct NotebookCellKind(i32);

lsp_enum! {
    impl NotebookCellKind {
        /// A markup-cell is formatted source that is used for display.
        const MARKUP = 1;
        /// A code-cell is source code.
        const CODE = 2;
    }
}

/// Capabilities specific to the notebook document support.
///
/// @since 3.17.0
#[derive(Debug, Default, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct NotebookDocumentClientCapabilities {
  /// Capabilities specific to notebook document synchronization
  ///
  /// @since 3.17.0
  pub synchronization: NotebookDocumentSyncClientCapabilities,
}

/// Notebook specific client capabilities.
///
/// @since 3.17.0
#[derive(Debug, Default, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct NotebookDocumentSyncClientCapabilities {
  /// Whether implementation supports dynamic registration. If this is
  /// set to `true` the client supports the new
  /// `(TextDocumentRegistrationOptions & StaticRegistrationOptions)`
  /// return value for the corresponding server capability as well.
  #[serde(skip_serializing_if = "Option::is_none")]
  pub dynamic_registration: Option<bool>,

  /// The client supports sending execution summary data per cell.
  #[serde(skip_serializing_if = "Option::is_none")]
  pub execution_summary_report: Option<bool>,
}

///  Options specific to a notebook plus its cells
///  to be synced to the server.
///
///  If a selector provides a notebook document
///  filter but no cell selector all cells of a
///  matching notebook document will be synced.
///
///  If a selector provides no notebook document
///  filter but only a cell selector all notebook
///  documents that contain at least one matching
///  cell will be synced.
///
///  @since 3.17.0
#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct NotebookDocumentSyncOptions {
  /// The notebooks to be synced
  pub notebook_selector: Vec<NotebookSelector>,
  /// Whether save notification should be forwarded to
  /// the server. Will only be honored if mode === `notebook`.
  #[serde(skip_serializing_if = "Option::is_none")]
  pub save:              Option<bool>,
}

/// Registration options specific to a notebook.
///
/// @since 3.17.0
#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct NotebookDocumentSyncRegistrationOptions {
  /// The notebooks to be synced
  pub notebook_selector: Vec<NotebookSelector>,
  /// Whether save notification should be forwarded to
  /// the server. Will only be honored if mode === `notebook`.
  #[serde(skip_serializing_if = "Option::is_none")]
  pub save:              Option<bool>,
  /// The id used to register the request. The id can be used to deregister
  /// the request again. See also Registration#id.
  #[serde(skip_serializing_if = "Option::is_none")]
  pub id:                Option<String>,
}

/// A notebook cell text document filter denotes a cell text
/// document by different properties.
///
/// @since 3.17.0
#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct NotebookCellTextDocumentFilter {
  /// A filter that matches against the notebook
  /// containing the notebook cell. If a string
  /// value is provided it matches against the
  /// notebook type. '*' matches every notebook.
  pub notebook: Notebook,
  /// A language id like `python`.
  ///
  /// Will be matched against the language id of the
  /// notebook cell document. '*' matches every language.
  #[serde(skip_serializing_if = "Option::is_none")]
  pub language: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", untagged)]
pub enum NotebookSelector {
  ByNotebook {
    /// The notebook to be synced. If a string
    /// value is provided it matches against the
    /// notebook type. '*' matches every notebook.
    notebook: Notebook,
    /// The cells of the matching notebook to be synced.
    #[serde(skip_serializing_if = "Option::is_none")]
    cells:    Option<Vec<NotebookCellSelector>>,
  },
  ByCells {
    /// The notebook to be synced. If a string
    /// value is provided it matches against the
    /// notebook type. '*' matches every notebook.
    #[serde(skip_serializing_if = "Option::is_none")]
    notebook: Option<Notebook>,
    /// The cells of the matching notebook to be synced.
    cells:    Vec<NotebookCellSelector>,
  },
}

#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct NotebookCellSelector {
  pub language: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(untagged)]
pub enum Notebook {
  String(String),
  NotebookDocumentFilter(NotebookDocumentFilter),
}

/// A notebook document filter denotes a notebook document by
/// different properties.
///
/// @since 3.17.0
#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", untagged)]
pub enum NotebookDocumentFilter {
  ByType {
    /// The type of the enclosing notebook.
    notebook_type: String,
    /// A Uri [scheme](#Uri.scheme), like `file` or `untitled`.
    #[serde(skip_serializing_if = "Option::is_none")]
    scheme:        Option<String>,
    /// A glob pattern.
    #[serde(skip_serializing_if = "Option::is_none")]
    pattern:       Option<String>,
  },
  ByScheme {
    /// The type of the enclosing notebook.
    #[serde(skip_serializing_if = "Option::is_none")]
    notebook_type: Option<String>,
    /// A Uri [scheme](#Uri.scheme), like `file` or `untitled`.
    scheme:        String,
    /// A glob pattern.
    #[serde(skip_serializing_if = "Option::is_none")]
    pattern:       Option<String>,
  },
  ByPattern {
    /// The type of the enclosing notebook.
    #[serde(skip_serializing_if = "Option::is_none")]
    notebook_type: Option<String>,
    /// A Uri [scheme](#Uri.scheme), like `file` or `untitled`.
    #[serde(skip_serializing_if = "Option::is_none")]
    scheme:        Option<String>,
    /// A glob pattern.
    pattern:       String,
  },
}

mod notification_params {
  use {
    crate::{
      protocol::lsp::{
        LSPObject,
        LspVersion,
        NotebookCell,
        NotebookDocument,
        TextDocumentContentChangeEvent,
        TextDocumentIdentifier,
        TextDocumentItem,
        Uri,
        VersionedTextDocumentIdentifier,
      },
    },
    serde::{
      Deserialize,
      Serialize,
    },
  };

  fn deserialize_lsp_version<'de, D>(deserializer: D) -> Result<LspVersion, D::Error>
  where
    D: serde::Deserializer<'de>,
  {
    let value = i32::deserialize(deserializer)?;
    Ok(value.into())
  }

  fn serialize_lsp_version<S>(version: &LspVersion, serializer: S) -> Result<S::Ok, S::Error>
  where
    S: serde::Serializer,
  {
    serializer.serialize_i32(version.raw())
  }

  /// The params sent in an open notebook document notification.
  ///
  /// @since 3.17.0
  #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
  #[serde(rename_all = "camelCase")]
  pub struct DidOpenNotebookDocumentParams {
    /// The notebook document that got opened.
    pub notebook_document:   NotebookDocument,
    /// The text documents that represent the content
    /// of a notebook cell.
    pub cell_text_documents: Vec<TextDocumentItem>,
  }

  /// The params sent in a change notebook document notification.
  ///
  /// @since 3.17.0
  #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
  #[serde(rename_all = "camelCase")]
  pub struct DidChangeNotebookDocumentParams {
    /// The notebook document that did change. The version number points
    /// to the version after all provided changes have been applied.
    pub notebook_document: VersionedNotebookDocumentIdentifier,

    /// The actual changes to the notebook document.
    ///
    /// The change describes single state change to the notebook document.
    /// So it moves a notebook document, its cells and its cell text document
    /// contents from state S to S'.
    ///
    /// To mirror the content of a notebook using change events use the
    /// following approach:
    /// - start with the same initial content
    /// - apply the `notebookDocument/didChange` notifications in the order you
    ///   receive them.
    pub change: NotebookDocumentChangeEvent,
  }

  /// A versioned notebook document identifier.
  ///
  /// @since 3.17.0
  #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
  #[serde(rename_all = "camelCase")]
  pub struct VersionedNotebookDocumentIdentifier {
    /// The version number of this notebook document.
    #[serde(deserialize_with = "deserialize_lsp_version", serialize_with = "serialize_lsp_version")]
    pub version: LspVersion,
    /// The notebook document's URI.
    pub uri:     Uri,
  }

  /// A change event for a notebook document.
  ///
  /// @since 3.17.0
  #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
  #[serde(rename_all = "camelCase")]
  pub struct NotebookDocumentChangeEvent {
    /// The changed meta data if any.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<LSPObject>,

    /// Changes to cells
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cells: Option<NotebookDocumentCellChange>,
  }

  #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
  #[serde(rename_all = "camelCase")]
  pub struct NotebookDocumentCellChange {
    /// Changes to the cell structure to add or
    /// remove cells.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub structure: Option<NotebookDocumentCellChangeStructure>,

    /// Changes to notebook cells properties like its
    /// kind, execution summary or metadata.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub data: Option<Vec<NotebookCell>>,

    /// Changes to the text content of notebook cells.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub text_content: Option<Vec<NotebookDocumentChangeTextContent>>,
  }

  #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
  #[serde(rename_all = "camelCase")]
  pub struct NotebookDocumentChangeTextContent {
    pub document: VersionedTextDocumentIdentifier,
    pub changes:  Vec<TextDocumentContentChangeEvent>,
  }

  #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
  #[serde(rename_all = "camelCase")]
  pub struct NotebookDocumentCellChangeStructure {
    /// The change to the cell array.
    pub array:     NotebookCellArrayChange,
    /// Additional opened cell text documents.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub did_open:  Option<Vec<TextDocumentItem>>,
    /// Additional closed cell text documents.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub did_close: Option<Vec<TextDocumentIdentifier>>,
  }

  /// A change describing how to move a `NotebookCell`
  /// array from state S to S'.
  ///
  /// @since 3.17.0
  #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
  #[serde(rename_all = "camelCase")]
  pub struct NotebookCellArrayChange {
    /// The start offset of the cell that changed.
    pub start: u32,

    /// The deleted cells
    pub delete_count: u32,

    /// The new cells, if any
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cells: Option<Vec<NotebookCell>>,
  }

  /// The params sent in a save notebook document notification.
  ///
  /// @since 3.17.0
  #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
  #[serde(rename_all = "camelCase")]
  pub struct DidSaveNotebookDocumentParams {
    /// The notebook document that got saved.
    pub notebook_document: NotebookDocumentIdentifier,
  }

  /// A literal to identify a notebook document in the client.
  ///
  /// @since 3.17.0
  #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
  #[serde(rename_all = "camelCase")]
  pub struct NotebookDocumentIdentifier {
    /// The notebook document's URI.
    pub uri: Uri,
  }

  /// The params sent in a close notebook document notification.
  ///
  /// @since 3.17.0
  #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
  #[serde(rename_all = "camelCase")]
  pub struct DidCloseNotebookDocumentParams {
    /// The notebook document that got closed.
    pub notebook_document: NotebookDocumentIdentifier,

    /// The text documents that represent the content
    /// of a notebook cell that got closed.
    pub cell_text_documents: Vec<TextDocumentIdentifier>,
  }
}

pub trait NotebookDocumentService<
  P: crate::database::storage::Partitions,
  T: crate::protocol::lsp::LanguageServer<P>,
>: Send + Sync + 'static
{
  /// The [`notebookDocument/didOpen`] notification is sent from the client to
  /// the server when a new notebook document is opened. It is only sent for
  /// notebooks selected by the `notebookDocumentSync` server capability.
  ///
  /// [`notebookDocument/didOpen`]: https://microsoft.github.io/language-server-protocol/specification/#notebookDocument_didChange
  fn notebook_did_open(
    &self,
    params: DidOpenNotebookDocumentParams,
    ctx: &mut TaskContext<P, T>,
    writer: &mut PartitionWriteContextRef<'_, P>,
  ) -> impl std::future::Future<Output = ()> + Send {
    let client_id = ctx.client_id();

    async move {
      let _ = writer;
      let notebook_uri = params.notebook_document.uri.clone();
      let version = params.notebook_document.version;

      let source_cache = ctx.source_cache();
      let mut cache = source_cache.write();

      if let Err(e) = cache.open_notebook(params.notebook_document) {
        otel::error!(
          "notebook_open_failed",
          format!("Failed to open notebook {}: {:?}", notebook_uri, e)
        );
        return;
      }

      for cell in params.cell_text_documents {
        match cache.upsert_with_version(
          cell.uri.clone(),
          cell.text,
          cell.version,
          client_id,
        ) {
          | Ok(_) => {
            cache.mark_file_open(cell.uri, client_id);
          },
          | Err(e) => {
            otel::error!(
              "notebook_cell_upsert_failed",
              format!("Failed to upsert notebook cell {}: {:?}", cell.uri, e)
            );
          },
        }
      }
    }
  }
  const NOTEBOOK_DID_OPEN_LANE: crate::scheduler::lanes::Lane =
    crate::scheduler::lanes::DEFAULT_LANE;

  /// The [`notebookDocument/didChange`] notification is sent from the client to
  /// the server when a notebook document changes. It is only sent for
  /// notebooks selected by the `notebookDocumentSync` server capability.
  ///
  /// [`notebookDocument/didChange`]: https://microsoft.github.io/language-server-protocol/specification#notebookDocument_didChange
  fn notebook_did_change(
    &self,
    params: DidChangeNotebookDocumentParams,
    ctx: &mut TaskContext<P, T>,
    writer: &mut PartitionWriteContextRef<'_, P>,
  ) -> impl std::future::Future<Output = ()> + Send {
    let client_id = ctx.client_id();
    let encoding = ctx.position_encoding();

    async move {
      let _ = writer;
      let notebook_uri = params.notebook_document.uri.clone();
      let version = params.notebook_document.version;

      let source_cache = ctx.source_cache();
      let mut cache = source_cache.write();

      {
        let mut notebook = match cache.get_notebook_mut(&notebook_uri) {
          | Some(nb) => nb,
          | None => {
            return;
          },
        };

        notebook.version = version;

        if let Some(metadata) = params.change.metadata {
          notebook.metadata = Some(metadata);
        }
      }

      if let Some(cells_change) = params.change.cells {
        if let Some(structure) = cells_change.structure {
          let start = structure.array.start as usize;
          let delete_count = structure.array.delete_count as usize;

          let mut removed_cells = Vec::new();
          {
            let Some(mut notebook) = cache.get_notebook_mut(&notebook_uri) else {
              otel::error!(
                "notebook_not_found",
                format!("Notebook {} not found during cell removal", notebook_uri)
              );
              return;
            };
            for _i in 0..delete_count {
              if start < notebook.cells.len() {
                removed_cells.push(notebook.cells.remove(start));
              }
            }
          }

          drop(removed_cells);

          if let Some(new_cells) = structure.array.cells {
            let Some(mut notebook) = cache.get_notebook_mut(&notebook_uri) else {
              otel::error!(
                "notebook_not_found",
                format!("Notebook {} not found during cell insertion", notebook_uri)
              );
              return;
            };
            for (i, cell) in new_cells.into_iter().enumerate() {
              notebook.cells.insert(start + i, cell);
            }
          }

          if let Some(opened_cells) = structure.did_open {
            for cell_doc in opened_cells {
              match cache.upsert_with_version(
                cell_doc.uri.clone(),
                cell_doc.text,
                cell_doc.version,
                client_id,
              ) {
                | Ok(_) => {
                  cache.mark_file_open(cell_doc.uri, client_id);
                },
                | Err(e) => {
                  otel::error!(
                    "notebook_cell_upsert_failed",
                    format!(
                      "Failed to upsert newly opened notebook cell {}: {:?}",
                      cell_doc.uri, e
                    )
                  );
                },
              }
            }
          }

          if let Some(closed_cells) = structure.did_close {
            for cell_id in closed_cells {
              cache.mark_file_closed(&cell_id.uri, client_id);
            }
          }
        }

        if let Some(data_changes) = cells_change.data {
          let Some(mut notebook) = cache.get_notebook_mut(&notebook_uri) else {
            otel::error!(
              "notebook_not_found",
              format!("Notebook {} not found during data change", notebook_uri)
            );
            return;
          };
          for updated_cell in data_changes {
            if let Some(cell) = notebook
              .cells
              .iter_mut()
              .find(|c| c.document == updated_cell.document)
            {
              *cell = updated_cell;
            }
          }
        }

        if let Some(text_changes) = cells_change.text_content {
          for text_change in text_changes {
            if let Err(e) = cache.apply_content_changes(
              &text_change.document.uri,
              text_change.document.version,
              text_change.changes,
              &encoding,
              client_id,
            ) {
              otel::error!(
                "notebook_cell_content_changes_failed",
                format!(
                  "Failed to apply content changes to notebook cell {}: {:?}",
                  text_change.document.uri, e
                )
              );
            }
          }
        }
      }
    }
  }
  const NOTEBOOK_DID_CHANGE_LANE: crate::scheduler::lanes::Lane =
    crate::scheduler::lanes::DEFAULT_LANE;

  /// The [`notebookDocument/didSave`] notification is sent from the client to
  /// the server when a notebook document is saved. It is only sent for
  /// notebooks selected by the `notebookDocumentSync` server capability.
  ///
  /// [`notebookDocument/didSave`]: https://microsoft.github.io/language-server-protocol/specification#notebookDocument_didSave
  fn notebook_did_save(
    &self,
    params: DidSaveNotebookDocumentParams,
    ctx: &mut TaskContext<P, T>,
    writer: &mut PartitionWriteContextRef<'_, P>,
  ) -> impl std::future::Future<Output = ()> + Send {
    async move {
      let _ = ctx;
      let _ = writer;
      let notebook_uri = params.notebook_document.uri;
    }
  }
  const NOTEBOOK_DID_SAVE_LANE: crate::scheduler::lanes::Lane =
    crate::scheduler::lanes::DEFAULT_LANE;

  /// The [`notebookDocument/didClose`] notification is sent from the client to
  /// the server when a notebook document is closed. It is only sent for
  /// notebooks selected by the `notebookDocumentSync` server capability.
  ///
  /// [`notebookDocument/didClose`]: https://microsoft.github.io/language-server-protocol/specification#notebookDocument_didClose
  fn notebook_did_close(
    &self,
    params: DidCloseNotebookDocumentParams,
    ctx: &mut TaskContext<P, T>,
    writer: &mut PartitionWriteContextRef<'_, P>,
  ) -> impl std::future::Future<Output = ()> + Send {
    let client_id = ctx.client_id();

    async move {
      let _ = writer;
      let notebook_uri = params.notebook_document.uri;

      let source_cache = ctx.source_cache();
      let mut cache = source_cache.write();

      for cell_id in params.cell_text_documents {
        cache.mark_file_closed(&cell_id.uri, client_id);
      }

      if let Err(e) = cache.close_notebook(&notebook_uri) {
        otel::error!(
          "notebook_close_failed",
          format!("Failed to close notebook {}: {:?}", notebook_uri, e)
        );
      }
    }
  }
  const NOTEBOOK_DID_CLOSE_LANE: crate::scheduler::lanes::Lane =
    crate::scheduler::lanes::DEFAULT_LANE;
}