laburnum 1.17.1

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

//! Call hierarchy and type hierarchy tools.
//!
//! Each of these chains an LSP `prepare*` step with a follow-up
//! (`incomingCalls` / `outgoingCalls` / `supertypes` / `subtypes`) so the
//! agent sees a single tool, not two. The prepare-step items are included
//! in the result so the agent can correlate calls/types to their target
//! definitions.

use {
  crate::{
    connect::lsp::{
      LspClient,
      request::{
        CallHierarchyIncomingCalls,
        CallHierarchyOutgoingCalls,
        CallHierarchyPrepare,
        TypeHierarchyPrepare,
        TypeHierarchySubtypes,
        TypeHierarchySupertypes,
      },
    },
    connect::mcp::{
      schema,
      tool::{
        Tool,
        ToolError,
        ToolRegistry,
      },
    },
    protocol::lsp::{
      CallHierarchyIncomingCall,
      CallHierarchyIncomingCallsParams,
      CallHierarchyItem,
      CallHierarchyOutgoingCall,
      CallHierarchyOutgoingCallsParams,
      CallHierarchyPrepareParams,
      TextDocumentPositionParams,
      TypeHierarchyItem,
      TypeHierarchyPrepareParams,
      TypeHierarchySubtypesParams,
      TypeHierarchySupertypesParams,
    },
  },
  serde::Serialize,
};

pub fn register(registry: &mut ToolRegistry) {
  registry.register::<CallHierarchyIncoming>();
  registry.register::<CallHierarchyOutgoing>();
  registry.register::<TypeHierarchySupertypesTool>();
  registry.register::<TypeHierarchySubtypesTool>();
}

// -- call hierarchy ----------------------------------------------------------

#[derive(Serialize)]
pub struct CallHierarchyIncomingEntry {
  pub item:  CallHierarchyItem,
  pub calls: Vec<CallHierarchyIncomingCall>,
}

#[derive(Serialize)]
pub struct CallHierarchyIncomingResult {
  pub items: Vec<CallHierarchyIncomingEntry>,
}

pub enum CallHierarchyIncoming {}

impl Tool for CallHierarchyIncoming {
  type Input = TextDocumentPositionParams;
  type Output = CallHierarchyIncomingResult;

  const NAME: &'static str = "call_hierarchy_incoming";
  const DESCRIPTION: &'static str =
    "Find callers of the function at the given position. Chains LSP \
     `textDocument/prepareCallHierarchy` and `callHierarchy/incomingCalls`.";

  fn input_schema() -> serde_json::Value {
    schema::text_document_position()
  }

  async fn call(
    client: &LspClient,
    input: Self::Input,
  ) -> Result<Self::Output, ToolError> {
    let items = client
      .send_request::<CallHierarchyPrepare>(CallHierarchyPrepareParams {
        text_document_position_params: input,
        work_done_progress_params:     Default::default(),
      })
      .await?
      .unwrap_or_default();

    let mut entries = Vec::with_capacity(items.len());
    for item in items {
      let calls = client
        .send_request::<CallHierarchyIncomingCalls>(
          CallHierarchyIncomingCallsParams {
            item:                      item.clone(),
            work_done_progress_params: Default::default(),
            partial_result_params:     Default::default(),
          },
        )
        .await?
        .unwrap_or_default();
      entries.push(CallHierarchyIncomingEntry { item, calls });
    }
    Ok(CallHierarchyIncomingResult { items: entries })
  }
}

#[derive(Serialize)]
pub struct CallHierarchyOutgoingEntry {
  pub item:  CallHierarchyItem,
  pub calls: Vec<CallHierarchyOutgoingCall>,
}

#[derive(Serialize)]
pub struct CallHierarchyOutgoingResult {
  pub items: Vec<CallHierarchyOutgoingEntry>,
}

pub enum CallHierarchyOutgoing {}

impl Tool for CallHierarchyOutgoing {
  type Input = TextDocumentPositionParams;
  type Output = CallHierarchyOutgoingResult;

  const NAME: &'static str = "call_hierarchy_outgoing";
  const DESCRIPTION: &'static str =
    "Find functions called by the function at the given position. Chains \
     LSP `textDocument/prepareCallHierarchy` and \
     `callHierarchy/outgoingCalls`.";

