py/emitnative: Implement floor-division and modulo for viper emitter.

dpgeorge · dpgeorge · commit a3afa8cfc469 · 2017-10-11T18:54:34.000+11:00
diff --git a/py/emitnative.c b/py/emitnative.c
@@ -123,6 +123,8 @@ STATIC byte mp_f_n_args[MP_F_NUMBER_OF] = {
     [MP_F_NEW_CELL] = 1,
     [MP_F_MAKE_CLOSURE_FROM_RAW_CODE] = 3,
     [MP_F_SETUP_CODE_STATE] = 5,
+    [MP_F_SMALL_INT_FLOOR_DIVIDE] = 2,
+    [MP_F_SMALL_INT_MODULO] = 2,
 };
 
 #include "py/asmx86.h"
@@ -1843,6 +1845,20 @@ STATIC void emit_native_binary_op(emit_t *emit, mp_binary_op_t op) {
             return;
         }
         #endif
+
+        // special cases for floor-divide and module because we dispatch to helper functions
+        if (op == MP_BINARY_OP_FLOOR_DIVIDE || op == MP_BINARY_OP_INPLACE_FLOOR_DIVIDE
+            || op == MP_BINARY_OP_MODULO || op == MP_BINARY_OP_INPLACE_MODULO) {
+            emit_pre_pop_reg_reg(emit, &vtype_rhs, REG_ARG_2, &vtype_lhs, REG_ARG_1);
+            if (op == MP_BINARY_OP_FLOOR_DIVIDE || op == MP_BINARY_OP_INPLACE_FLOOR_DIVIDE) {
+                emit_call(emit, MP_F_SMALL_INT_FLOOR_DIVIDE);
+            } else {
+                emit_call(emit, MP_F_SMALL_INT_MODULO);
+            }
+            emit_post_push_reg(emit, VTYPE_INT, REG_RET);
+            return;
+        }
+
         int reg_rhs = REG_ARG_3;
         emit_pre_pop_reg_flexible(emit, &vtype_rhs, &reg_rhs, REG_RET, REG_ARG_2);
         emit_pre_pop_reg(emit, &vtype_lhs, REG_ARG_2);
diff --git a/py/nativeglue.c b/py/nativeglue.c
@@ -29,6 +29,7 @@
 #include <assert.h>
 
 #include "py/runtime.h"
+#include "py/smallint.h"
 #include "py/emitglue.h"
 #include "py/bc.h"
 
@@ -170,6 +171,8 @@ void *const mp_fun_table[MP_F_NUMBER_OF] = {
     mp_obj_new_cell,
     mp_make_closure_from_raw_code,
     mp_setup_code_state,
+    mp_small_int_floor_divide,
+    mp_small_int_modulo,
 };
 
 /*
diff --git a/py/runtime0.h b/py/runtime0.h
@@ -185,6 +185,8 @@ typedef enum {
     MP_F_NEW_CELL,
     MP_F_MAKE_CLOSURE_FROM_RAW_CODE,
     MP_F_SETUP_CODE_STATE,
+    MP_F_SMALL_INT_FLOOR_DIVIDE,
+    MP_F_SMALL_INT_MODULO,
     MP_F_NUMBER_OF,
 } mp_fun_kind_t;
 
diff --git a/tests/micropython/viper_binop_divmod.py b/tests/micropython/viper_binop_divmod.py
@@ -0,0 +1,18 @@
+# test floor-division and modulo operators
+
+@micropython.viper
+def div(x:int, y:int) -> int:
+    return x // y
+
+@micropython.viper
+def mod(x:int, y:int) -> int:
+    return x % y
+
+def dm(x, y):
+    print(div(x, y), mod(x, y))
+
+for x in (-6, 6):
+    for y in range(-7, 8):
+        if y == 0:
+            continue
+        dm(x, y)
diff --git a/tests/micropython/viper_binop_divmod.py.exp b/tests/micropython/viper_binop_divmod.py.exp
@@ -0,0 +1,28 @@
+0 -6
+1 0
+1 -1
+1 -2
+2 0
+3 0
+6 0
+-6 0
+-3 0
+-2 0
+-2 2
+-2 4
+-1 0
+-1 1
+-1 -1
+-1 0
+-2 -4
+-2 -2
+-2 0
+-3 0
+-6 0
+6 0
+3 0
+2 0
+1 2
+1 1
+1 0
+0 6

-Original file line number
+Diff line change
@@ @@ -0,0 +1,28 @@ @@
 +0 -6
 +1 0
 +1 -1
 +1 -2
 +2 0
 +3 0
 +6 0
 +-6 0
 +-3 0
 +-2 0
 +-2 2
 +-2 4
 +-1 0
 +-1 1
 +-1 -1
 +-1 0
 +-2 -4
 +-2 -2
 +-2 0
 +-3 0
 +-6 0
 +6 0
 +3 0
 +2 0
 +1 2
 +1 1
 +1 0
 +0 6