From c9ba487373de52cac4ea18e788c9c07ce2930dcd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?fn=20=E2=8C=83=20=E2=8C=A5?=
 <70830482+FnControlOption@users.noreply.github.com>
Date: Tue, 9 Nov 2021 09:26:00 -0800
Subject: [PATCH] Add custom `NodeVisitor` subclass for `Style/VariableNames`

---
 spec/ameba/rule/style/variable_names_spec.cr |  1 -
 src/ameba/rule/style/variable_names.cr       | 29 ++++++++++++--------
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/spec/ameba/rule/style/variable_names_spec.cr b/spec/ameba/rule/style/variable_names_spec.cr
index f150e8a9..2571d0e0 100644
--- a/spec/ameba/rule/style/variable_names_spec.cr
+++ b/spec/ameba/rule/style/variable_names_spec.cr
@@ -38,7 +38,6 @@ module Ameba
         class Greeting
           def initialize(@badNamed = nil)
                        # ^ error: Var name should be underscore-cased: @bad_named, not @badNamed
-                       # ^ error: Var name should be underscore-cased: bad_named, not badNamed
           end
         end
         CRYSTAL
diff --git a/src/ameba/rule/style/variable_names.cr b/src/ameba/rule/style/variable_names.cr
index 48d89805..2176dbec 100644
--- a/src/ameba/rule/style/variable_names.cr
+++ b/src/ameba/rule/style/variable_names.cr
@@ -35,17 +35,9 @@ module Ameba::Rule::Style
       issue_for node, MSG % {expected, node.name}
     end
 
-    # TODO: Handle special case where instance variable is method parameter.
-    #       For example, this:
-    #
-    #         def initialize(@foo)
-    #         end
-    #
-    #       is represented in the AST as:
-    #
-    #         def initialize(foo)
-    #           @foo = foo
-    #         end
+    def test(source : Source)
+      VarVisitor.new self, source
+    end
 
     def test(source, node : Crystal::Var)
       check_node source, node
@@ -58,5 +50,20 @@ module Ameba::Rule::Style
     def test(source, node : Crystal::ClassVar)
       check_node source, node
     end
+
+    private class VarVisitor < AST::NodeVisitor
+      private getter var_locations = [] of Crystal::Location
+
+      def visit(node : Crystal::Var)
+        !var_locations.includes?(node.location) && super
+      end
+
+      def visit(node : Crystal::InstanceVar | Crystal::ClassVar)
+        if (location = node.location)
+          var_locations << location
+        end
+        super
+      end
+    end
   end
 end