[GR-28711] Fix method invalidation when a call site is used before include/prepend M and a method is added to M later

eregon · eregon · commit cdfb356ad572 · 2021-05-20T19:11:50.000Z
PullRequest: truffleruby/2659
diff --git a/spec/ruby/core/module/include_spec.rb b/spec/ruby/core/module/include_spec.rb
@@ -263,6 +263,30 @@ def foo
     foo.call.should == 'm'
   end
 
+
+  it "updates the method when a module included after a call is later updated" do
+    m_module = Module.new
+    a_class = Class.new do
+      def foo
+        'a'
+      end
+    end
+    b_class = Class.new(a_class)
+    b = b_class.new
+    foo = -> { b.foo }
+    foo.call.should == 'a'
+
+    b_class.include m_module
+    foo.call.should == 'a'
+
+    m_module.module_eval do
+      def foo
+        "m"
+      end
+    end
+    foo.call.should == 'm'
+  end
+
   it "updates the method when a nested included module is updated" do
     a_class = Class.new do
       def foo
diff --git a/spec/ruby/core/module/prepend_spec.rb b/spec/ruby/core/module/prepend_spec.rb
@@ -47,9 +47,8 @@ def foo
         'a'
       end
     end
-    b_class = Class.new(a_class)
-    b = b_class.new
-    foo = -> { b.foo }
+    a = a_class.new
+    foo = -> { a.foo }
     foo.call.should == 'a'
     a_class.class_eval do
       prepend m_module
@@ -65,9 +64,8 @@ def foo
         'a'
       end
     end
-    b_class = Class.new(a_class)
-    b = b_class.new
-    foo = -> { b.foo }
+    a = a_class.new
+    foo = -> { a.foo }
     foo.call.should == 'a'
     m_module.module_eval do
       def foo
@@ -77,6 +75,124 @@ def foo
     foo.call.should == 'm'
   end
 
+  it "updates the method when there is a base included method and the prepended module overrides it" do
+    base_module = Module.new do
+      def foo
+        'a'
+      end
+    end
+    a_class = Class.new do
+      include base_module
+    end
+    a = a_class.new
+    foo = -> { a.foo }
+    foo.call.should == 'a'
+
+    m_module = Module.new do
+      def foo
+        "m"
+      end
+    end
+    a_class.prepend m_module
+    foo.call.should == 'm'
+  end
+
+  it "updates the method when there is a base included method and the prepended module is later updated" do
+    base_module = Module.new do
+      def foo
+        'a'
+      end
+    end
+    a_class = Class.new do
+      include base_module
+    end
+    a = a_class.new
+    foo = -> { a.foo }
+    foo.call.should == 'a'
+
+    m_module = Module.new
+    a_class.prepend m_module
+    foo.call.should == 'a'
+
+    m_module.module_eval do
+      def foo
+        "m"
+      end
+    end
+    foo.call.should == 'm'
+  end
+
+  it "updates the method when a module prepended after a call is later updated" do
+    m_module = Module.new
+    a_class = Class.new do
+      def foo
+        'a'
+      end
+    end
+    a = a_class.new
+    foo = -> { a.foo }
+    foo.call.should == 'a'
+
+    a_class.prepend m_module
+    foo.call.should == 'a'
+
+    m_module.module_eval do
+      def foo
+        "m"
+      end
+    end
+    foo.call.should == 'm'
+  end
+
+  it "updates the method when a module is prepended after another and the method is defined later on that module" do
+    m_module = Module.new do
+      def foo
+        'a'
+      end
+    end
+    a_class = Class.new
+    a_class.prepend m_module
+    a = a_class.new
+    foo = -> { a.foo }
+    foo.call.should == 'a'
+
+    n_module = Module.new
+    a_class.prepend n_module
+    foo.call.should == 'a'
+
+    n_module.module_eval do
+      def foo
+        "n"
+      end
+    end
+    foo.call.should == 'n'
+  end
+
+  it "updates the method when a module is included in a prepended module and the method is defined later" do
+    a_class = Class.new
+    base_module = Module.new do
+      def foo
+        'a'
+      end
+    end
+    a_class.prepend base_module
+    a = a_class.new
+    foo = -> { a.foo }
+    foo.call.should == 'a'
+
+    m_module = Module.new
+    n_module = Module.new
+    m_module.include n_module
+    a_class.prepend m_module
+
+    n_module.module_eval do
+      def foo
+        "n"
+      end
+    end
+    foo.call.should == 'n'
+  end
+
   it "updates the method when a new module with an included module is prepended" do
     a_class = Class.new do
       def foo
