py/lexer: Add support for nested f-strings within f-strings.

It turns out that it's relatively simple to support nested f-strings, which
is what this commit implements.

The way the MicroPython f-string parser works at the moment is:
1. it extracts the f-string arguments (things in curly braces) into a
   temporary buffer (a vstr)
2. once the f-string ends (reaches its closing quote) the lexer switches to
   tokenizing the temporary buffer
3. once the buffer is empty it switches back to the stream.

The temporary buffer can easily hold f-strings itself (ie nested f-strings)
and they can be re-parsed by the lexer using the same algorithm.  The only
thing stopping that from working is that the temporary buffer can't be
reused for the nested f-string because it's currently being parsed.

This commit fixes that by adding a second temporary buffer, which is the
"injection" buffer.  That allows arbitrary number of nestings with a simple
modification to the original algorithm:
1. when an f-string is encountered the string is parsed and its arguments
   are extracted into `fstring_args`
2. when the f-string finishes, `fstring_args` is inserted into the current
   position in `inject_chrs` (which is the start of that buffer if no
   injection is ongoing)
3. `fstring_args` is now cleared and ready for any further f-strings
   (nested or not)
4. the lexer switches to `inject_chrs` if it's not already reading from it
5. if an f-string appeared inside the f-string then it is in `inject_chrs`
   and can be processed as before, extracting its arguments into
   `fstring_args`, which can then be inserted again into `inject_chrs`
6. once `inject_chrs` is exhausted (meaning that all levels of f-strings
   have been fully processed) the lexer switched back to tokenizing the
   stream.

Amazingly, this scheme supports arbitrary numbers of nestings of f-strings
using the same quote style.

This adds some code size and a bit more memory usage for the lexer.  In
particular for a single (non-nested) f-string it now makes an extra copy of
the `fstring_args` data, when copying it across to `inject_chrs`.
Otherwise, memory use only goes up with the complexity of nested f-strings.

Signed-off-by: Damien George <damien@micropython.org>

Author: dpgeorge

py/lexer.c

