From 09fdac6be987fada88b0c1eaf74f63ebeb4d9fe8 Mon Sep 17 00:00:00 2001
From: Sijawusz Pur Rahnama <sija@sija.pl>
Date: Mon, 29 May 2023 16:33:45 +0200
Subject: [PATCH] Refactor `Lint::Documentation` rule to use a custom visitor

---
 src/ameba/rule/lint/documentation.cr | 32 ++++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/ameba/rule/lint/documentation.cr b/src/ameba/rule/lint/documentation.cr
index efd7d58a..534181a0 100644
--- a/src/ameba/rule/lint/documentation.cr
+++ b/src/ameba/rule/lint/documentation.cr
@@ -22,10 +22,10 @@ module Ameba::Rule::Lint
     MSG = "Missing documentation"
 
     def test(source)
-      AST::ScopeVisitor.new self, source
+      DocumentationVisitor.new self, source
     end
 
-    def test(source, node : Crystal::ClassDef | Crystal::ModuleDef | Crystal::EnumDef | Crystal::Def | Crystal::Macro, scope : AST::Scope)
+    def test(source, node)
       return unless node.visibility.public?
 
       case node
@@ -38,5 +38,33 @@ module Ameba::Rule::Lint
 
       issue_for(node, MSG) unless node.doc.presence
     end
+
+    # :nodoc:
+    private class DocumentationVisitor < AST::BaseVisitor
+      NODES = {
+        ClassDef,
+        ModuleDef,
+        EnumDef,
+        Def,
+        Macro,
+      }
+
+      @visibility : Crystal::Visibility = :public
+
+      def visit(node : Crystal::VisibilityModifier)
+        @visibility = node.modifier
+        true
+      end
+
+      {% for name in NODES %}
+        def visit(node : Crystal::{{ name }})
+          node.visibility = @visibility
+          @visibility = :public
+
+          @rule.test @source, node
+          true
+        end
+      {% end %}
+    end
   end
 end