From 4ede70368722f4702320412f393a7c703cdf4276 Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Tue, 7 Apr 2020 12:23:08 +1000
Subject: [PATCH] py/parse: Support constant folding of power operator for
 integers.

Constant expression like "2 ** 3" will now be folded, and the special form
"X = const(2 ** 3)" will now compile because the argument to the const is
now a constant.

Fixes issue #5865.
---
 py/parse.c                       | 13 ++++++++++---
 tests/basics/int_constfolding.py |  4 ++++
 tests/micropython/const_error.py |  2 +-
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/py/parse.c b/py/parse.c
index 35f06afea6..b86b919a86 100644
--- a/py/parse.c
+++ b/py/parse.c
@@ -618,8 +618,9 @@ STATIC bool fold_constants(parser_t *parser, uint8_t rule_id, size_t num_args) {
     mp_obj_t arg0;
     if (rule_id == RULE_expr
         || rule_id == RULE_xor_expr
-        || rule_id == RULE_and_expr) {
-        // folding for binary ops: | ^ &
+        || rule_id == RULE_and_expr
+        || rule_id == RULE_power) {
+        // folding for binary ops: | ^ & **
         mp_parse_node_t pn = peek_result(parser, num_args - 1);
         if (!mp_parse_node_get_int_maybe(pn, &arg0)) {
             return false;
@@ -629,8 +630,10 @@ STATIC bool fold_constants(parser_t *parser, uint8_t rule_id, size_t num_args) {
             op = MP_BINARY_OP_OR;
         } else if (rule_id == RULE_xor_expr) {
             op = MP_BINARY_OP_XOR;
-        } else {
+        } else if (rule_id == RULE_and_expr) {
             op = MP_BINARY_OP_AND;
+        } else {
+            op = MP_BINARY_OP_POWER;
         }
         for (ssize_t i = num_args - 2; i >= 0; --i) {
             pn = peek_result(parser, i);
@@ -638,6 +641,10 @@ STATIC bool fold_constants(parser_t *parser, uint8_t rule_id, size_t num_args) {
             if (!mp_parse_node_get_int_maybe(pn, &arg1)) {
                 return false;
             }
+            if (op == MP_BINARY_OP_POWER && mp_obj_int_sign(arg1) < 0) {
+                // ** can't have negative rhs
+                return false;
+            }
             arg0 = mp_binary_op(op, arg0, arg1);
         }
     } else if (rule_id == RULE_shift_expr
diff --git a/tests/basics/int_constfolding.py b/tests/basics/int_constfolding.py
index 158897f553..3168685ff5 100644
--- a/tests/basics/int_constfolding.py
+++ b/tests/basics/int_constfolding.py
@@ -30,6 +30,10 @@ print(-123 // 7, -123 % 7)
 print(123 // -7, 123 % -7)
 print(-123 // -7, -123 % -7)
 
+# power
+print(2 ** 3)
+print(3 ** 4)
+
 # won't fold so an exception can be raised at runtime
 try:
     1 << -1
diff --git a/tests/micropython/const_error.py b/tests/micropython/const_error.py
index 395fe0f776..d35be530a7 100644
--- a/tests/micropython/const_error.py
+++ b/tests/micropython/const_error.py
@@ -19,7 +19,7 @@ test_syntax("A = const(1); A = const(2)")
 # these operations are not supported within const
 test_syntax("A = const(1 @ 2)")
 test_syntax("A = const(1 / 2)")
-test_syntax("A = const(1 ** 2)")
+test_syntax("A = const(1 ** -2)")
 test_syntax("A = const(1 << -2)")
 test_syntax("A = const(1 >> -2)")
 test_syntax("A = const(1 % 0)")