From ef3ef6f4f89a1f14c1885e5b65b55fdacf3448e4 Mon Sep 17 00:00:00 2001
From: Alessandro Gatti <a.gatti@frob.it>
Date: Wed, 9 Jul 2025 01:41:39 +0200
Subject: [PATCH] py/parse: Recognise const int assignments with type hints.

This commit extends the parser's behaviour to also recognise constant
integer assignments even when the assignment itself has an "int" type
annotation.

Now declarations like "var: int = const(val)" can be treated as
constants by the compiler and thus be folded in the rest of the program
as it sees fit.  Before this change, that line would generate a variable
creation and its value assignment, without any folding being performed.

This fixes #15608.

Signed-off-by: Alessandro Gatti <a.gatti@frob.it>
---
 py/parse.c                                | 28 +++++++++----
 tests/micropython/const2_annotated.py     | 48 +++++++++++++++++++++++
 tests/micropython/const2_annotated.py.exp |  4 ++
 tests/micropython/const_annotated.py      | 30 ++++++++++++++
 tests/micropython/const_annotated.py.exp  |  5 +++
 5 files changed, 108 insertions(+), 7 deletions(-)
 create mode 100644 tests/micropython/const2_annotated.py
 create mode 100644 tests/micropython/const2_annotated.py.exp
 create mode 100644 tests/micropython/const_annotated.py
 create mode 100644 tests/micropython/const_annotated.py.exp

diff --git a/py/parse.c b/py/parse.c
index db89fb58450e3..fc14afca0a476 100644
--- a/py/parse.c
+++ b/py/parse.c
@@ -705,6 +705,15 @@ static bool fold_logical_constants(parser_t *parser, uint8_t rule_id, size_t *nu
     return false;
 }
 
