py/qstr: Add support for sorted qstr pools.

jimmo · jimmo · commit 354d5ba2e348 · 2023-10-09T13:37:43.000+11:00
This provides a significant performance boost for qstr_find_strn, which is
called a lot during parsing and loading of .mpy files, as well as
interning of string objects (which happens in most string methods that
return new strings).

This work was funded through GitHub Sponsors.

Signed-off-by: Jim Mussared &lt;jim.mussared@gmail.com&gt;
diff --git a/py/makeqstrdata.py b/py/makeqstrdata.py
@@ -220,6 +220,66 @@
     "values",
     "write",
     "zip",
+    # Additional QSTRs that must have index <255 (these are not part of
+    # the .mpy compatibility list though).
+    "__bool__",
+    "__pos__",
+    "__neg__",
+    "__invert__",
+    "__abs__",
+    "__float__",
+    "__complex__",
+    "__sizeof__",
+    "__lt__",
+    "__gt__",
+    "__eq__",
+    "__le__",
+    "__ge__",
+    "__ne__",
+    "__contains__",
+    "__iadd__",
+    "__isub__",
+    "__imul__",
+    "__imatmul__",
+    "__ifloordiv__",
+    "__itruediv__",
+    "__imod__",
+    "__ipow__",
+    "__ior__",
+    "__ixor__",
+    "__iand__",
+    "__ilshift__",
+    "__irshift__",
+    "__add__",
+    "__sub__",
+    "__mul__",
+    "__matmul__",
+    "__floordiv__",
+    "__truediv__",
+    "__mod__",
+    "__divmod__",
+    "__pow__",
+    "__or__",
+    "__xor__",
+    "__and__",
+    "__lshift__",
+    "__rshift__",
+    "__radd__",
+    "__rsub__",
+    "__rmul__",
+    "__rmatmul__",
+    "__rfloordiv__",
+    "__rtruediv__",
+    "__rmod__",
+    "__rpow__",
+    "__ror__",
+    "__rxor__",
+    "__rand__",
+    "__rlshift__",
+    "__rrshift__",
+    "__get__",
+    "__set__",
+    "__delete__",
 ]
 
 
@@ -244,22 +304,13 @@ def esc_char(m):
     return re.sub(r"[^A-Za-z0-9_]", esc_char, qst)
 
 
+static_qstr_list_ident = list(map(qstr_escape, static_qstr_list))
+
+
 def parse_input_headers(infiles):
     qcfgs = {}
     qstrs = {}
 
-    # add static qstrs
-    for qstr in static_qstr_list:
-        # work out the corresponding qstr name
-        ident = qstr_escape(qstr)
-
-        # don't add duplicates
-        assert ident not in qstrs
-
-        # add the qstr to the list, with order number to retain original order in file
-        order = len(qstrs) - 300000
-        qstrs[ident] = (order, ident, qstr)
-
     # read the qstrs in from the input files
     for infile in infiles:
         with open(infile, "rt") as f:
@@ -294,22 +345,12 @@ def parse_input_headers(infiles):
                 ident = qstr_escape(qstr)
 
                 # don't add duplicates
+                if ident in static_qstr_list_ident:
+                    continue
                 if ident in qstrs:
                     continue
 
-                # add the qstr to the list, with order number to retain original order in file
-                order = len(qstrs)
-                # but put special method names like __add__ at the top of list, so
-                # that their id's fit into a byte
-                if ident == "":
-                    # Sort empty qstr above all still
-                    order = -200000
-                elif ident == "__dir__":
-                    # Put __dir__ after empty qstr for builtin dir() to work
-                    order = -190000
-                elif ident.startswith("__"):
-                    order -= 100000
-                qstrs[ident] = (order, ident, qstr)
+                qstrs[ident] = (ident, qstr)
 
     if not qcfgs:
         sys.stderr.write("ERROR: Empty preprocessor output - check for errors above\n")
@@ -348,12 +389,17 @@ def print_qstr_data(qcfgs, qstrs):
     print("")
 
     # add NULL qstr with no hash or data
-    print('QDEF(MP_QSTRnull, 0, 0, "")')
+    print('QDEF0(MP_QSTRnull, 0, 0, "")')
+
+    # add static qstrs to the first unsorted pool
+    for qstr in static_qstr_list:
+        qbytes = make_bytes(cfg_bytes_len, cfg_bytes_hash, qstr)
+        print("QDEF0(MP_QSTR_%s, %s)" % (qstr_escape(qstr), qbytes))
 