@@ -94,9 +210,8 @@ def foo
       include n_module
     end
 
-    b_class = Class.new(a_class)
-    b = b_class.new
-    foo = -> { b.foo }
+    a = a_class.new
+    foo = -> { a.foo }
 
     foo.call.should == 'a'
 
diff --git a/src/main/java/org/truffleruby/core/module/ModuleFields.java b/src/main/java/org/truffleruby/core/module/ModuleFields.java
@@ -12,6 +12,7 @@
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Deque;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -21,6 +22,7 @@
 import java.util.Map.Entry;
 import java.util.NoSuchElementException;
 import java.util.Set;
+import java.util.WeakHashMap;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.locks.ReentrantLock;
@@ -113,6 +115,17 @@ public static void debugModuleChain(RubyModule module) {
     // Concurrency: only modified during boot
     private final Map<String, Assumption> inlinedBuiltinsAssumptions = new HashMap<>();
 
+    /** A weak set of all classes include-ing or prepend-ing this module M, and of modules which where already prepended
+     * before to such classes. Used by method invalidation so that lookups before the include/prepend is done can be
+     * invalidated when a method is added later in M. When M is an included module, defining a method on M only needs to
+     * invalidate classes in which M is included. However, when M is a prepended module defining a method on M needs to
+     * invalidate classes and other modules which were prepended before M (since the method could be defined there, and
+     * then the assumption on M would not be checked). Not that there is no need to be transitive here, include/prepend
+     * snapshot the modules to include alongside M, and later M.include/M.prepend have no effect on classes already
+     * including M. In other words,
+     * {@code module A; end; module B; end; class C; end; C.include A; A.include B; C.ancestors.include?(B) => false} */
+    private final Set<RubyModule> includedBy;
+
     public ModuleFields(
             RubyLanguage language,
             SourceSection sourceSection,
@@ -128,6 +141,9 @@ public ModuleFields(
         this.constantsUnmodifiedAssumption = new CyclicAssumption("constants are unmodified");
         classVariables = new ClassVariableStorage(language);
         start = new PrependMarker(this);
+        this.includedBy = rubyModule instanceof RubyClass
+                ? null
+                : Collections.newSetFromMap(Collections.synchronizedMap(new WeakHashMap<>()));
     }
 
     public RubyConstant getAdoptedByLexicalParent(
@@ -295,11 +311,14 @@ public void include(RubyContext context, Node currentNode, RubyModule module) {
 
     private void performIncludes(ModuleChain inclusionPoint, Deque<RubyModule> moduleAncestors) {
         while (!moduleAncestors.isEmpty()) {
-            RubyModule mod = moduleAncestors.pop();
-            inclusionPoint.insertAfter(mod);
-            // Module#include only adds modules between the current class and the super class,
-            // so invalidating the current class is enough as all affected lookups would go through the current class.
-            newMethodsVersion(mod.fields.getMethodNames());
+            RubyModule toInclude = moduleAncestors.pop();
+            inclusionPoint.insertAfter(toInclude);
+            if (rubyModule instanceof RubyClass) { // M.include(N) just registers N but does nothing until C.include/prepend(M)
+                toInclude.fields.includedBy.add(rubyModule);
+                // Module#include only adds modules between the current class and the super class,
+                // so invalidating the current class is enough as all affected lookups would go through the current class.
+                newMethodsVersion(toInclude.fields.getMethodNames());
+            }
         }
     }
 
@@ -327,31 +346,50 @@ public void prepend(RubyContext context, Node currentNode, RubyModule module) {
 
         SharedObjects.propagate(context.getLanguageSlow(), rubyModule, module);
 
-        // Previous calls on instances of the current class must have looked up through the first prepended module,
-        // so invalidate that one.
-        final ModuleFields moduleFieldsToInvalidate = getFirstModuleChain().getActualModule().fields;
+        /* We need to invalidate all prepended modules and the class, because call sites which looked up methods before
+         * only check the class or one of the prepend module (if the method is defined there). */
+        final List<RubyModule> prependedModulesAndClass = rubyModule instanceof RubyClass
+                ? getPrependedModulesAndClass()
+                : null;
 
         ModuleChain mod = module.fields.start;
         ModuleChain cur = start;
         while (mod != null &&
                 !(mod instanceof ModuleFields && ((ModuleFields) mod).rubyModule instanceof RubyClass)) {
             if (!(mod instanceof PrependMarker)) {
-                final RubyModule actualModule = mod.getActualModule();
-                if (!ModuleOperations.includesModule(rubyModule, actualModule)) {
-                    cur.insertAfter(actualModule);
-                    moduleFieldsToInvalidate.newMethodsVersion(actualModule.fields.getMethodNames());
+                final RubyModule toPrepend = mod.getActualModule();
+                if (!ModuleOperations.includesModule(rubyModule, toPrepend)) {
+                    cur.insertAfter(toPrepend);
+                    if (rubyModule instanceof RubyClass) { // M.prepend(N) just registers N but does nothing until C.prepend/include(M)
+                        final List<String> methodsToInvalidate = toPrepend.fields.getMethodNames();
+                        for (RubyModule moduleToInvalidate : prependedModulesAndClass) {
+                            toPrepend.fields.includedBy.add(moduleToInvalidate);
+                            moduleToInvalidate.fields.newMethodsVersion(methodsToInvalidate);
+                        }
+                    }
                     cur = cur.getParentModule();
                 }
             }
             mod = mod.getParentModule();
         }
 
         // If there were already prepended modules, invalidate the first of them
-        moduleFieldsToInvalidate.newHierarchyVersion();
+        newHierarchyVersion();
 
         invalidateBuiltinsAssumptions();
     }
 
+    private List<RubyModule> getPrependedModulesAndClass() {
+        final List<RubyModule> prependedModulesAndClass = new ArrayList<>();
+        ModuleChain chain = getFirstModuleChain();
+        while (chain != this) {
+            prependedModulesAndClass.add(chain.getActualModule());
+            chain = chain.getParentModule();
+        }
+        prependedModulesAndClass.add(rubyModule);
+        return prependedModulesAndClass;
+    }
+
     /** Set the value of a constant, possibly redefining it. */
     @TruffleBoundary
     public RubyConstant setConstant(RubyContext context, Node currentNode, String name, Object value) {
@@ -463,6 +501,11 @@ public void addMethod(RubyContext context, Node currentNode, InternalMethod meth
             if (previousMethodEntry != null) {
                 previousMethodEntry.invalidate(rubyModule, method.getName());
             }
+
+            if (includedBy != null && !includedBy.isEmpty()) {
+                invalidateIncludedBy(method.getName());
+            }
+
             // invalidate assumptions to not use an AST-inlined methods
             changedMethod(method.getName());
             if (refinedModule != null) {
@@ -733,17 +776,27 @@ public void newHierarchyVersion() {
         }
     }
 
-    public void newMethodsVersion(List<String> methodsToInvalidate) {
-        for (String entryToInvalidate : methodsToInvalidate) {
-            while (true) {
-                final MethodEntry methodEntry = methods.get(entryToInvalidate);
-                if (methodEntry == null) {
-                    break;
-                } else {
-                    methodEntry.invalidate(rubyModule, entryToInvalidate);
-                    if (methods.replace(entryToInvalidate, methodEntry, methodEntry.withNewAssumption())) {
-                        break;
-                    }
+    private void invalidateIncludedBy(String method) {
+        for (RubyModule module : includedBy) {
+            module.fields.newMethodVersion(method);
+        }
+    }
+
+    public void newMethodsVersion(Collection<String> methodsToInvalidate) {
+        for (String name : methodsToInvalidate) {
+            newMethodVersion(name);
+        }
+    }
+
+    private void newMethodVersion(String methodToInvalidate) {
+        while (true) {
+            final MethodEntry methodEntry = methods.get(methodToInvalidate);
+            if (methodEntry == null) {
+                return;
+            } else {
+                if (methods.replace(methodToInvalidate, methodEntry, methodEntry.withNewAssumption())) {
+                    methodEntry.invalidate(rubyModule, methodToInvalidate);
+                    return;
                 }
             }
         }
