Extend override check to also check arity

Override now cares about parents method arity: if parents arity
doesn't match raises an error.

1 files changed, 78 insertions, 31 deletions

diff --git a/lib/gitlab/utils/override.rb b/lib/gitlab/utils/override.rb
index c412961ea3f..c87e97d0213 100644
--- a/lib/gitlab/utils/override.rb
+++ b/lib/gitlab/utils/override.rb
@@ -4,16 +4,11 @@ module Gitlab
   module Utils
     module Override
       class Extension
-        def self.verify_class!(klass, method_name)
-          instance_method_defined?(klass, method_name) ||
-            raise(
-              NotImplementedError.new(
-                "#{klass}\##{method_name} doesn't exist!"))
-        end
-
-        def self.instance_method_defined?(klass, name, include_super: true)
-          klass.instance_methods(include_super).include?(name) ||
-            klass.private_instance_methods(include_super).include?(name)
+        def self.verify_class!(klass, method_name, arity)
+          extension = new(klass)
+          parents = extension.parents_for(klass)
+          extension.verify_method!(
+            klass: klass, parents: parents, method_name: method_name, sub_method_arity: arity)
         end
 
         attr_reader :subject
@@ -22,35 +17,77 @@ module Gitlab
           @subject = subject
         end
 
-        def add_method_name(method_name)
-          method_names << method_name
-        end
-
-        def add_class(klass)
-          classes << klass
+        def parents_for(klass)
+          index = klass.ancestors.index(subject)
+          klass.ancestors.drop(index + 1)
         end
 
         def verify!
           classes.each do |klass|
-            index = klass.ancestors.index(subject)
-            parents = klass.ancestors.drop(index + 1)
-
-            method_names.each do |method_name|
-              parents.any? do |parent|
-                self.class.instance_method_defined?(
-                  parent, method_name, include_super: false)
-              end ||
-                raise(
-                  NotImplementedError.new(
-                    "#{klass}\##{method_name} doesn't exist!"))
+            parents = parents_for(klass)
+
+            method_names.each_pair do |method_name, arity|
+              verify_method!(
+                klass: klass,
+                parents: parents,
+                method_name: method_name,
+                sub_method_arity: arity)
             end
           end
         end
 
+        def verify_method!(klass:, parents:, method_name:, sub_method_arity:)
+          overridden_parent = parents.find do |parent|
+            instance_method_defined?(parent, method_name)
+          end
+
+          raise NotImplementedError.new("#{klass}\##{method_name} doesn't exist!") unless overridden_parent
+
+          super_method_arity = find_direct_method(overridden_parent, method_name).arity
+
+          unless arity_compatible?(sub_method_arity, super_method_arity)
+            raise NotImplementedError.new("#{subject}\##{method_name} has arity of #{sub_method_arity}, but #{overridden_parent}\##{method_name} has arity of #{super_method_arity}")
+          end
+        end
+
+        def add_method_name(method_name, arity = nil)
+          method_names[method_name] = arity
+        end
+
+        def add_class(klass)
+          classes << klass
+        end
+
+        def verify_override?(method_name)
+          method_names.has_key?(method_name)
+        end
+
         private
 
+        def instance_method_defined?(klass, name)
+          klass.instance_methods(false).include?(name) ||
+            klass.private_instance_methods(false).include?(name)
+        end
+
+        def find_direct_method(klass, name)
+          method = klass.instance_method(name)
+          method = method.super_method until method && klass == method.owner
+          method
+        end
+
+        def arity_compatible?(sub_method_arity, super_method_arity)
+          if sub_method_arity >= 0 && super_method_arity >= 0
+            # Regular arguments
+            sub_method_arity == super_method_arity
+          else
+            # It's too complex to check this case, just allow sub-method having negative arity
+            # But we don't allow sub_method_arity > 0 yet super_method_arity < 0
+            sub_method_arity < 0
+          end
+        end
+
         def method_names
-          @method_names ||= []
+          @method_names ||= {}
         end
 
         def classes
@@ -80,11 +117,21 @@ module Gitlab
       def override(method_name)
         return unless ENV['STATIC_VERIFICATION']
 
+        Override.extensions[self] ||= Extension.new(self)
+        Override.extensions[self].add_method_name(method_name)
+      end
+
+      def method_added(method_name)
+        super
+
+        return unless ENV['STATIC_VERIFICATION']
+        return unless Override.extensions[self]&.verify_override?(method_name)
+
+        method_arity = instance_method(method_name).arity
         if is_a?(Class)
-          Extension.verify_class!(self, method_name)
+          Extension.verify_class!(self, method_name, method_arity)
         else # We delay the check for modules
-          Override.extensions[self] ||= Extension.new(self)
-          Override.extensions[self].add_method_name(method_name)
+          Override.extensions[self].add_method_name(method_name, method_arity)
         end
       end