-    # go through each qstr and print it out
-    for order, ident, qstr in sorted(qstrs.values(), key=lambda x: x[0]):
+    # add remaining qstrs to the sorted (by value) pool
+    for ident, qstr in sorted(qstrs.values(), key=lambda x: x[1]):
         qbytes = make_bytes(cfg_bytes_len, cfg_bytes_hash, qstr)
-        print("QDEF(MP_QSTR_%s, %s)" % (ident, qbytes))
+        print("QDEF1(MP_QSTR_%s, %s)" % (ident, qbytes))
 
 
 def do_work(infiles):
diff --git a/py/qstr.c b/py/qstr.c
@@ -79,37 +79,83 @@ size_t qstr_compute_hash(const byte *data, size_t len) {
 }
 
 #if MICROPY_QSTR_BYTES_IN_HASH
+const qstr_hash_t mp_qstr_const_hashes_special[] = {
+    #ifndef NO_QSTR
+#define QDEF0(id, hash, len, str) hash,
+#define QDEF1(id, hash, len, str)
+    #include "genhdr/qstrdefs.generated.h"
+#undef QDEF0
+#undef QDEF1
+    #endif
+};
 const qstr_hash_t mp_qstr_const_hashes[] = {
     #ifndef NO_QSTR
-#define QDEF(id, hash, len, str) hash,
+#define QDEF0(id, hash, len, str)
+#define QDEF1(id, hash, len, str) hash,
     #include "genhdr/qstrdefs.generated.h"
-#undef QDEF
+#undef QDEF0
+#undef QDEF1
     #endif
 };
 #endif
 
+const qstr_len_t mp_qstr_const_lengths_special[] = {
+    #ifndef NO_QSTR
+#define QDEF0(id, hash, len, str) len,
+#define QDEF1(id, hash, len, str)
+    #include "genhdr/qstrdefs.generated.h"
+#undef QDEF0
+#undef QDEF1
+    #endif
+};
 const qstr_len_t mp_qstr_const_lengths[] = {
     #ifndef NO_QSTR
-#define QDEF(id, hash, len, str) len,
+#define QDEF0(id, hash, len, str)
+#define QDEF1(id, hash, len, str) len,
     #include "genhdr/qstrdefs.generated.h"
-#undef QDEF
+#undef QDEF0
+#undef QDEF1
     #endif
 };
 
