From 4c740f394afb66a69093e1b093e19b20a72fffcb Mon Sep 17 00:00:00 2001
From: Sijawusz Pur Rahnama <sija@sija.pl>
Date: Thu, 8 Jun 2023 01:58:58 +0200
Subject: [PATCH] Implement `Scope#visibility`

---
 spec/ameba/ast/visitors/scope_visitor_spec.cr | 37 +++++++++++++++++++
 spec/spec_helper.cr                           |  3 ++
 src/ameba/ast/scope.cr                        |  8 ++++
 src/ameba/ast/visitors/scope_visitor.cr       | 15 +++++++-
 4 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/spec/ameba/ast/visitors/scope_visitor_spec.cr b/spec/ameba/ast/visitors/scope_visitor_spec.cr
index a15402ee..b4f2432c 100644
--- a/spec/ameba/ast/visitors/scope_visitor_spec.cr
+++ b/spec/ameba/ast/visitors/scope_visitor_spec.cr
@@ -2,6 +2,17 @@ require "../../../spec_helper"
 
 module Ameba::AST
   describe ScopeVisitor do
+    {% for type in %w[class module enum].map(&.id) %}
+      it "creates a scope for the {{ type }} def" do
+        rule = ScopeRule.new
+        ScopeVisitor.new rule, Source.new <<-CRYSTAL
+          {{ type }} Foo
+          end
+          CRYSTAL
+        rule.scopes.size.should eq 1
+      end
+    {% end %}
+
     it "creates a scope for the def" do
       rule = ScopeRule.new
       ScopeVisitor.new rule, Source.new <<-CRYSTAL
@@ -54,5 +65,31 @@ module Ameba::AST
         outer_block.outer_scope.should be_nil
       end
     end
+
+    context "#visibility" do
+      it "is being properly set" do
+        rule = ScopeRule.new
+        ScopeVisitor.new rule, Source.new <<-CRYSTAL
+          private class Foo
+          end
+          CRYSTAL
+        rule.scopes.size.should eq 1
+        rule.scopes.first.visibility.should eq Crystal::Visibility::Private
+      end
+
+      it "is being inherited from the outer scope(s)" do
+        rule = ScopeRule.new
+        ScopeVisitor.new rule, Source.new <<-CRYSTAL
+          private class Foo
+            class Bar
+              def baz
+              end
+            end
+          end
+          CRYSTAL
+        rule.scopes.size.should eq 3
+        rule.scopes.each &.visibility.should eq Crystal::Visibility::Private
+      end
+    end
   end
 end
diff --git a/spec/spec_helper.cr b/spec/spec_helper.cr
index 4a18746b..365febd1 100644
--- a/spec/spec_helper.cr
+++ b/spec/spec_helper.cr
@@ -43,6 +43,9 @@ module Ameba
       description "Internal rule to test scopes"
     end
 
+    def test(source, node : Crystal::VisibilityModifier, scope : AST::Scope)
+    end
+
     def test(source, node : Crystal::ASTNode, scope : AST::Scope)
       @scopes << scope
     end
diff --git a/src/ameba/ast/scope.cr b/src/ameba/ast/scope.cr
index b74f8b62..3c343c6e 100644
--- a/src/ameba/ast/scope.cr
+++ b/src/ameba/ast/scope.cr
@@ -7,6 +7,9 @@ module Ameba::AST
     # Whether the scope yields.
     setter yields = false
 
+    # Scope visibility level
+    setter visibility : Crystal::Visibility?
+
     # Link to local variables
     getter variables = [] of Variable
 
@@ -172,6 +175,11 @@ module Ameba::AST
       false
     end
 
+    # Returns visibility of the current scope (could be inherited from the outer scope).
+    def visibility
+      @visibility || outer_scope.try(&.visibility)
+    end
+
     # Returns `true` if current scope is a def, `false` otherwise.
     def def?
       node.is_a?(Crystal::Def)
diff --git a/src/ameba/ast/visitors/scope_visitor.cr b/src/ameba/ast/visitors/scope_visitor.cr
index a5c578ae..023094a0 100644
--- a/src/ameba/ast/visitors/scope_visitor.cr
+++ b/src/ameba/ast/visitors/scope_visitor.cr
@@ -25,6 +25,7 @@ module Ameba::AST
     @scope_queue = [] of Scope
     @current_scope : Scope
     @current_assign : Crystal::ASTNode?
+    @visibility_modifier : Crystal::Visibility?
     @skip : Array(Crystal::ASTNode.class)?
 
     def initialize(@rule, @source, skip = nil)
@@ -36,12 +37,18 @@ module Ameba::AST
 
     private def on_scope_enter(node)
       return if skip?(node)
-      @current_scope = Scope.new(node, @current_scope)
+
+      scope = Scope.new(node, @current_scope)
+      scope.visibility = @visibility_modifier
+
+      @current_scope = scope
     end
 
     private def on_scope_end(node)
       @scope_queue << @current_scope
 
+      @visibility_modifier = nil
+
       # go up if this is not a top level scope
       return unless outer_scope = @current_scope.outer_scope
       @current_scope = outer_scope
@@ -64,6 +71,12 @@ module Ameba::AST
       end
     {% end %}
 
+    # :nodoc:
+    def visit(node : Crystal::VisibilityModifier)
+      @visibility_modifier = node.modifier
+      true
+    end
+
     # :nodoc:
     def visit(node : Crystal::Yield)
       @current_scope.yields = true