From ae54fbf1668959e50632034495631b0c5952ae9d Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Sat, 22 Apr 2017 14:58:01 +1000
Subject: [PATCH] py/compile: Add COMP_RETURN_IF_EXPR option to enable
 return-if-else opt.

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.
---
 py/compile.c  | 3 ++-
 py/mpconfig.h | 6 ++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/py/compile.c b/py/compile.c
index 3aa70e8bf2..ae5c3b2dc8 100644
--- 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
index b6e24d2055..9122c99163 100644
--- 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                                                  */