From e97d20641c1be8203812214fd76b19fcee48082a Mon Sep 17 00:00:00 2001
From: Teppei Shintani <teppest1215@gmail.com>
Date: Mon, 10 Jun 2024 10:00:25 +0900
Subject: [PATCH] integrate DefineDataVisitor

---
 lib/rbs_inline_data/cli.rb               |   3 +-
 lib/rbs_inline_data/parser.rb            | 158 +++++++++++------------
 sig/generated/rbs_inline_data/parser.rbs |  35 +++--
 3 files changed, 101 insertions(+), 95 deletions(-)

diff --git a/lib/rbs_inline_data/cli.rb b/lib/rbs_inline_data/cli.rb
index d23fbcc..74f2d52 100644
--- a/lib/rbs_inline_data/cli.rb
+++ b/lib/rbs_inline_data/cli.rb
@@ -23,7 +23,8 @@ def run(args)
       targets.uniq!
 
       targets.each do |file|
-        definitions = Parser.parse(file)
+        result = Prism.parse_file(file.to_s)
+        definitions = Parser.parse(result)
         Writer.write(file, definitions)
       end
     end
diff --git a/lib/rbs_inline_data/parser.rb b/lib/rbs_inline_data/parser.rb
index 134f0c7..218669b 100644
--- a/lib/rbs_inline_data/parser.rb
+++ b/lib/rbs_inline_data/parser.rb
@@ -1,78 +1,7 @@
 require 'prism'
 
 module RbsInlineData