  fn input_schema() -> serde_json::Value {
    schema::text_document_position()
  }

  async fn call(
    client: &LspClient,
    input: Self::Input,
  ) -> Result<Self::Output, ToolError> {
    let items = client
      .send_request::<CallHierarchyPrepare>(CallHierarchyPrepareParams {
        text_document_position_params: input,
        work_done_progress_params:     Default::default(),
      })
      .await?
      .unwrap_or_default();

    let mut entries = Vec::with_capacity(items.len());
    for item in items {
      let calls = client
        .send_request::<CallHierarchyOutgoingCalls>(
          CallHierarchyOutgoingCallsParams {
            item:                      item.clone(),
            work_done_progress_params: Default::default(),
            partial_result_params:     Default::default(),
          },
        )
        .await?
        .unwrap_or_default();
      entries.push(CallHierarchyOutgoingEntry { item, calls });
    }
    Ok(CallHierarchyOutgoingResult { items: entries })
  }
}

// -- type hierarchy ----------------------------------------------------------

#[derive(Serialize)]
pub struct TypeHierarchyEntry {
  pub item:  TypeHierarchyItem,
  pub types: Vec<TypeHierarchyItem>,
}

#[derive(Serialize)]
pub struct TypeHierarchyResult {
  pub items: Vec<TypeHierarchyEntry>,
}

pub enum TypeHierarchySupertypesTool {}

impl Tool for TypeHierarchySupertypesTool {
  type Input = TextDocumentPositionParams;
  type Output = TypeHierarchyResult;

  const NAME: &'static str = "type_hierarchy_supertypes";
  const DESCRIPTION: &'static str =
    "Find supertypes of the type at the given position. Chains LSP \
     `textDocument/prepareTypeHierarchy` and `typeHierarchy/supertypes`.";

  fn input_schema() -> serde_json::Value {
    schema::text_document_position()
  }

  async fn call(
    client: &LspClient,
    input: Self::Input,
  ) -> Result<Self::Output, ToolError> {
    let items = client
      .send_request::<TypeHierarchyPrepare>(TypeHierarchyPrepareParams {
        text_document_position_params: input,
        work_done_progress_params:     Default::default(),
      })
      .await?
      .unwrap_or_default();

    let mut entries = Vec::with_capacity(items.len());
    for item in items {
      let types = client
        .send_request::<TypeHierarchySupertypes>(TypeHierarchySupertypesParams {
          item:                      item.clone(),
          work_done_progress_params: Default::default(),
          partial_result_params:     Default::default(),
        })
        .await?
        .unwrap_or_default();
      entries.push(TypeHierarchyEntry { item, types });
    }
    Ok(TypeHierarchyResult { items: entries })
  }
}

pub enum TypeHierarchySubtypesTool {}

impl Tool for TypeHierarchySubtypesTool {
  type Input = TextDocumentPositionParams;
  type Output = TypeHierarchyResult;

  const NAME: &'static str = "type_hierarchy_subtypes";
  const DESCRIPTION: &'static str =
    "Find subtypes of the type at the given position. Chains LSP \
     `textDocument/prepareTypeHierarchy` and `typeHierarchy/subtypes`.";

  fn input_schema() -> serde_json::Value {
    schema::text_document_position()
  }

  async fn call(
    client: &LspClient,
    input: Self::Input,
  ) -> Result<Self::Output, ToolError> {
    let items = client
      .send_request::<TypeHierarchyPrepare>(TypeHierarchyPrepareParams {
        text_document_position_params: input,
        work_done_progress_params:     Default::default(),
      })
      .await?
      .unwrap_or_default();

    let mut entries = Vec::with_capacity(items.len());
    for item in items {
      let types = client
        .send_request::<TypeHierarchySubtypes>(TypeHierarchySubtypesParams {
          item:                      item.clone(),
          work_done_progress_params: Default::default(),
          partial_result_params:     Default::default(),
        })
        .await?
        .unwrap_or_default();
      entries.push(TypeHierarchyEntry { item, types });
    }
    Ok(TypeHierarchyResult { items: entries })
  }
}