From b318ebf1015ced6354f8bbaf035308214b3f5c5d Mon Sep 17 00:00:00 2001
From: Jan Klusacek <honza.klu@gmail.com>
Date: Tue, 9 Jan 2018 22:47:35 +0100
Subject: [PATCH] py/modbuiltins: Add support for rounding integers.

As per CPython semantics.  This feature is controlled by
MICROPY_PY_BUILTINS_ROUND_INT which is disabled by default.
---
 py/modbuiltins.c                     | 31 +++++++++++++++++++++++++++-
 py/mpconfig.h                        |  5 +++++
 tests/basics/builtin_round_int.py    | 18 ++++++++++++++++
 tests/basics/builtin_round_intbig.py | 17 +++++++++++++++
 4 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 tests/basics/builtin_round_int.py
 create mode 100644 tests/basics/builtin_round_intbig.py

diff --git a/py/modbuiltins.c b/py/modbuiltins.c
index 0d511338b0..b216e021f3 100644
--- a/py/modbuiltins.c
+++ b/py/modbuiltins.c
@@ -445,7 +445,36 @@ MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_repr_obj, mp_builtin_repr);
 STATIC mp_obj_t mp_builtin_round(size_t n_args, const mp_obj_t *args) {
     mp_obj_t o_in = args[0];
     if (MP_OBJ_IS_INT(o_in)) {
-        return o_in;
+        if (n_args <= 1) {
+            return o_in;
+        }
+
+        #if !MICROPY_PY_BUILTINS_ROUND_INT
+        mp_raise_NotImplementedError(NULL);
+        #else
+        mp_int_t num_dig = mp_obj_get_int(args[1]);
+        if (num_dig >= 0) {
+            return o_in;
+        }
+
+        mp_obj_t mult = mp_binary_op(MP_BINARY_OP_POWER, MP_OBJ_NEW_SMALL_INT(10), MP_OBJ_NEW_SMALL_INT(-num_dig));
+        mp_obj_t half_mult =  mp_binary_op(MP_BINARY_OP_FLOOR_DIVIDE, mult, MP_OBJ_NEW_SMALL_INT(2));
+        mp_obj_t modulo = mp_binary_op(MP_BINARY_OP_MODULO, o_in, mult);
+        mp_obj_t rounded = mp_binary_op(MP_BINARY_OP_SUBTRACT, o_in, modulo);
+        if (mp_obj_is_true(mp_binary_op(MP_BINARY_OP_MORE, half_mult, modulo))) {
+            return rounded;
+        } else if (mp_obj_is_true(mp_binary_op(MP_BINARY_OP_MORE, modulo, half_mult))) {
+            return mp_binary_op(MP_BINARY_OP_ADD, rounded, mult);
+        } else {
+            // round to even number
+            mp_obj_t floor = mp_binary_op(MP_BINARY_OP_FLOOR_DIVIDE, o_in, mult);
+            if (mp_obj_is_true(mp_binary_op(MP_BINARY_OP_AND, floor, MP_OBJ_NEW_SMALL_INT(1)))) {
+                return mp_binary_op(MP_BINARY_OP_ADD, rounded, mult);
+            } else {
+                return rounded;
+            }
+        }
+        #endif
     }
 #if MICROPY_PY_BUILTINS_FLOAT
     mp_float_t val = mp_obj_get_float(o_in);
diff --git a/py/mpconfig.h b/py/mpconfig.h
index c42fe78536..100e2a9810 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -802,6 +802,11 @@ typedef double mp_float_t;
 #define MICROPY_PY_BUILTINS_RANGE_BINOP (0)
 #endif
 
+// Whether to support rounding of integers (incl bignum); eg round(123,-1)=120
+#ifndef MICROPY_PY_BUILTINS_ROUND_INT
+#define MICROPY_PY_BUILTINS_ROUND_INT (0)
+#endif
+
 // Whether to support timeout exceptions (like socket.timeout)
 #ifndef MICROPY_PY_BUILTINS_TIMEOUTERROR
 #define MICROPY_PY_BUILTINS_TIMEOUTERROR (0)
diff --git a/tests/basics/builtin_round_int.py b/tests/basics/builtin_round_int.py
new file mode 100644
index 0000000000..a2017622a7
--- /dev/null
+++ b/tests/basics/builtin_round_int.py
@@ -0,0 +1,18 @@
+# test round() with integer values and second arg
+
+# rounding integers is an optional feature so test for it
+try:
+    round(1, -1)
+except NotImplementedError:
+    print('SKIP')
+    raise SystemExit
+
+tests = [
+    (1, False), (1, True),
+    (124, -1), (125, -1), (126, -1),
+    (5, -1), (15, -1), (25, -1),
+    (12345, 0), (12345, -1), (12345, 1),
+    (-1234, 0), (-1234, -1), (-1234, 1),
+]
+for t in tests:
+    print(round(*t))
diff --git a/tests/basics/builtin_round_intbig.py b/tests/basics/builtin_round_intbig.py
new file mode 100644
index 0000000000..adf9d29f2f
--- /dev/null
+++ b/tests/basics/builtin_round_intbig.py
@@ -0,0 +1,17 @@
+# test round() with large integer values and second arg
+
+# rounding integers is an optional feature so test for it
+try:
+    round(1, -1)
+except NotImplementedError:
+    print('SKIP')
+    raise SystemExit
+
+i = 2**70
+
+tests = [
+    (i, 0), (i, -1), (i, -10), (i, 1),
+    (-i, 0), (-i, -1), (-i, -10), (-i, 1),
+]
+for t in tests:
+    print(round(*t))