Implementation of math.dist()

Author: otethal

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/MathModuleBuiltins.java

@@ -43,14 +43,17 @@
 import com.oracle.graal.python.builtins.PythonBuiltinClassType;
 import com.oracle.graal.python.builtins.PythonBuiltins;
 import com.oracle.graal.python.builtins.objects.PNone;
+import com.oracle.graal.python.builtins.objects.common.SequenceNodes;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
 import com.oracle.graal.python.builtins.objects.floats.PFloat;
+import com.oracle.graal.python.builtins.objects.function.PArguments;
 import com.oracle.graal.python.builtins.objects.function.PKeyword;
 import com.oracle.graal.python.builtins.objects.ints.PInt;
 import com.oracle.graal.python.builtins.objects.object.PythonObjectLibrary;
 import com.oracle.graal.python.builtins.objects.tuple.PTuple;
 import com.oracle.graal.python.nodes.ErrorMessages;
 import com.oracle.graal.python.nodes.PGuards;
+import com.oracle.graal.python.nodes.builtins.TupleNodes;
 import com.oracle.graal.python.nodes.call.special.LookupAndCallBinaryNode;
 import com.oracle.graal.python.nodes.call.special.LookupAndCallUnaryNode;
 import com.oracle.graal.python.nodes.control.GetIteratorExpressionNode.GetIteratorNode;
@@ -2686,4 +2689,58 @@ public Object doGeneric(VirtualFrame frame, Object iterable, Object start) {
             }
         }
     }
+
+    @Builtin(name = "dist", minNumOfPositionalArgs = 2, numOfPositionalOnlyArgs = 2, parameterNames = {"p", "q"})
+    @GenerateNodeFactory
+    public abstract static class DistNode extends PythonBuiltinNode {
+
+        @Child private TupleNodes.ConstructTupleNode tupleCtor = TupleNodes.ConstructTupleNode.create();
+        @Child private SequenceNodes.GetObjectArrayNode getObjectArray = SequenceNodes.GetObjectArrayNode.create();
+
+        @Specialization
+        public double doGeneric(VirtualFrame frame, Object p, Object q,
+                        @CachedLibrary(limit = "4") PythonObjectLibrary lib) {
+            // adapted from CPython math_dist_impl and vector_norm
+            Object[] ps = getObjectArray.execute(tupleCtor.execute(frame, p));
+            Object[] qs = getObjectArray.execute(tupleCtor.execute(frame, q));
+            int len = ps.length;
+            if (len != qs.length) {
+                throw raise(ValueError, ErrorMessages.BOTH_POINTS_MUST_HAVE_THE_SAME_NUMBER_OF_DIMENSIONS);
+            }
+            double[] diffs = new double[len];
+            double max = 0.0;
+            boolean foundNan = false;
+            for (int i = 0; i < len; ++i) {
+                double a = lib.asJavaDoubleWithState(ps[i], PArguments.getThreadState(frame));
+                double b = lib.asJavaDoubleWithState(qs[i], PArguments.getThreadState(frame));
+                double x = Math.abs(a - b);
+                diffs[i] = x;
+                foundNan |= Double.isNaN(x);
+                if (x > max) {
+                    max = x;
+                }
+            }
+            if (Double.isInfinite(max)) {
+                return max;
+            }
+            if (foundNan) {
+                return Double.NaN;
+            }
+            if (max == 0.0 || len <= 1) {
+                return max;
+            }
+
+            double csum = 1.0;
+            double frac = 0.0;
+            for (int i = 0; i < len; ++i) {
+                double x = diffs[i];
+                x /= max;
+                x = x * x;
+                double oldcsum = csum;
+                csum += x;
+                frac += (oldcsum - csum) + x;
+            }
+            return max * Math.sqrt(csum - 1.0 + frac);
+        }
+    }
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/SequenceNodes.java

@@ -125,5 +125,9 @@ static Object[] doGeneric(Object seq,
                         @Cached SequenceStorageNodes.ToArrayNode toArrayNode) {
             return toArrayNode.execute(getSequenceStorageNode.execute(seq));
         }
+
+        public static GetObjectArrayNode create() {
+            return SequenceNodesFactory.GetObjectArrayNodeGen.create();
+        }
     }
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/ints/PInt.java

@@ -246,14 +246,26 @@ public boolean canBeJavaDouble() {
     public double asJavaDouble(
                     @CachedLibrary("this") PythonObjectLibrary lib,
                     @Exclusive @Cached CastToJavaDoubleNode castToDouble,
-                    @Exclusive @Cached() ConditionProfile hasIndexFunc,
+                    @Exclusive @Cached ConditionProfile hasIndexFunc,
                     @Exclusive @Cached PRaiseNode raise) {
         if (hasIndexFunc.profile(lib.canBeIndex(this))) {
             return castToDouble.execute(lib.asIndex(this));
         }
         throw raise.raise(TypeError, ErrorMessages.MUST_BE_REAL_NUMBER, this);
     }
 
+    @ExportMessage
+    public double asJavaDoubleWithState(ThreadState threadState,
+                    @CachedLibrary("this") PythonObjectLibrary lib,
+                    @Exclusive @Cached CastToJavaDoubleNode castToDouble,
+                    @Exclusive @Cached ConditionProfile hasIndexFunc,
+                    @Exclusive @Cached PRaiseNode raise) {
+        if (hasIndexFunc.profile(lib.canBeIndex(this))) {
+            return castToDouble.execute(lib.asIndexWithState(this, threadState));
+        }
+        throw raise.raise(TypeError, ErrorMessages.MUST_BE_REAL_NUMBER, this);
+    }
+
     @SuppressWarnings("static-method")
     @ExportMessage
     public boolean canBeJavaLong() {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/ErrorMessages.java

@@ -86,6 +86,7 @@ public abstract class ErrorMessages {
     public static final String BASES_MUST_BE_TYPES = "bases must be types";
     public static final String BASES_ITEM_CAUSES_INHERITANCE_CYCLE = "a __bases__ item causes an inheritance cycle";
     public static final String BOOL_SHOULD_RETURN_BOOL = "__bool__ should return bool, returned %p";
+    public static final String BOTH_POINTS_MUST_HAVE_THE_SAME_NUMBER_OF_DIMENSIONS = "both points must have the same number of dimensions";
     public static final String BUFFER_INDICES_MUST_BE_INTS = "buffer indices must be integers, not %p";
     public static final String BYTE_STR_IS_TOO_LARGE = "byte string is too large";
     public static final String BYTEARRAY_OUT_OF_BOUNDS = "bytearray index out of range";