[Truffle] Compilation of Array#inject.

jruby · Mar 5, 2015 · f57425f · f57425f
1 parent 6dc1b69
commit f57425f
Showing 1 changed file with 67 additions and 11 deletions.
diff --git a/truffle/src/main/java/org/jruby/truffle/nodes/core/ArrayNodes.java b/truffle/src/main/java/org/jruby/truffle/nodes/core/ArrayNodes.java
@@ -1309,11 +1309,11 @@ public InjectNode(InjectNode prev) {
             dispatch = prev.dispatch;
         }
 
-        @Specialization(guards = "isObject")
-        public Object injectObject(VirtualFrame frame, RubyArray array, Object initial, RubyProc block) {
+        @Specialization(guards = "isIntegerFixnum")
+        public Object injectIntegerFixnum(VirtualFrame frame, RubyArray array, Object initial, RubyProc block) {
             int count = 0;
 
-            final Object[] store = (Object[]) array.getStore();
+            final int[] store = (int[]) array.getStore();
 
             Object accumulator = initial;
 
@@ -1334,20 +1334,76 @@ public Object injectObject(VirtualFrame frame, RubyArray array, Object initial,
             return accumulator;
         }
 
-        @Specialization
-        public Object inject(VirtualFrame frame, RubyArray array, Object initial, RubyProc block) {
-            notDesignedForCompilation();
+        @Specialization(guards = "isLongFixnum")
+        public Object injectLongFixnum(VirtualFrame frame, RubyArray array, Object initial, RubyProc block) {
+            int count = 0;
 
-            final Object[] store = array.slowToArray();
+            final long[] store = (long[]) array.getStore();
 
-            if (store.length < 2) {
-                throw new UnsupportedOperationException();
+            Object accumulator = initial;
+
+            try {
+                for (int n = 0; n < array.getSize(); n++) {
+                    if (CompilerDirectives.inInterpreter()) {
+                        count++;
+                    }
+
+                    accumulator = yield(frame, block, accumulator, store[n]);
+                }
+            } finally {
+                if (CompilerDirectives.inInterpreter()) {
+                    getRootNode().reportLoopCount(count);
+                }
             }
 
+            return accumulator;
+        }
+
+        @Specialization(guards = "isFloat")
+        public Object injectFloat(VirtualFrame frame, RubyArray array, Object initial, RubyProc block) {
+            int count = 0;
+
+            final double[] store = (double[]) array.getStore();
+
             Object accumulator = initial;
 
-            for (int n = 0; n < array.getSize(); n++) {
-                accumulator = yield(frame, block, accumulator, store[n]);
+            try {
+                for (int n = 0; n < array.getSize(); n++) {
+                    if (CompilerDirectives.inInterpreter()) {
+                        count++;
+                    }
+
+                    accumulator = yield(frame, block, accumulator, store[n]);
+                }
+            } finally {
+                if (CompilerDirectives.inInterpreter()) {
+                    getRootNode().reportLoopCount(count);
+                }
+            }
+
+            return accumulator;
+        }
+
+        @Specialization(guards = "isObject")
+        public Object injectObject(VirtualFrame frame, RubyArray array, Object initial, RubyProc block) {
+            int count = 0;
+
+            final Object[] store = (Object[]) array.getStore();
+
+            Object accumulator = initial;
+
+            try {
+                for (int n = 0; n < array.getSize(); n++) {
+                    if (CompilerDirectives.inInterpreter()) {
+                        count++;
+                    }
+
+                    accumulator = yield(frame, block, accumulator, store[n]);
+                }
+            } finally {
+                if (CompilerDirectives.inInterpreter()) {
+                    getRootNode().reportLoopCount(count);
+                }
             }
 
             return accumulator;