py/compile: Add COMP_RETURN_IF_EXPR option to enable return-if-else opt.

dpgeorge · dpgeorge · commit ae54fbf16689 · 2017-04-22T14:58:01.000+10:00
With this optimisation enabled the compiler optimises the if-else
expression within a return statement.  The optimisation reduces bytecode
size by 2 bytes for each use of such a return-if-else statement.  Since
such a statement is not often used, and costs bytes for the code, the
feature is disabled by default.

For example the following code:

    def f(x):
        return 1 if x else 2

compiles to this bytecode with the optimisation disabled (left column is
bytecode offset in bytes):

    00 LOAD_FAST 0
    01 POP_JUMP_IF_FALSE 8
    04 LOAD_CONST_SMALL_INT 1
    05 JUMP 9
    08 LOAD_CONST_SMALL_INT 2
    09 RETURN_VALUE

and to this bytecode with the optimisation enabled:

    00 LOAD_FAST 0
    01 POP_JUMP_IF_FALSE 6
    04 LOAD_CONST_SMALL_INT 1
    05 RETURN_VALUE
    06 LOAD_CONST_SMALL_INT 2
    07 RETURN_VALUE

So the JUMP to RETURN_VALUE is optimised and replaced by RETURN_VALUE,
saving 2 bytes and making the code a bit faster.
diff --git a/py/compile.c b/py/compile.c
@@ -980,7 +980,8 @@ STATIC void compile_return_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
     if (MP_PARSE_NODE_IS_NULL(pns->nodes[0])) {
         // no argument to 'return', so return None
         EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
-    } else if (MP_PARSE_NODE_IS_STRUCT_KIND(pns->nodes[0], PN_test_if_expr)) {
+    } else if (MICROPY_COMP_RETURN_IF_EXPR
+        && MP_PARSE_NODE_IS_STRUCT_KIND(pns->nodes[0], PN_test_if_expr)) {
         // special case when returning an if-expression; to match CPython optimisation
         mp_parse_node_struct_t *pns_test_if_expr = (mp_parse_node_struct_t*)pns->nodes[0];
         mp_parse_node_struct_t *pns_test_if_else = (mp_parse_node_struct_t*)pns_test_if_expr->nodes[1];
diff --git a/py/mpconfig.h b/py/mpconfig.h
@@ -353,6 +353,12 @@
 #define MICROPY_COMP_TRIPLE_TUPLE_ASSIGN (0)
 #endif
 
+// Whether to enable optimisation of: return a if b else c
+// Costs about 80 bytes (Thumb2) and saves 2 bytes of bytecode for each use
+#ifndef MICROPY_COMP_RETURN_IF_EXPR
+#define MICROPY_COMP_RETURN_IF_EXPR (0)
+#endif
+
 /*****************************************************************************/
 /* Internal debugging stuff                                                  */
 