+#if MICROPY_COMP_CONST
+static bool is_const_atom_expr_normal(mp_parse_node_t node) {
+    return MP_PARSE_NODE_IS_STRUCT_KIND(node, RULE_atom_expr_normal)
+           && MP_PARSE_NODE_IS_ID(((mp_parse_node_struct_t *)node)->nodes[0])
+           && MP_PARSE_NODE_LEAF_ARG(((mp_parse_node_struct_t *)node)->nodes[0]) == MP_QSTR_const
+           && MP_PARSE_NODE_IS_STRUCT_KIND(((mp_parse_node_struct_t *)node)->nodes[1], RULE_trailer_paren);
+}
+#endif
+
 static bool fold_constants(parser_t *parser, uint8_t rule_id, size_t num_args) {
     // this code does folding of arbitrary integer expressions, eg 1 + 2 * 3 + 4
     // it does not do partial folding, eg 1 + 2 + x -> 3 + x
@@ -797,21 +806,26 @@ static bool fold_constants(parser_t *parser, uint8_t rule_id, size_t num_args) {
         if (!MP_PARSE_NODE_IS_NULL(pn1)
             && !(MP_PARSE_NODE_IS_STRUCT_KIND(pn1, RULE_expr_stmt_augassign)
                  || MP_PARSE_NODE_IS_STRUCT_KIND(pn1, RULE_expr_stmt_assign_list))) {
-            // this node is of the form <x> = <y>
+            // this node is of the form <x> = <y> or <x> : int = <y>
             mp_parse_node_t pn0 = peek_result(parser, 1);
             if (MP_PARSE_NODE_IS_ID(pn0)
-                && MP_PARSE_NODE_IS_STRUCT_KIND(pn1, RULE_atom_expr_normal)
-                && MP_PARSE_NODE_IS_ID(((mp_parse_node_struct_t *)pn1)->nodes[0])
-                && MP_PARSE_NODE_LEAF_ARG(((mp_parse_node_struct_t *)pn1)->nodes[0]) == MP_QSTR_const
-                && MP_PARSE_NODE_IS_STRUCT_KIND(((mp_parse_node_struct_t *)pn1)->nodes[1], RULE_trailer_paren)
-                ) {
+                && ((MP_PARSE_NODE_IS_STRUCT_KIND(pn1, RULE_annassign) &&
+                     MP_PARSE_NODE_IS_ID(((mp_parse_node_struct_t *)pn1)->nodes[0]) &&
+                     MP_PARSE_NODE_LEAF_ARG(((mp_parse_node_struct_t *)pn1)->nodes[0]) == MP_QSTR_int &&
+                     is_const_atom_expr_normal(((mp_parse_node_struct_t *)pn1)->nodes[1]))
+                    || is_const_atom_expr_normal(pn1))) {
                 // code to assign dynamic constants: id = const(value)
 
                 // get the id
                 qstr id = MP_PARSE_NODE_LEAF_ARG(pn0);
 
                 // get the value
-                mp_parse_node_t pn_value = ((mp_parse_node_struct_t *)((mp_parse_node_struct_t *)pn1)->nodes[1])->nodes[0];
+                mp_parse_node_t pn_value;
+                if (MP_PARSE_NODE_IS_STRUCT_KIND(pn1, RULE_annassign)) {
+                    pn_value = ((mp_parse_node_struct_t *)(((mp_parse_node_struct_t *)(((mp_parse_node_struct_t *)pn1)->nodes[1]))->nodes[1]))->nodes[0];
+                } else {
+                    pn_value = ((mp_parse_node_struct_t *)((mp_parse_node_struct_t *)pn1)->nodes[1])->nodes[0];
+                }
                 if (!mp_parse_node_is_const(pn_value)) {
                     mp_obj_t exc = mp_obj_new_exception_msg(&mp_type_SyntaxError,
                         MP_ERROR_TEXT("not a constant"));
diff --git a/tests/micropython/const2_annotated.py b/tests/micropython/const2_annotated.py
new file mode 100644
index 0000000000000..fde653add7550
--- /dev/null
+++ b/tests/micropython/const2_annotated.py
@@ -0,0 +1,48 @@
+# check that consts are not replaced in anything except standalone identifiers
+
+from micropython import const
+
+X: int = const(1)
+Y: int = const(2)
+Z: int = const(3)
+
+# import that uses a constant
+import micropython as X
+
+print(globals()["X"])
+
+
+# function name that matches a constant
+def X():
+    print("function X", X)
+
+
+globals()["X"]()
+
+
+# arguments that match a constant
+def f(X, *Y, **Z):
+    pass
+
+
+f(1)
+
+
+# class name that matches a constant
+class X:
+    def f(self):
+        print("class X", X)
+
+
+globals()["X"]().f()
+
+
+# constant within a class
+class A:
+    C1: int = const(4)
+
+    def X(self):
+        print("method X", Y, C1, self.C1)
+
+
+A().X()
diff --git a/tests/micropython/const2_annotated.py.exp b/tests/micropython/const2_annotated.py.exp
new file mode 100644
index 0000000000000..0568f91ce2860
--- /dev/null
+++ b/tests/micropython/const2_annotated.py.exp
@@ -0,0 +1,4 @@
+<module 'micropython'>
+function X 1
+class X 1
+method X 2 4 4
diff --git a/tests/micropython/const_annotated.py b/tests/micropython/const_annotated.py
new file mode 100644
index 0000000000000..3af7dd3fceda3
--- /dev/null
+++ b/tests/micropython/const_annotated.py
@@ -0,0 +1,30 @@
+# test constant optimisation
+# This test will only work when MICROPY_COMP_CONST is enabled.
+
+from micropython import const
+
+X: int = const(123)
+Y: int = const(X + 456)
+
+print(X, Y + 1)
+
+
+def f():
+    print(X, Y + 1)
+
+
+f()
+
+_X: int = const(12)
+_Y: int = const(_X + 34)
+
+print(_X, _Y)
+
+
+class A:
+    Z: int = const(1)
+    _Z: int = const(2)
+    print(Z, _Z)
+
+
+print(hasattr(A, "Z"), hasattr(A, "_Z"))
diff --git a/tests/micropython/const_annotated.py.exp b/tests/micropython/const_annotated.py.exp
new file mode 100644
index 0000000000000..ece6a5cb2e552
--- /dev/null
+++ b/tests/micropython/const_annotated.py.exp
@@ -0,0 +1,5 @@
+123 580
+123 580
+12 46
+1 2
+True False

<!DOCTYPE html PUBLIC '-//W3C//DTD XHTML 1.0 Transitional//EN' 'http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd'>
<html xmlns='http://www.w3.org/1999/xhtml'>
<head>
<title>pFad - Phonifier reborn</title>
<meta http-equiv='Content-Type' content='text/html; charset=utf-8' />
</head>
<body>
<h1>Pfad - The Proxy pFad of &#169; 2024 Garber Painting. All rights reserved.</h1>


<!-- Disclaimer -->
<p>Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.</p>
<br>
<p>Alternative Proxies:</p><p><a href="http://rainy.clevelandohioweatherforecast.com/php-proxy/index.php?q=https://patch-diff.githubusercontent.com/raw/micropython/micropython/pull/17643.patch" target="_blank">Alternative Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/index.php?u=https://patch-diff.githubusercontent.com/raw/micropython/micropython/pull/17643.patch" target="_blank">pFad Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/v3index.php?u=https://patch-diff.githubusercontent.com/raw/micropython/micropython/pull/17643.patch" target="_blank">pFad v3 Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/v4index.php?u=https://patch-diff.githubusercontent.com/raw/micropython/micropython/pull/17643.patch" target="_blank">pFad v4 Proxy</a></p></body>
</html>