@@ -145,27 +145,17 @@ static void next_char(mp_lexer_t *lex) {
     lex->chr0 = lex->chr1;
     lex->chr1 = lex->chr2;
 
-    // and add the next byte from either the fstring args or the reader
+    // and add the next byte from either inject_chrs or the reader
     mp_uint_t chr2;
 fetch_next_byte:
     #if MICROPY_PY_FSTRINGS
-    if (lex->fstring_args_idx) {
-        // if there are saved chars, then we're currently injecting fstring args
-        if (lex->fstring_args_idx < lex->fstring_args.len) {
-            chr2 = lex->fstring_args.buf[lex->fstring_args_idx++];
-        } else {
-            // no more fstring arg bytes
-            chr2 = '\0';
-        }
-
-        if (lex->chr0 == '\0') {
-            // consumed all fstring data, restore saved input queue
-            lex->chr0 = lex->chr0_saved;
-            lex->chr1 = lex->chr1_saved;
-            chr2 = lex->chr2_saved;
-            // stop consuming fstring arg data
-            vstr_reset(&lex->fstring_args);
-            lex->fstring_args_idx = 0;
+    if (lex->inject_chrs_idx) {
+        // if there are saved chars, then we're currently injecting them
+        chr2 = lex->inject_chrs.buf[lex->inject_chrs_idx++];
+        if (lex->inject_chrs_idx >= lex->inject_chrs.len) {
+            // consumed all injected characters, switch back to the input stream
+            vstr_reset(&lex->inject_chrs);
+            lex->inject_chrs_idx = 0;
         }
     } else
     #endif
@@ -346,8 +336,7 @@ static void parse_string_literal(mp_lexer_t *lex, bool is_raw, bool is_fstring)
     #if MICROPY_PY_FSTRINGS
     if (is_fstring) {
         // assume there's going to be interpolation, so prep the injection data
-        // fstring_args_idx==0 && len(fstring_args)>0 means we're extracting the args.
-        // only when fstring_args_idx>0 will we consume the arg data
+        // len(fstring_args)>0 means we're extracting the args.
         // lex->fstring_args is reset when finished, so at this point there are two cases:
         // - lex->fstring_args is empty: start of a new f-string
         // - lex->fstring_args is non-empty: concatenation of adjacent f-strings
@@ -570,19 +559,26 @@ static bool skip_whitespace(mp_lexer_t *lex, bool stop_at_newline) {
 
 void mp_lexer_to_next(mp_lexer_t *lex) {
     #if MICROPY_PY_FSTRINGS
-    if (lex->fstring_args.len && lex->fstring_args_idx == 0) {
+    if (lex->fstring_args.len) {
         // moving onto the next token means the literal string is complete.
         // switch into injecting the format args.
         vstr_add_byte(&lex->fstring_args, ')');
-        lex->chr0_saved = lex->chr0;
-        lex->chr1_saved = lex->chr1;
-        lex->chr2_saved = lex->chr2;
-        lex->chr0 = lex->fstring_args.buf[0];
-        lex->chr1 = lex->fstring_args.buf[1];
-        lex->chr2 = lex->fstring_args.buf[2];
-        // we've already extracted 3 chars, but setting this non-zero also
-        // means we'll start consuming the fstring data
-        lex->fstring_args_idx = 3;
+        if (lex->inject_chrs_idx == 0) {
+            // switch from stream to inject_chrs
+            char *s = vstr_add_len(&lex->inject_chrs, 3);
+            s[0] = lex->chr0;
+            s[1] = lex->chr1;
+            s[2] = lex->chr2;
+        } else {
+            // already consuming from inject_chrs, rewind cached chars to insert new ones
+            assert(lex->inject_chrs_idx >= 3);
+            lex->inject_chrs_idx -= 3;
+        }
+        vstr_ins_strn(&lex->inject_chrs, lex->inject_chrs_idx, lex->fstring_args.buf, lex->fstring_args.len);
+        vstr_reset(&lex->fstring_args);
+        lex->chr0 = lex->inject_chrs.buf[lex->inject_chrs_idx++];
+        lex->chr1 = lex->inject_chrs.buf[lex->inject_chrs_idx++];
+        lex->chr2 = lex->inject_chrs.buf[lex->inject_chrs_idx++];
     }
     #endif
 
@@ -867,8 +863,9 @@ mp_lexer_t *mp_lexer_new(qstr src_name, mp_reader_t reader) {
     lex->indent_level = m_new(uint16_t, lex->alloc_indent_level);
     vstr_init(&lex->vstr, 32);
     #if MICROPY_PY_FSTRINGS
+    vstr_init(&lex->inject_chrs, 0);
+    lex->inject_chrs_idx = 0;
     vstr_init(&lex->fstring_args, 0);
-    lex->fstring_args_idx = 0;
     #endif
 
     // store sentinel for first indentation level
@@ -925,6 +922,7 @@ void mp_lexer_free(mp_lexer_t *lex) {
         lex->reader.close(lex->reader.data);
         vstr_clear(&lex->vstr);
         #if MICROPY_PY_FSTRINGS
+        vstr_clear(&lex->inject_chrs);
         vstr_clear(&lex->fstring_args);
         #endif
         m_del(uint16_t, lex->indent_level, lex->alloc_indent_level);

py/lexer.h

@@ -164,9 +164,6 @@ typedef struct _mp_lexer_t {
 
     uint32_t chr0;              // first cached byte from source (32-bits for efficient access)
     uint8_t chr1, chr2;         // subsequent cached bytes from source
-    #if MICROPY_PY_FSTRINGS
-    uint8_t chr0_saved, chr1_saved, chr2_saved; // current cached bytes from alt source
-    #endif
 
     size_t line;                // current source line
     size_t column;              // current source column
@@ -183,8 +180,9 @@ typedef struct _mp_lexer_t {
     mp_token_kind_t tok_kind;   // token kind
     vstr_t vstr;                // token data
     #if MICROPY_PY_FSTRINGS
+    vstr_t inject_chrs;         // characters currently being injected into the stream
+    size_t inject_chrs_idx;     // current index into inject_chrs
     vstr_t fstring_args;        // extracted arguments to pass to .format()
-    size_t fstring_args_idx;    // how many bytes of fstring_args have been read
     #endif
 } mp_lexer_t;
 

tests/basics/string_fstring_nested.py

@@ -0,0 +1,9 @@
+# Test nesting of f-strings within f-strings.
+
+x = 1
+
+# 2-level nesting, with padding.
+print(f"a{f'b{x:2}c':>5}d")
+
+# 4-level nesting using the different styles of quotes.
+print(f"""a{f'''b{f"c{f'd{x}e'}f"}g'''}h""")

tests/basics/string_fstring_nested_py312.py

@@ -0,0 +1,7 @@
+# Test nesting of f-strings within f-strings.
+# These test rely on Python 3.12+ to use the same quote style for nesting.
+
+x = 1
+
+# 8-level nesting using the same quote style.
+print(f"a{f"b{f"c{f"d{f"e{f"f{f"g{f"h{x}i"}j"}k"}l"}m"}n"}o"}p")

tests/basics/string_fstring_nested_py312.py.exp

@@ -0,0 +1 @@
+abcdefgh1ijklmnop