Support array and inline-def arguments to private_class_method

rust/rubydex/src/indexing/ruby_indexer.rs

@@ -1274,6 +1274,7 @@ impl<'a> RubyIndexer<'a> {
         }
     }
 
+    #[allow(clippy::too_many_lines)]
     fn handle_singleton_method_visibility(
         &mut self,
         node: &ruby_prism::CallNode,
@@ -1307,7 +1308,10 @@ impl<'a> RubyIndexer<'a> {
             return;
         };
 
-        for argument in &arguments.arguments() {
+        let args = arguments.arguments();
+        let arg_count = args.len();
+
+        for argument in &args {
             match argument {
                 ruby_prism::Node::SymbolNode { .. } | ruby_prism::Node::StringNode { .. } => {
                     self.create_method_visibility_definition(
@@ -1316,6 +1320,88 @@ impl<'a> RubyIndexer<'a> {
                         DefinitionFlags::SINGLETON_METHOD_VISIBILITY,
                     );
                 }
+                ruby_prism::Node::ArrayNode { .. } if arg_count == 1 => {
+                    let array = argument.as_array_node().unwrap();
+                    for element in &array.elements() {
+                        match element {
+                            ruby_prism::Node::SymbolNode { .. } | ruby_prism::Node::StringNode { .. } => {
+                                self.create_method_visibility_definition(
+                                    &element,
+                                    visibility,
+                                    DefinitionFlags::SINGLETON_METHOD_VISIBILITY,
+                                );
+                            }
+                            ruby_prism::Node::DefNode { .. } => {
+                                let def_node = element.as_def_node().unwrap();
+                                if def_node.receiver().is_none() {
+                                    self.local_graph.add_diagnostic(
+                                        Rule::InvalidMethodVisibility,
+                                        Offset::from_prism_location(&element.location()),
+                                        format!("`{call_name}` requires a singleton method definition"),
+                                    );
+                                    self.visit(&element);
+                                    continue;
+                                }
+                                let name_loc = def_node.name_loc();
+                                let name = Self::location_to_string(&name_loc);
+                                self.create_method_visibility_definition_from_name(
+                                    &name,
+                                    &name_loc,
+                                    visibility,
+                                    DefinitionFlags::SINGLETON_METHOD_VISIBILITY,
+                                );
+                                self.visit(&element);
+                            }
+                            _ => {
+                                self.local_graph.add_diagnostic(
+                                    Rule::InvalidMethodVisibility,
+                                    Offset::from_prism_location(&element.location()),
+                                    format!(
+                                        "`{call_name}` array element must be a Symbol, String, or method definition"
+                                    ),
+                                );
+                                self.visit(&element);
+                            }
+                        }
+                    }
+                }
+                ruby_prism::Node::DefNode { .. } => {
+                    let def_node = argument.as_def_node().unwrap();
+                    if def_node.receiver().is_none() {
+                        self.local_graph.add_diagnostic(
+                            Rule::InvalidMethodVisibility,
+                            Offset::from_prism_location(&argument.location()),
+                            format!("`{call_name}` requires a singleton method definition"),
+                        );
+                        self.visit(&argument);
+                        continue;
+                    }
+                    let name_loc = def_node.name_loc();
+                    let name = Self::location_to_string(&name_loc);
+                    self.create_method_visibility_definition_from_name(
+                        &name,
+                        &name_loc,
+                        visibility,
+                        DefinitionFlags::SINGLETON_METHOD_VISIBILITY,
+                    );
+                    self.visit(&argument);
+                }
+                arg if Self::is_attr_call(&arg) => {
+                    self.local_graph.add_diagnostic(
+                        Rule::InvalidMethodVisibility,
+                        Offset::from_prism_location(&arg.location()),
+                        format!("`{call_name}` does not accept `attr_*` arguments"),
+                    );
+                    self.visit(&arg);
+                }
+                ruby_prism::Node::ArrayNode { .. } => {
+                    self.local_graph.add_diagnostic(
+                        Rule::InvalidMethodVisibility,
+                        Offset::from_prism_location(&argument.location()),
+                        format!("`{call_name}` array argument must be the only argument"),
+                    );
+                    self.visit(&argument);
+                }
                 _ => {
                     self.local_graph.add_diagnostic(
                         Rule::InvalidMethodVisibility,
@@ -1390,11 +1476,8 @@ impl<'a> RubyIndexer<'a> {
         let (name, location) = match arg {
             ruby_prism::Node::SymbolNode { .. } => {
                 let symbol = arg.as_symbol_node().unwrap();
-                if let Some(value_loc) = symbol.value_loc() {
-                    (Self::location_to_string(&value_loc), value_loc)
-                } else {
-                    return;
-                }
+                let Some(value_loc) = symbol.value_loc() else { return };
+                (Self::location_to_string(&value_loc), value_loc)
             }
             ruby_prism::Node::StringNode { .. } => {
                 let string = arg.as_string_node().unwrap();
@@ -1404,8 +1487,18 @@ impl<'a> RubyIndexer<'a> {
             _ => return,
         };
 
+        self.create_method_visibility_definition_from_name(&name, &location, visibility, flags);
+    }
+
+    fn create_method_visibility_definition_from_name(
+        &mut self,
+        name: &str,
+        location: &ruby_prism::Location,
+        visibility: Visibility,
+        flags: DefinitionFlags,
+    ) {
         let str_id = self.local_graph.intern_string(format!("{name}()"));
-        let arg_offset = Offset::from_prism_location(&location);
+        let arg_offset = Offset::from_prism_location(location);
         let definition = Definition::MethodVisibility(Box::new(MethodVisibilityDefinition::new(
             str_id,
             visibility,

rust/rubydex/src/indexing/ruby_indexer_tests.rs

@@ -2237,6 +2237,9 @@ mod visibility_tests {
 
             module Foo
               private_class_method NOT_INDEXED
+              attr_reader :a_attr_target
+              private_class_method attr_reader(:bad)
+              private_class_method def inline; end
             end
             ",
         );
@@ -2247,6 +2250,8 @@ mod visibility_tests {
                 "invalid-method-visibility: `private_class_method` called at top level (1:1-1:34)",
                 "invalid-method-visibility: `private_class_method` called at top level (2:1-2:39)",
                 "invalid-method-visibility: `private_class_method` called with a non-literal argument (6:24-6:35)",
+                "invalid-method-visibility: `private_class_method` does not accept `attr_*` arguments (8:24-8:41)",
+                "invalid-method-visibility: `private_class_method` requires a singleton method definition (9:24-9:39)",
             ]
         );
     }
@@ -2306,6 +2311,159 @@ mod visibility_tests {
         );
         assert_constant_references_eq!(&context, ["NESTED_REF"]);
     }
+
+    #[test]
+    fn index_private_class_method_inline_def() {
+        let context = index_source(
+            r"
+            class Foo
+              private_class_method def self.inline; end
+            end
+            ",
+        );
+
+        assert_no_local_diagnostics!(&context);
+
+        assert_definition_at!(&context, "2:33-2:39", MethodVisibility, |def| {
+            assert!(def.flags().is_singleton_method_visibility());
+            assert_string_eq!(&context, def.str_id(), "inline()");
+            assert_eq!(def.visibility(), &Visibility::Private);
+        });
+    }
+
+    #[test]
+    fn index_private_class_method_array_form() {
+        let context = index_source(
+            r#"
+            class Foo
+              def self.flat; end
+              def self.flat2; end
+              def self.mixed; end
+
+              private_class_method [:flat, :flat2]
+              public_class_method [:mixed, "flat"]
+              private_class_method [def self.dyn; end]
+            end
+            "#,
+        );
+
+        assert_no_local_diagnostics!(&context);
+
+        assert_definition_at!(&context, "6:26-6:30", MethodVisibility, |def| {
+            assert!(def.flags().is_singleton_method_visibility());
+            assert_string_eq!(&context, def.str_id(), "flat()");
+            assert_eq!(def.visibility(), &Visibility::Private);
+        });
+        assert_definition_at!(&context, "6:33-6:38", MethodVisibility, |def| {
+            assert!(def.flags().is_singleton_method_visibility());
+            assert_string_eq!(&context, def.str_id(), "flat2()");
+            assert_eq!(def.visibility(), &Visibility::Private);
+        });
+        assert_definition_at!(&context, "7:25-7:30", MethodVisibility, |def| {
+            assert!(def.flags().is_singleton_method_visibility());
+            assert_string_eq!(&context, def.str_id(), "mixed()");
+            assert_eq!(def.visibility(), &Visibility::Public);
+        });
+        assert_definition_at!(&context, "7:32-7:38", MethodVisibility, |def| {
+            assert!(def.flags().is_singleton_method_visibility());
+            assert_string_eq!(&context, def.str_id(), "flat()");
+            assert_eq!(def.visibility(), &Visibility::Public);
+        });
+        assert_definition_at!(&context, "8:34-8:37", MethodVisibility, |def| {
+            assert!(def.flags().is_singleton_method_visibility());
+            assert_string_eq!(&context, def.str_id(), "dyn()");
+            assert_eq!(def.visibility(), &Visibility::Private);
+        });
+    }
+
+    #[test]
+    fn index_private_class_method_array_continues_past_invalid_element() {
+        let context = index_source(
+            r"
+            class Foo
+              def self.flat; end
+              def self.later; end
+
+              private_class_method [:flat, SOME_CONST, :later]
+            end
+            ",
+        );
+
+        assert_local_diagnostics_eq!(
+            &context,
+            vec![
+                "invalid-method-visibility: `private_class_method` array element must be a Symbol, String, or method definition (5:32-5:42)"
+            ]
+        );
+
+        assert_definition_at!(&context, "5:26-5:30", MethodVisibility, |def| {
+            assert!(def.flags().is_singleton_method_visibility());
+            assert_string_eq!(&context, def.str_id(), "flat()");
+            assert_eq!(def.visibility(), &Visibility::Private);
+        });
+        assert_definition_at!(&context, "5:45-5:50", MethodVisibility, |def| {
+            assert!(def.flags().is_singleton_method_visibility());
+            assert_string_eq!(&context, def.str_id(), "later()");
+            assert_eq!(def.visibility(), &Visibility::Private);
+        });
+    }
+
+    #[test]
+    fn index_private_class_method_array_rejects_receiverless_def() {
+        let context = index_source(
+            r"
+            class Foo
+              private_class_method [def instance_method; end]
+            end
+            ",
+        );
+
+        assert_local_diagnostics_eq!(
+            &context,
+            vec![
+                "invalid-method-visibility: `private_class_method` requires a singleton method definition (2:25-2:49)"
+            ]
+        );
+
+        for def in context.graph().definitions().values() {
+            assert!(
+                !matches!(def, Definition::MethodVisibility(d) if d.flags().is_singleton_method_visibility()),
+                "should not record visibility for receiverless def in array"
+            );
+        }
+
+        assert_definition_at!(&context, "2:25-2:49", Method, |def| {
+            assert_def_str_eq!(&context, def, "instance_method()");
+            assert!(def.receiver().is_none());
+        });
+    }
+
+    #[test]
+    fn index_private_class_method_array_not_sole_arg_diagnostic() {
+        let context = index_source(
+            r"
+            class Foo
+              def self.a; end
+              def self.b; end
+
+              private_class_method [:a], :b
+            end
+            ",
+        );
+
+        assert_local_diagnostics_eq!(
+            &context,
+            vec![
+                "invalid-method-visibility: `private_class_method` array argument must be the only argument (5:24-5:28)"
+            ]
+        );
+
+        assert_definition_at!(&context, "5:31-5:32", MethodVisibility, |def| {
+            assert!(def.flags().is_singleton_method_visibility());
+            assert_string_eq!(&context, def.str_id(), "b()");
+            assert_eq!(def.visibility(), &Visibility::Private);
+        });
+    }
 }
 
 mod attr_accessor_tests {

rust/rubydex/src/resolution_tests.rs

@@ -5460,4 +5460,45 @@ mod visibility_resolution_tests {
         assert_no_diagnostics!(&context);
         assert_visibility_eq!(context, "Foo::<Foo>#missing()", Visibility::Private);
     }
+
+    #[test]
+    fn retroactive_singleton_method_visibility_inline_def() {
+        let mut context = GraphTest::new();
+        context.index_uri(
+            "file:///foo.rb",
+            r"
+            class Foo
+              private_class_method def self.bar; end
+            end
+            ",
+        );
+        context.resolve();
+
+        assert_no_diagnostics!(&context);
+        assert_visibility_eq!(context, "Foo::<Foo>#bar()", Visibility::Private);
+    }
+
+    #[test]
+    fn retroactive_singleton_method_visibility_array_form() {
+        let mut context = GraphTest::new();
+        context.index_uri(
+            "file:///foo.rb",
+            r#"
+            class Foo
+              def self.a; end
+              def self.b; end
+              def self.c; end
+
+              private_class_method [:a, "b"]
+              public_class_method [:c]
+            end
+            "#,
+        );
+        context.resolve();
+
+        assert_no_diagnostics!(&context);
+        assert_visibility_eq!(context, "Foo::<Foo>#a()", Visibility::Private);
+        assert_visibility_eq!(context, "Foo::<Foo>#b()", Visibility::Private);
+        assert_visibility_eq!(context, "Foo::<Foo>#c()", Visibility::Public);
+    }
 }