py: Add simple way of looking up constants in compiler.

Working towards trying to support compile-time constants (see discussion
in issue #227), this patch allows the compiler to look inside arbitrary
uPy objects at compile time.  The objects to search are given by the
macro MICROPY_EXTRA_CONSTANTS (so they must be constant/ROM objects),
and the constant folding occures on forms base.attr (both base and attr
must be id's).

It works, but it breaks strict CPython compatibility, since the lookup
will succeed even without importing the namespace.

2 files changed, 42 insertions, 4 deletions

diff --git a/py/compile.c b/py/compile.c
index ca3ee9d6d1..4738c0ae55 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -78,6 +78,19 @@ STATIC void compile_syntax_error(compiler_t *comp, mp_parse_node_t pn, const cha
     comp->had_error = true;
 }
 
+STATIC const mp_map_elem_t mp_constants_table[] = {
+    // Extra constants as defined by a port
+    MICROPY_EXTRA_CONSTANTS
+};
+
+STATIC const mp_map_t mp_constants_map = {
+    .all_keys_are_qstrs = 1,
+    .table_is_fixed_array = 1,
+    .used = sizeof(mp_constants_table) / sizeof(mp_map_elem_t),
+    .alloc = sizeof(mp_constants_table) / sizeof(mp_map_elem_t),
+    .table = (mp_map_elem_t*)mp_constants_table,
+};
+
 mp_parse_node_t fold_constants(mp_parse_node_t pn) {
     if (MP_PARSE_NODE_IS_STRUCT(pn)) {
         mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
@@ -168,10 +181,12 @@ mp_parse_node_t fold_constants(mp_parse_node_t pn) {
                 }
                 break;
 
-#if MICROPY_EMIT_CPYTHON
             case PN_power:
-                // can overflow; enabled only to compare with CPython
-                if (MP_PARSE_NODE_IS_SMALL_INT(pns->nodes[0]) && MP_PARSE_NODE_IS_NULL(pns->nodes[1]) && !MP_PARSE_NODE_IS_NULL(pns->nodes[2])) {
+                if (0) {
+#if MICROPY_EMIT_CPYTHON
+                } else if (MP_PARSE_NODE_IS_SMALL_INT(pns->nodes[0]) && MP_PARSE_NODE_IS_NULL(pns->nodes[1]) && !MP_PARSE_NODE_IS_NULL(pns->nodes[2])) {
+                    // int**x
+                    // can overflow; enabled only to compare with CPython
                     mp_parse_node_struct_t* pns2 = (mp_parse_node_struct_t*)pns->nodes[2];
                     if (MP_PARSE_NODE_IS_SMALL_INT(pns2->nodes[0])) {
                         int power = MP_PARSE_NODE_LEAF_SMALL_INT(pns2->nodes[0]);
@@ -184,9 +199,27 @@ mp_parse_node_t fold_constants(mp_parse_node_t pn) {
                             pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, ans);
                         }
                     }
+#endif
+                } else if (MP_PARSE_NODE_IS_ID(pns->nodes[0]) && MP_PARSE_NODE_IS_STRUCT_KIND(pns->nodes[1], PN_trailer_period) && MP_PARSE_NODE_IS_NULL(pns->nodes[2])) {
+                    // id.id
+                    // look it up in constant table, see if it can be replaced with an integer
+                    mp_parse_node_struct_t* pns1 = (mp_parse_node_struct_t*)pns->nodes[1];
+                    assert(MP_PARSE_NODE_IS_ID(pns1->nodes[0]));
+                    qstr q_base = MP_PARSE_NODE_LEAF_ARG(pns->nodes[0]);
+                    qstr q_attr = MP_PARSE_NODE_LEAF_ARG(pns1->nodes[0]);
+                    mp_map_elem_t *elem = mp_map_lookup((mp_map_t*)&mp_constants_map, MP_OBJ_NEW_QSTR(q_base), MP_MAP_LOOKUP);
+                    if (elem != NULL) {
+                        mp_obj_t dest[2];
+                        mp_load_method_maybe(elem->value, q_attr, dest);
+                        if (MP_OBJ_IS_SMALL_INT(dest[0]) && dest[1] == NULL) {
+                            machine_int_t val = MP_OBJ_SMALL_INT_VALUE(dest[0]);
+                            if (MP_PARSE_FITS_SMALL_INT(val)) {
+                                pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, val);
+                            }
+                        }
+                    }
                 }
                 break;
-#endif
         }
     }
 
diff --git a/py/mpconfig.h b/py/mpconfig.h
index 1f2862ab0e..b120c4bb4a 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -150,6 +150,11 @@ typedef double mp_float_t;
 #define MICROPY_EXTRA_BUILTIN_MODULES
 #endif
 
+// Additional constant definitions for the compiler - see compile.c:mp_constants_table.
+#ifndef MICROPY_EXTRA_CONSTANTS
+#define MICROPY_EXTRA_CONSTANTS
+#endif
+
 /*****************************************************************************/
 /* Miscellaneous settings                                                    */