-  class Parser
-    class DefineDataVisitor < Prism::Visitor
-      # @rbs @definitions: Array[RbsInlineData::Parser::TypedDefinition]
-      # @rbs @surronding_class_or_module: Array[Symbol]
-
-      def initialize(definitions)
-        @definitions = definitions
-        @surronding_class_or_module = []
-      end
-
-      # @rbs override
-      def visit_class_node(node)
-        record_surrounding_class_or_module(node) { super }
-      end
-
-      # @rbs override
-      def visit_module_node(node)
-        record_surrounding_class_or_module(node) { super }
-      end
-
-      # @rbs override
-      def visit_constant_write_node(node)
-        if define_data?(node)
-          definition = extract_definition(node)
-          @definitions << definition if definition
-        end
-
-        super
-      end
-
-      private
-
-      #:: (Prism::ClassNode | Prism::ModuleNode) { (Prism::ClassNode | Prism::ModuleNode) -> void } -> void
-      def record_surrounding_class_or_module(node)
-        @surronding_class_or_module.push(node.constant_path.name)
-        yield(node)
-      ensure
-        @surronding_class_or_module.pop
-      end
-
-      #:: (Prism::ConstantWriteNode) -> bool
-      def define_data?(node)
-        node in {
-          value: Prism::CallNode[
-            receiver: (
-              Prism::ConstantReadNode[name: :Data] |
-              Prism::ConstantPathNode[parent: nil, name: :Data]
-            ),
-            name: :define,
-          ]
-        }
-      end
-
-      #:: (Prism::ConstantWriteNode) -> RbsInlineData::Parser::TypedDefinition?
-      def extract_definition(node)
-        source = node.slice
-        _, class_name, field_text = source.match(/\A([a-zA-Z]+) = Data\.define\(([\n\s\w\W]+)\)\z/).to_a
-        return nil if field_text.nil? || class_name.nil?
-
-        class_name = @surronding_class_or_module.join("::") + "::" + class_name
-
-        fields = field_text.split("\n").map(&:strip).map do |str|
-          str.match(/:(\w+), #:: ([\w\[\]]+)/)&.to_a
-        end.compact.map { |_, field_name, type| TypedField.new(field_name: field_name, type: type) }
-
-        TypedDefinition.new(
-          class_name: class_name,
-          fields: fields,
-        )
-      end
-    end
-
+  class Parser < Prism::Visitor
     # @rbs skip
     TypedDefinition = Data.define(
       :class_name, #:: String
@@ -84,26 +13,83 @@ def extract_definition(node)
       :type, #:: String
     )
 
-    # @rbs @file: Pathname
+    # @rbs @definitions: Array[RbsInlineData::Parser::TypedDefinition]
+    # @rbs @surronding_class_or_module: Array[Symbol]
 
-    #:: (Pathname) -> void
-    def initialize(file)
-      @file = file
+    #:: (Array[RbsInlineData::Parser::TypedDefinition]) -> void
+    def initialize(definitions)
+      @definitions = definitions
+      @surronding_class_or_module = []
     end
 
-    #:: (Pathname) -> Array[RbsInlineData::Parser::TypedDefinition]
-    def self.parse(file)
-      new(file).parse
-    end
-
-    #:: () -> Array[RbsInlineData::Parser::TypedDefinition]
-    def parse
-      result = Prism.parse_file(@file.to_s)
-
+    #:: (Prism::ParseResult) -> Array[RbsInlineData::Parser::TypedDefinition]
+    def self.parse(result)
       # @type var definitions: Array[RbsInlineData::Parser::TypedDefinition]
       definitions = []
-      result.value.accept(DefineDataVisitor.new(definitions))
+      instance = self.new(definitions)
+      result.value.accept(instance)
       definitions
     end
+
+    # @rbs override
+    def visit_class_node(node)
+      record_surrounding_class_or_module(node) { super }
+    end
+
+    # @rbs override
+    def visit_module_node(node)
+      record_surrounding_class_or_module(node) { super }
+    end
+
+    # @rbs override
+    def visit_constant_write_node(node)
+      if define_data?(node)
+        definition = extract_definition(node)
+        @definitions << definition if definition
+      end
+
+      super
+    end
+
+    private
+
+    #:: (Prism::ClassNode | Prism::ModuleNode) { (Prism::ClassNode | Prism::ModuleNode) -> void } -> void
+    def record_surrounding_class_or_module(node)
+      @surronding_class_or_module.push(node.constant_path.name)
+      yield(node)
+    ensure
+      @surronding_class_or_module.pop
+    end
+
+    #:: (Prism::ConstantWriteNode) -> bool
+    def define_data?(node)
+      node in {
+        value: Prism::CallNode[
+          receiver: (
+            Prism::ConstantReadNode[name: :Data] |
+            Prism::ConstantPathNode[parent: nil, name: :Data]
+          ),
+          name: :define,
+        ]
+      }
+    end
+
+    #:: (Prism::ConstantWriteNode) -> RbsInlineData::Parser::TypedDefinition?
+    def extract_definition(node)
+      source = node.slice
+      _, class_name, field_text = source.match(/\A([a-zA-Z]+) = Data\.define\(([\n\s\w\W]+)\)\z/).to_a
+      return nil if field_text.nil? || class_name.nil?
+
+      class_name = @surronding_class_or_module.join("::") + "::" + class_name
+
+      fields = field_text.split("\n").map(&:strip).map do |str|
+        str.match(/:(\w+), #:: ([\w\[\]]+)/)&.to_a
+      end.compact.map { |_, field_name, type| TypedField.new(field_name: field_name, type: type) }
+
+      TypedDefinition.new(
+        class_name: class_name,
+        fields: fields,
+      )
+    end
   end
 end
diff --git a/sig/generated/rbs_inline_data/parser.rbs b/sig/generated/rbs_inline_data/parser.rbs
index 4d91466..99b610f 100644
--- a/sig/generated/rbs_inline_data/parser.rbs
+++ b/sig/generated/rbs_inline_data/parser.rbs
@@ -1,7 +1,7 @@
 # Generated from lib/rbs_inline_data/parser.rb with RBS::Inline
 
 module RbsInlineData
-  class Parser
+  class Parser < Prism::Visitor
     class DefineDataVisitor < Prism::Visitor
       @definitions: Array[RbsInlineData::Parser::TypedDefinition]
 
@@ -30,15 +30,34 @@ module RbsInlineData
       def extract_definition: (Prism::ConstantWriteNode) -> RbsInlineData::Parser::TypedDefinition?
     end
 
-    @file: Pathname
+    @definitions: Array[RbsInlineData::Parser::TypedDefinition]
 
-    # :: (Pathname) -> void
-    def initialize: (Pathname) -> void
+    @surronding_class_or_module: Array[Symbol]
 
-    # :: (Pathname) -> Array[RbsInlineData::Parser::TypedDefinition]
-    def self.parse: (Pathname) -> Array[RbsInlineData::Parser::TypedDefinition]
+    # :: (Array[RbsInlineData::Parser::TypedDefinition]) -> void
+    def initialize: (Array[RbsInlineData::Parser::TypedDefinition]) -> void
 
-    # :: () -> Array[RbsInlineData::Parser::TypedDefinition]
-    def parse: () -> Array[RbsInlineData::Parser::TypedDefinition]
+    # :: (Prism::ParseResult) -> Array[RbsInlineData::Parser::TypedDefinition]
+    def self.parse: (Prism::ParseResult) -> Array[RbsInlineData::Parser::TypedDefinition]
+
+    # @rbs override
+    def visit_class_node: ...
+
+    # @rbs override
+    def visit_module_node: ...
+
+    # @rbs override
+    def visit_constant_write_node: ...
+
+    private
+
+    # :: (Prism::ClassNode | Prism::ModuleNode) { (Prism::ClassNode | Prism::ModuleNode) -> void } -> void
+    def record_surrounding_class_or_module: (Prism::ClassNode | Prism::ModuleNode) { (Prism::ClassNode | Prism::ModuleNode) -> void } -> void
+
+    # :: (Prism::ConstantWriteNode) -> bool
+    def define_data?: (Prism::ConstantWriteNode) -> bool
+
+    # :: (Prism::ConstantWriteNode) -> RbsInlineData::Parser::TypedDefinition?
+    def extract_definition: (Prism::ConstantWriteNode) -> RbsInlineData::Parser::TypedDefinition?
   end
 end