Fix comparisons of arrays containing NaN

Author: msimacek

graalpython/com.oracle.graal.python.test/src/tests/unittest_tags/test_array.txt

@@ -92,6 +92,7 @@
 *graalpython.lib-python.3.test.test_array.DoubleTest.test_iterator_pickle
 *graalpython.lib-python.3.test.test_array.DoubleTest.test_len
 *graalpython.lib-python.3.test.test_array.DoubleTest.test_mul
+*graalpython.lib-python.3.test.test_array.DoubleTest.test_nan
 *graalpython.lib-python.3.test.test_array.DoubleTest.test_obsolete_write_lock
 *graalpython.lib-python.3.test.test_array.DoubleTest.test_pickle
 *graalpython.lib-python.3.test.test_array.DoubleTest.test_pickle_for_empty_array
@@ -147,6 +148,7 @@
 *graalpython.lib-python.3.test.test_array.FloatTest.test_iterator_pickle
 *graalpython.lib-python.3.test.test_array.FloatTest.test_len
 *graalpython.lib-python.3.test.test_array.FloatTest.test_mul
+*graalpython.lib-python.3.test.test_array.FloatTest.test_nan
 *graalpython.lib-python.3.test.test_array.FloatTest.test_obsolete_write_lock
 *graalpython.lib-python.3.test.test_array.FloatTest.test_pickle
 *graalpython.lib-python.3.test.test_array.FloatTest.test_pickle_for_empty_array

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/array/ArrayBuiltins.java

@@ -98,6 +98,7 @@
 import com.oracle.truffle.api.dsl.Cached;
 import com.oracle.truffle.api.dsl.Fallback;
 import com.oracle.truffle.api.dsl.GenerateNodeFactory;
+import com.oracle.truffle.api.dsl.ImportStatic;
 import com.oracle.truffle.api.dsl.NodeFactory;
 import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.frame.VirtualFrame;
@@ -221,9 +222,10 @@ protected ArgumentClinicProvider getArgumentClinic() {
 
     @Builtin(name = __EQ__, minNumOfPositionalArgs = 2)
     @GenerateNodeFactory
+    @ImportStatic(BufferFormat.class)
     abstract static class EqNode extends PythonBinaryBuiltinNode {
 
-        @Specialization(guards = "shouldCompareBytes(left, right)")
+        @Specialization(guards = {"left.getFormat() == right.getFormat()", "!isFloatingPoint(left.getFormat())"})
         static boolean eqBytes(PArray left, PArray right) {
             if (left.getLength() != right.getLength()) {
                 return false;
@@ -237,7 +239,7 @@ static boolean eqBytes(PArray left, PArray right) {
             return true;
         }
 
-        @Specialization(guards = "!shouldCompareBytes(left, right)")
+        @Specialization(guards = "left.getFormat() != right.getFormat()")
         static boolean eqItems(PArray left, PArray right,
                         @CachedLibrary(limit = "4") PythonObjectLibrary lib,
                         @Cached ArrayNodes.GetValueNode getLeft,
@@ -253,21 +255,44 @@ static boolean eqItems(PArray left, PArray right,
             return true;
         }
 
+        // Separate specialization for float/double is needed because of NaN comparisons
+        @Specialization(guards = {"left.getFormat() == right.getFormat()", "isFloatingPoint(left.getFormat())"})
+        static boolean eqDoubles(PArray left, PArray right,
+                        @Cached ArrayNodes.GetValueNode getLeft,
+                        @Cached ArrayNodes.GetValueNode getRight) {
+            if (left.getLength() != right.getLength()) {
+                return false;
+            }
+            for (int i = 0; i < left.getLength(); i++) {
+                double leftValue = (Double) getLeft.execute(left, i);
+                double rightValue = (Double) getRight.execute(right, i);
+                if (leftValue != rightValue) {
+                    return false;
+                }
+            }
+            return true;
+        }
+
         @Specialization(guards = "!isArray(right)")
         @SuppressWarnings("unused")
         static Object eq(PArray left, Object right) {
             return PNotImplemented.NOT_IMPLEMENTED;
         }
 
+        protected static boolean shouldCompareDouble(PArray left, PArray right) {
+            return left.getFormat() == right.getFormat() && (left.getFormat() == BufferFormat.DOUBLE || left.getFormat() == BufferFormat.FLOAT);
+        }
+
         protected static boolean shouldCompareBytes(PArray left, PArray right) {
             return left.getFormat() == right.getFormat() && left.getFormat() != BufferFormat.DOUBLE && left.getFormat() != BufferFormat.FLOAT;
         }
     }
 
+    @ImportStatic(BufferFormat.class)
     abstract static class AbstractComparisonNode extends PythonBinaryBuiltinNode {
 
-        @Specialization
-        boolean cmp(VirtualFrame frame, PArray left, PArray right,
+        @Specialization(guards = "!isFloatingPoint(left.getFormat()) || (left.getFormat() != right.getFormat())")
+        boolean cmpItems(VirtualFrame frame, PArray left, PArray right,
                         @CachedLibrary(limit = "4") PythonObjectLibrary lib,
                         @Cached("createComparison()") BinaryComparisonNode compareNode,
                         @Cached("createIfTrueNode()") CoerceToBooleanNode coerceToBooleanNode,
@@ -284,6 +309,24 @@ boolean cmp(VirtualFrame frame, PArray left, PArray right,
             return compareLengths(left.getLength(), right.getLength());
         }
 
+        // Separate specialization for float/double is needed because of NaN comparisons
+        @Specialization(guards = {"isFloatingPoint(left.getFormat())", "left.getFormat() == right.getFormat()"})
+        boolean cmpDoubles(VirtualFrame frame, PArray left, PArray right,
+                        @Cached("createComparison()") BinaryComparisonNode compareNode,
+                        @Cached("createIfTrueNode()") CoerceToBooleanNode coerceToBooleanNode,
+                        @Cached ArrayNodes.GetValueNode getLeft,
+                        @Cached ArrayNodes.GetValueNode getRight) {
+            int commonLength = Math.min(left.getLength(), right.getLength());
+            for (int i = 0; i < commonLength; i++) {
+                double leftValue = (Double) getLeft.execute(left, i);
+                double rightValue = (Double) getRight.execute(right, i);
+                if (leftValue != rightValue) {
+                    return coerceToBooleanNode.executeBoolean(frame, compareNode.executeWith(frame, leftValue, rightValue));
+                }
+            }
+            return compareLengths(left.getLength(), right.getLength());
+        }
+
         @Specialization(guards = "!isArray(right)")
         @SuppressWarnings("unused")
         static Object cmp(PArray left, Object right) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/util/BufferFormat.java

@@ -135,4 +135,8 @@ private static BufferFormat fromCharCommon(char fmtchar) {
         return null;
     }
 
+    public static boolean isFloatingPoint(BufferFormat format) {
+        return format == FLOAT || format == DOUBLE;
+    }
+
 }