-const qstr_pool_t mp_qstr_const_pool = {
+const qstr_pool_t mp_qstr_const_pool_special = {
     NULL,               // no previous pool
     0,                  // no previous pool
     MICROPY_ALLOC_QSTR_ENTRIES_INIT,
-    MP_QSTRnumber_of,   // corresponds to number of strings in array just below
+    MP_QSTRnumber_of_special,   // corresponds to number of strings in array just below
+    false,              // not sorted
+    #if MICROPY_QSTR_BYTES_IN_HASH
+    (qstr_hash_t *)mp_qstr_const_hashes_special,
+    #endif
+    (qstr_len_t *)mp_qstr_const_lengths_special,
+    {
+        #ifndef NO_QSTR
+#define QDEF0(id, hash, len, str) str,
+#define QDEF1(id, hash, len, str)
+        #include "genhdr/qstrdefs.generated.h"
+#undef QDEF0
+#undef QDEF1
+        #endif
+    },
+};
+
+const qstr_pool_t mp_qstr_const_pool = {
+    &mp_qstr_const_pool_special,
+    MP_QSTRnumber_of_special,
+    MICROPY_ALLOC_QSTR_ENTRIES_INIT,
+    MP_QSTRnumber_of - MP_QSTRnumber_of_special,   // corresponds to number of strings in array just below
+    true,              // sorted
     #if MICROPY_QSTR_BYTES_IN_HASH
     (qstr_hash_t *)mp_qstr_const_hashes,
     #endif
     (qstr_len_t *)mp_qstr_const_lengths,
     {
         #ifndef NO_QSTR
-#define QDEF(id, hash, len, str) str,
+#define QDEF0(id, hash, len, str)
+#define QDEF1(id, hash, len, str) str,
         #include "genhdr/qstrdefs.generated.h"
-#undef QDEF
+#undef QDEF0
+#undef QDEF1
         #endif
     },
 };
@@ -205,6 +251,24 @@ STATIC qstr qstr_add(
     return MP_STATE_VM(last_pool)->total_prev_len + at;
 }
 
+STATIC inline int qstr_strncmp(const char *a, size_t a_len, const char *b) {
+    while (a_len && *b) {
+        if (*a != *b) {
+            return (int)*a - (int)*b;
+        }
+        a++;
+        b++;
+        a_len--;
+    }
+    if (a_len == 0 && *b) {
+        return -1;
+    }
+    if (a_len && !*b) {
+        return 1;
+    }
+    return 0;
+}
+
 qstr qstr_find_strn(const char *str, size_t str_len) {
     // work out hash of str
     #if MICROPY_QSTR_BYTES_IN_HASH
@@ -213,7 +277,30 @@ qstr qstr_find_strn(const char *str, size_t str_len) {
 
     // search pools for the data
     for (const qstr_pool_t *pool = MP_STATE_VM(last_pool); pool != NULL; pool = pool->prev) {
-        for (mp_uint_t at = 0, top = pool->len; at < top; at++) {
+        size_t low = 0;
+        size_t high = pool->len - 1;
+
+        // binary search inside the pool
+        if (pool->sorted) {
+            while (high - low > 1) {
+                size_t mid = (low + high) / 2;
+                int cmp = qstr_strncmp(str, str_len, pool->qstrs[mid]);// strncmp();//pool->hashes[mid] - str_hash;
+                if (cmp < 0) {
+                    high = mid;
+                } else {
+                    low = mid;
+                    if (MP_UNLIKELY(cmp == 0)) {
+                        // while (low > 0 && pool->hashes[low - 1] == str_hash) {
+                        //     low--;
+                        // }
+                        break;
+                    }
+                }
+            }
+        }
+
+        // sequential search for the remaining strings
+        for (mp_uint_t at = low; at < high + 1; at++) {
             if (
                 #if MICROPY_QSTR_BYTES_IN_HASH
                 pool->hashes[at] == str_hash &&
diff --git a/py/qstr.h b/py/qstr.h
@@ -38,9 +38,21 @@
 // first entry in enum will be MP_QSTRnull=0, which indicates invalid/no qstr
 enum {
     #ifndef NO_QSTR
-#define QDEF(id, hash, len, str) id,
+#define QDEF0(id, hash, len, str) id,
+#define QDEF1(id, hash, len, str)
     #include "genhdr/qstrdefs.generated.h"
-#undef QDEF
+#undef QDEF0
+#undef QDEF1
+    #endif
+    MP_QSTRnumber_of_special,
+    MP_QSTRstart_of_main = MP_QSTRnumber_of_special - 1,
+
+    #ifndef NO_QSTR
+#define QDEF0(id, hash, len, str)
+#define QDEF1(id, hash, len, str) id,
+    #include "genhdr/qstrdefs.generated.h"
+#undef QDEF0
+#undef QDEF1
     #endif
     MP_QSTRnumber_of, // no underscore so it can't clash with any of the above
 };
@@ -70,6 +82,7 @@ typedef struct _qstr_pool_t {
     size_t total_prev_len;
     size_t alloc;
     size_t len;
+    bool sorted;
     #if MICROPY_QSTR_BYTES_IN_HASH
     qstr_hash_t *hashes;
     #endif
diff --git a/tools/mpy-tool.py b/tools/mpy-tool.py
@@ -1406,7 +1406,7 @@ def freeze_mpy(base_qstrs, compiled_modules):
         if q is None or q.qstr_esc in base_qstrs or q.qstr_esc in new:
             continue
         new[q.qstr_esc] = (len(new), q.qstr_esc, q.str, bytes_cons(q.str, "utf8"))
-    new = sorted(new.values(), key=lambda x: x[0])
+    new = sorted(new.values(), key=lambda x: x[2])
 
     print('#include "py/mpconfig.h"')
     print('#include "py/objint.h"')
@@ -1489,6 +1489,7 @@ def freeze_mpy(base_qstrs, compiled_modules):
     print("    MP_QSTRnumber_of, // previous pool size")
     print("    %u, // allocated entries" % qstr_pool_alloc)
     print("    %u, // used entries" % len(new))
+    print("    true, // sorted")
     if config.MICROPY_QSTR_BYTES_IN_HASH:
         print("    (qstr_hash_t *)mp_qstr_frozen_const_hashes,")
     print("    (qstr_len_t *)mp_qstr_frozen_const_lengths,")
@@ -1781,7 +1782,9 @@ def main():
 
     # set config values for qstrs, and get the existing base set of qstrs
     if args.qstr_header:
-        qcfgs, base_qstrs = qstrutil.parse_input_headers([args.qstr_header])
+        special_qstrs = qstrutil.static_qstr_list_ident
+        qcfgs, extra_qstrs = qstrutil.parse_input_headers([args.qstr_header])
+        base_qstrs = set(special_qstrs) | set(extra_qstrs.keys())
         config.MICROPY_QSTR_BYTES_IN_LEN = int(qcfgs["BYTES_IN_LEN"])
         config.MICROPY_QSTR_BYTES_IN_HASH = int(qcfgs["BYTES_IN_HASH"])
     else:
