User can use include: [] instead of modules

lib/diver_down/helper.rb

@@ -45,12 +45,18 @@ def self.resolve_singleton_class(obj)
       end
     end
 
-    # @param obj [Object, Module, Class]
-    # @return ['instance', 'class']
+    # @param obj [Object]
+    # @return [Boolean]
     def self.module?(obj)
       Module === obj
     end
 
+    # @param obj [Object]
+    # @return [Boolean]
+    def self.class?(obj)
+      Class === obj
+    end
+
     # @param str [String]
     # @return [Module]
     def self.constantize(str)

lib/diver_down/trace/module_set.rb

@@ -4,56 +4,72 @@ module DiverDown
   module Trace
     # A class to quickly determine if a TracePoint is a module to be traced.
     class ModuleSet
-      require 'diver_down/trace/module_set/base_strategy'
-      require 'diver_down/trace/module_set/array_strategy'
-      require 'diver_down/trace/module_set/const_source_location_strategy'
+      require 'diver_down/trace/module_set/array_module_set'
+      require 'diver_down/trace/module_set/const_source_location_module_set'
 
-      # @param [DiverDown::Trace::ModuleSet::BaseStrategy] strategy
-      def initialize(strategy)
-        @strategy = strategy
+      # @param [Array<Module, String>, Set<Module, String>, nil] modules
+      # @param [Array<String>] include
+      def initialize(modules: nil, include: nil)
+        @cache = {}
+        @array_module_set = DiverDown::Trace::ModuleSet::ArrayModuleSet.new(modules) unless modules.nil?
+        @const_source_location_module_set = DiverDown::Trace::ModuleSet::ConstSourceLocationModuleSet.new(include:) unless include.nil?
       end
 
       # @param [Module] mod
       # @return [Boolean]
       def include?(mod)
-        result = @strategy[mod]
+        unless @cache.key?(mod)
+          if DiverDown::Helper.class?(mod)
+            # class
+            begin
+              dig_superclass(mod)
+            rescue TypeError => e
+              # https://github.com/ruby/ruby/blob/f42164e03700469a7000b4f00148a8ca01d75044/object.c#L2232
+              return false if e.message == 'uninitialized class'
 
-        if result.nil?
-          # If the strategy doesn't know, check the superclass
-          superclass_include(mod)
-        else
-          result
+              raise
+            end
+          else
+            # module
+            @cache[mod] = !!_include?(mod)
+          end
         end
+
+        @cache.fetch(mod)
       end
 
       private
 
-      def superclass_include(mod)
-        return false unless Class === mod
+      def _include?(mod)
+        @array_module_set&.include?(mod) ||
+          @const_source_location_module_set&.include?(mod)
+      end
+
+      def dig_superclass(mod)
+        stack = []
+        current = mod
+        included = nil
 
-        stack = [mod]
-        current = mod.superclass
-        found = nil
+        until current.nil?
+          if @cache.key?(current)
+            included = @cache.fetch(current)
+            break
+          else
+            stack.push(current)
+            included = _include?(current)
 
-        while current && found.nil?
-          found = @strategy[current] # nil or boolean
-          stack.push(current)
-          current = current.superclass
+            break if included
+
+            current = current.superclass
+          end
         end
 
         # Convert nil to boolean
-        found = !!found
+        included = !!included
 
         stack.each do
-          @strategy[_1] = found
+          @cache[_1] = included unless @cache.key?(_1)
         end
-
-        found
-      rescue TypeError => e
-        # https://github.com/ruby/ruby//blob/f42164e03700469a7000b4f00148a8ca01d75044/object.c#L2232
-        return false if e.message == 'uninitialized class'
-
-        raise
       end
     end
   end

lib/diver_down/trace/module_set/array_strategy.rb → lib/diver_down/trace/module_set/array_module_set.rb

@@ -3,12 +3,10 @@
 module DiverDown
   module Trace
     class ModuleSet
-      class ArrayStrategy < BaseStrategy
+      class ArrayModuleSet
         # @param [Array<Module, String>, #each] modules
         def initialize(modules)
-          super()
-
-          @set = {}
+          @map = {}
 
           modules.each do
             mod = if DiverDown::Helper.module?(_1)
@@ -18,21 +16,14 @@ def initialize(modules)
                     DiverDown::Helper.constantize(_1)
                   end
 
