refactor or-assignment nodes by extracting a RunTwiceBranchProfile class

Author: Nicolas Laurent

src/main/java/org/truffleruby/language/constants/OrAssignConstantNode.java

@@ -10,14 +10,14 @@
 package org.truffleruby.language.constants;
 
 import com.oracle.truffle.api.CompilerDirectives;
-import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.profiles.ConditionProfile;
 import org.truffleruby.core.cast.BooleanCastNode;
 import org.truffleruby.core.cast.BooleanCastNodeGen;
 import org.truffleruby.core.module.RubyModule;
 import org.truffleruby.language.RubyConstant;
 import org.truffleruby.language.RubyContextSourceNode;
+import org.truffleruby.utils.RunTwiceBranchProfile;
 
 /** e.g. for {@code module::FOO ||= 1}
  *
@@ -31,25 +31,9 @@ public class OrAssignConstantNode extends RubyContextSourceNode {
     @Child protected WriteConstantNode writeConstant;
     @Child private BooleanCastNode cast;
 
-    private enum ExecuteCounter {
-        NEVER,
-        ONCE,
-        MANY;
-
-        public ExecuteCounter next() {
-            if (this == NEVER) {
-                return ONCE;
-            } else {
-                return MANY;
-            }
-        }
-
-    }
-
-    @CompilationFinal private ExecuteCounter executeCounter = ExecuteCounter.NEVER;
     private final ConditionProfile triviallyUndefined = ConditionProfile.create();
     private final ConditionProfile defined = ConditionProfile.create();
-    private final ConditionProfile truthy = ConditionProfile.create();
+    private final RunTwiceBranchProfile writeTwiceProfile = new RunTwiceBranchProfile();
 
     public OrAssignConstantNode(ReadConstantNode readConstant, WriteConstantNode writeConstant) {
         this.readConstant = readConstant;
@@ -59,17 +43,11 @@ public OrAssignConstantNode(ReadConstantNode readConstant, WriteConstantNode wri
     @Override
     public Object execute(VirtualFrame frame) {
 
-        if (executeCounter != ExecuteCounter.MANY) {
-            // Make sure to reset the profiles after the first execution, as we expect the first execution to write
-            // the constant and subsequent execution to simply read it.
-            CompilerDirectives.transferToInterpreterAndInvalidate();
-            executeCounter = executeCounter.next();
-        }
-
         if (triviallyUndefined.profile(readConstant.isModuleTriviallyUndefined(frame, getLanguage(), getContext()))) {
             // It might not be defined because of autoloaded constants (maybe other reasons?),
             // simply attempt writing (which will trigger autoloads if required).
             // Since we didn't evaluate the module part yet, no side-effects can occur.
+            writeTwiceProfile.enter();
             return writeConstant.execute(frame);
         }
 
@@ -82,12 +60,16 @@ public Object execute(VirtualFrame frame) {
 
         // Next we check if the constant itself is defined, and if it is, we get its value.
         final RubyConstant constant = readConstant.getConstantIfDefined(module);
-        final Object value = defined.profile(constant != null)
+
+        final boolean isDefined = defined.profile(constant != null);
+
+        final Object value = isDefined
                 ? readConstant.getConstant(module, constant)
                 : null;
 
         // Write if the constant is undefined or if its value is falsy.
-        if (!defined.profile(constant != null) || !truthy.profile(castToBoolean(value))) {
+        if (!isDefined || !castToBoolean(value)) {
+            writeTwiceProfile.enter();
             return writeConstant.execute(frame, module);
         } else {
             return value;

src/main/java/org/truffleruby/language/control/OrLazyValueDefinedNode.java

@@ -15,41 +15,24 @@
 import org.truffleruby.language.RubyContextSourceNode;
 import org.truffleruby.language.RubyNode;
 
-import com.oracle.truffle.api.CompilerDirectives;
-import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.profiles.ConditionProfile;
+import org.truffleruby.utils.RunTwiceBranchProfile;
 
 /** OrLazyValueDefinedNode is used as the 'or' node for ||=, because we know from idiomatic Ruby usage that this is
  * often used to lazy initialize a value. In that case normal counting profiling gives a misleading result. With the RHS
  * having been executed once (the lazy initialization) it will be compiled expecting it to be used again. We know that
- * it's unlikely to be used again, so only compile it in when it's been used more than once, by using a small saturating
- * counter. */
+ * it's unlikely to be used again, so only compile it in when it's been used more than once, by using a
+ * {@link RunTwiceBranchProfile}. */
 public class OrLazyValueDefinedNode extends RubyContextSourceNode {
 
     @Child private RubyNode left;
     @Child private RubyNode right;
 
     @Child private BooleanCastNode leftCast = BooleanCastNode.create();
 
-    private enum RightUsage {
-        NEVER,
-        ONCE,
-        MANY;
-
-        public RightUsage next() {
-            if (this == NEVER) {
-                return ONCE;
-            } else {
-                return MANY;
-            }
-        }
-
-    }
-
-    @CompilationFinal private RightUsage rightUsage = RightUsage.NEVER;
-
-    private final ConditionProfile conditionProfile = ConditionProfile.createCountingProfile();
+    private final RunTwiceBranchProfile rightTwiceProfile = new RunTwiceBranchProfile();
+    private final ConditionProfile countingProfile = ConditionProfile.createCountingProfile();
 
     public OrLazyValueDefinedNode(RubyNode left, RubyNode right) {
         this.left = left;
@@ -60,15 +43,10 @@ public OrLazyValueDefinedNode(RubyNode left, RubyNode right) {
     public Object execute(VirtualFrame frame) {
         final Object leftValue = left.execute(frame);
 
-        if (conditionProfile.profile(leftCast.executeToBoolean(leftValue))) {
+        if (countingProfile.profile(leftCast.executeToBoolean(leftValue))) {
             return leftValue;
         } else {
-            if (rightUsage != RightUsage.MANY) {
-                // Don't compile the RHS of a lazy initialization unless it has been used many times
-                CompilerDirectives.transferToInterpreterAndInvalidate();
-                rightUsage = rightUsage.next();
-            }
-
+            rightTwiceProfile.enter();
             return right.execute(frame);
         }
     }

src/main/java/org/truffleruby/utils/RunTwiceBranchProfile.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2021 Oracle and/or its affiliates. All rights reserved. This
+ * code is released under a tri EPL/GPL/LGPL license. You can use it,
+ * redistribute it and/or modify it under the terms of the:
+ *
+ * Eclipse Public License version 2.0, or
+ * GNU General Public License version 2, or
+ * GNU Lesser General Public License version 2.1.
+ */
+package org.truffleruby.utils;
+
+import com.oracle.truffle.api.CompilerDirectives;
+import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
+
+public final class RunTwiceBranchProfile {
+
+    private enum ExecuteCounter {
+        NEVER,
+        ONCE,
+        MANY;
+
+        public ExecuteCounter next() {
+            if (this == NEVER) {
+                return ONCE;
+            } else {
+                return MANY;
+            }
+        }
+    }
+
+    @CompilationFinal private ExecuteCounter executeCounter = ExecuteCounter.NEVER;
+
+    public void enter() {
+        if (executeCounter != ExecuteCounter.MANY) {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            executeCounter = executeCounter.next();
+        }
+    }
+}