From c3d69f7c69b5c2ce9a6936e1a098ce9384115a3f Mon Sep 17 00:00:00 2001
From: Jane Lewis <jane@astral.sh>
Date: Mon, 15 Apr 2024 20:35:48 -0700
Subject: [PATCH] Rework code action support checker

---
 crates/ruff_server/src/server.rs              | 37 +++++++------------
 .../src/server/api/requests/code_action.rs    | 12 ++----
 .../api/requests/code_action_resolve.rs       | 18 +++++----
 3 files changed, 26 insertions(+), 41 deletions(-)

diff --git a/crates/ruff_server/src/server.rs b/crates/ruff_server/src/server.rs
index d148830b5789d..4c45ece7473c9 100644
--- a/crates/ruff_server/src/server.rs
+++ b/crates/ruff_server/src/server.rs
@@ -224,7 +224,7 @@ impl Server {
                 CodeActionOptions {
                     code_action_kinds: Some(
                         SupportedCodeAction::all()
-                            .flat_map(|action| action.kinds().iter())
+                            .map(SupportedCodeAction::to_kind)
                             .cloned()
                             .collect(),
                     ),
@@ -285,22 +285,24 @@ pub(crate) enum SupportedCodeAction {
 }
 
 impl SupportedCodeAction {
-    /// Returns the possible LSP code action kind(s) that map to this code action.
-    fn kinds(self) -> &'static [CodeActionKind] {
-        static QUICKFIX: [CodeActionKind; 1] = [CodeActionKind::QUICKFIX];
-        static SOURCE_FIX_ALL: [CodeActionKind; 2] =
-            [CodeActionKind::SOURCE_FIX_ALL, crate::SOURCE_FIX_ALL_RUFF];
-        static SOURCE_ORGANIZE_IMPORTS: [CodeActionKind; 2] = [
-            CodeActionKind::SOURCE_ORGANIZE_IMPORTS,
-            crate::SOURCE_ORGANIZE_IMPORTS_RUFF,
-        ];
+    /// Returns the LSP code action kind that map to this code action.
+    fn to_kind(self) -> &'static CodeActionKind {
+        static QUICK_FIX: CodeActionKind = CodeActionKind::QUICKFIX;
+        static SOURCE_FIX_ALL: CodeActionKind = crate::SOURCE_FIX_ALL_RUFF;
+        static SOURCE_ORGANIZE_IMPORTS: CodeActionKind = crate::SOURCE_ORGANIZE_IMPORTS_RUFF;
         match self {
-            Self::QuickFix => &QUICKFIX,
+            Self::QuickFix => &QUICK_FIX,
             Self::SourceFixAll => &SOURCE_FIX_ALL,
             Self::SourceOrganizeImports => &SOURCE_ORGANIZE_IMPORTS,
         }
     }
 
+    fn from_kind(kind: CodeActionKind) -> impl Iterator<Item = Self> {
+        Self::all().filter(move |supported_kind| {
+            supported_kind.to_kind().as_str().starts_with(kind.as_str())
+        })
+    }
+
     /// Returns all code actions kinds that the server currently supports.
     fn all() -> impl Iterator<Item = Self> {
         [
@@ -311,16 +313,3 @@ impl SupportedCodeAction {
         .into_iter()
     }
 }
-
-impl TryFrom<CodeActionKind> for SupportedCodeAction {
-    type Error = ();
-
-    fn try_from(kind: CodeActionKind) -> std::result::Result<Self, Self::Error> {
-        for supported_kind in Self::all() {
-            if supported_kind.kinds().contains(&kind) {
-                return Ok(supported_kind);
-            }
-        }
-        Err(())
-    }
-}
diff --git a/crates/ruff_server/src/server/api/requests/code_action.rs b/crates/ruff_server/src/server/api/requests/code_action.rs
index 3af967f137593..96eb297f55280 100644
--- a/crates/ruff_server/src/server/api/requests/code_action.rs
+++ b/crates/ruff_server/src/server/api/requests/code_action.rs
@@ -166,14 +166,8 @@ fn supported_code_actions(
         return SupportedCodeAction::all().collect();
     };
 
-    SupportedCodeAction::all()
-        .filter(move |action| {
-            action_filter.iter().any(|filter| {
-                action
-                    .kinds()
-                    .iter()
-                    .any(|kind| kind.as_str().starts_with(filter.as_str()))
-            })
-        })
+    action_filter
+        .into_iter()
+        .flat_map(SupportedCodeAction::from_kind)
         .collect()
 }
diff --git a/crates/ruff_server/src/server/api/requests/code_action_resolve.rs b/crates/ruff_server/src/server/api/requests/code_action_resolve.rs
index c752a10827383..f9c47267d9f9e 100644
--- a/crates/ruff_server/src/server/api/requests/code_action_resolve.rs
+++ b/crates/ruff_server/src/server/api/requests/code_action_resolve.rs
@@ -30,14 +30,16 @@ impl super::BackgroundDocumentRequestHandler for CodeActionResolve {
     ) -> Result<types::CodeAction> {
         let document = snapshot.document();
 
-        let action_kind: SupportedCodeAction = action
-            .kind
-            .clone()
-            .ok_or(anyhow::anyhow!("No kind was given for code action"))
-            .with_failure_code(ErrorCode::InvalidParams)?
-            .try_into()
-            .map_err(|()| anyhow::anyhow!("Code action was of an invalid kind"))
-            .with_failure_code(ErrorCode::InvalidParams)?;
+        let action_kind = SupportedCodeAction::from_kind(
+            action
+                .kind
+                .clone()
+                .ok_or(anyhow::anyhow!("No kind was given for code action"))
+                .with_failure_code(ErrorCode::InvalidParams)?,
+        )
+        .next()
+        .ok_or(anyhow::anyhow!("Code action was of an invalid kind"))
+        .with_failure_code(ErrorCode::InvalidParams)?;
 
         action.edit = match action_kind {
             SupportedCodeAction::SourceFixAll => Some(