-            self[mod] = true
+            @map[mod] = true
           end
         end
 
         # @param [Module] mod
         # @return [Boolean, nil]
-        def [](mod)
-          @set[mod]
-        end
-
-        # @param [Module] mod
-        # @param value [Boolean]
-        # @return [void]
-        def []=(mod, value)
-          @set[mod] = value
+        def include?(mod)
+          @map[mod]
         end
       end
     end

lib/diver_down/trace/module_set/const_source_location_module_set.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module DiverDown
+  module Trace
+    class ModuleSet
+      class ConstSourceLocationModuleSet
+        # @param [Array<String>, Set<String>] include
+        def initialize(include: [])
+          @include = include.to_set
+        end
+
+        # @param [Module] mod
+        # @return [Boolean]
+        def include?(mod)
+          module_name = DiverDown::Helper.normalize_module_name(mod)
+
+          path, = begin
+            Object.const_source_location(module_name)
+          rescue NameError, TypeError
+            nil
+          end
+
+          path && @include.include?(path)
+        end
+      end
+    end
+  end
+end

lib/diver_down/trace/tracer.rb

@@ -27,12 +27,10 @@ def initialize(module_set: [], target_files: nil, filter_method_id_path: nil, mo
 
         @module_set = if module_set.is_a?(DiverDown::Trace::ModuleSet)
                         module_set
-                      elsif module_set.is_a?(DiverDown::Trace::ModuleSet::BaseStrategy)
-                        DiverDown::Trace::ModuleSet.new(module_set)
+                      elsif module_set.respond_to?(:to_set)
+                        DiverDown::Trace::ModuleSet.new(modules: module_set)
                       else
-                        DiverDown::Trace::ModuleSet.new(
-                          DiverDown::Trace::ModuleSet::ArrayStrategy.new(module_set)
-                        )
+                        DiverDown::Trace::ModuleSet.new(**module_set)
                       end
 
         @target_file_set = target_files&.to_set

spec/diver_down/helper_spec.rb

@@ -145,6 +145,25 @@ def self.name
       end
     end
 
+    describe '.class?' do
+      it 'returns true given class' do
+        expect(described_class.class?(Class.new)).to be(true)
+      end
+
+      it 'returns true given singleton class' do
+        expect(described_class.class?(Class.new.singleton_class)).to be(true)
+      end
+
+      it 'returns false given module' do
+        expect(described_class.class?(Module.new)).to be(false)
+      end
+
+      it 'returns false given instance' do
+        klass = Class.new
+        expect(described_class.class?(klass.new)).to be(false)
+      end
+    end
+
     describe '.constantize' do
       it 'returns constant given string' do
         expect(described_class.constantize('String')).to eq(String)

spec/diver_down/trace/module_set_spec.rb

@@ -3,13 +3,13 @@
 RSpec.describe DiverDown::Trace::ModuleSet do
   describe 'InstanceMethods' do
     describe '#include?' do
-      context 'array strategy' do
+      context 'with modules' do
         it 'checks module or module name' do
           stub_const('A', Module.new)
           stub_const('B', Module.new)
 
           set = described_class.new(
-            DiverDown::Trace::ModuleSet::ArrayStrategy.new([A])
+            modules: [A]
           )
 
           expect(set.include?(A)).to be(true)
@@ -22,7 +22,7 @@
           stub_const('C', Class.new(B))
 
           set = described_class.new(
-            DiverDown::Trace::ModuleSet::ArrayStrategy.new([B])
+            modules: [B]
           )
 
           expect(set.include?(A)).to be(false)
@@ -48,9 +48,7 @@ class ::A
             # rubocop:enable Lint/ConstantDefinitionInBlock
 
             set = described_class.new(
-              DiverDown::Trace::ModuleSet::ConstSourceLocationStrategy.new(
-                include: [__FILE__]
-              )
+              include: [__FILE__]
             )
 
             expect(set.include?(A)).to be(true)
@@ -66,9 +64,7 @@ class ::B < ::A; end
             # rubocop:enable Lint/ConstantDefinitionInBlock
 
             set = described_class.new(
-              DiverDown::Trace::ModuleSet::ConstSourceLocationStrategy.new(
-                include: [__FILE__]
-              )
+              include: [__FILE__]
             )
 
             expect(set.include?(A)).to be(true)