Avoid the overhead of the conversion in UTF-8 and ASCII locales

Add _Py_LocaleNeedsWcharConversion(), which memoizes the check of the
locale encoding, and test it before scanning the string.  Also keep
using wcscoll() and wcsxfrm() when no conversion is needed, so that
locale.strxfrm() returns the same keys as before.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: serhiy-storchaka

Include/internal/pycore_fileutils.h

@@ -243,6 +243,7 @@ extern void _Py_ResetForceASCII(void);
 // Convert between Unicode code points and the native wchar_t form for
 // C library functions which use the latter (such as the curses library).
 // Export for '_curses'.
+PyAPI_FUNC(int) _Py_LocaleNeedsWcharConversion(void);
 PyAPI_FUNC(int) _Py_UnicodeToLocaleWchar_InPlace(wchar_t *str, Py_ssize_t size);
 PyAPI_FUNC(void) _Py_LocaleWcharToUnicode_InPlace(wchar_t *str, Py_ssize_t size);
 #endif

Include/internal/pycore_runtime_structs.h

@@ -52,6 +52,12 @@ struct pyhash_runtime_state {
 
 struct _fileutils_state {
     int force_ascii;
+    /* Memoized result of _Py_LocaleNeedsWcharConversion() with the locale
+       encoding it was computed for (see Python/fileutils.c). */
+    struct _Py_wchar_conversion_state {
+        char codeset[32];
+        int needed;
+    } wchar;
 };
 
 #include "pycore_interpframe_structs.h" // _PyInterpreterFrame

Modules/_localemodule.c

@@ -407,6 +407,7 @@ _locale_strcoll_impl(PyObject *module, PyObject *os1, PyObject *os2)
 /*[clinic end generated code: output=82ddc6d62c76d618 input=693cd02bcbf38dd8]*/
 {
 #ifdef _Py_NON_UNICODE_WCHAR_T
+  if (_Py_LocaleNeedsWcharConversion()) {
     /* wchar_t is not Unicode in non-UTF-8 locales and the collation
        tables are keyed by the raw locale values, so wcscoll() cannot
        be used with Unicode wchar_t.  Encode to the locale encoding
@@ -424,7 +425,9 @@ _locale_strcoll_impl(PyObject *module, PyObject *os1, PyObject *os2)
     }
     Py_DECREF(b1);
     return result;
-#else
+  }
+#endif
+    {
     PyObject *result = NULL;
     wchar_t *ws1 = NULL, *ws2 = NULL;
 
@@ -442,7 +445,7 @@ _locale_strcoll_impl(PyObject *module, PyObject *os1, PyObject *os2)
     if (ws1) PyMem_Free(ws1);
     if (ws2) PyMem_Free(ws2);
     return result;
-#endif
+    }
 }
 #endif
 
@@ -462,6 +465,7 @@ _locale_strxfrm_impl(PyObject *module, PyObject *str)
 /*[clinic end generated code: output=3081866ebffc01af input=1378bbe6a88b4780]*/
 {
 #ifdef _Py_NON_UNICODE_WCHAR_T
+  if (_Py_LocaleNeedsWcharConversion()) {
     /* wchar_t is not Unicode in non-UTF-8 locales and the collation
        tables are keyed by the raw locale values, so wcsxfrm() cannot
        be used with Unicode wchar_t.  Encode to the locale encoding
@@ -498,7 +502,9 @@ _locale_strxfrm_impl(PyObject *module, PyObject *str)
 error:
     Py_DECREF(bytes);
     return result;
-#else
+  }
+#endif
+    {
     Py_ssize_t n1;
     wchar_t *s = NULL, *buf = NULL;
     wchar_t dummy[1];
@@ -586,7 +592,7 @@ _locale_strxfrm_impl(PyObject *module, PyObject *str)
     PyMem_Free(buf);
     PyMem_Free(s);
     return result;
-#endif /* _Py_NON_UNICODE_WCHAR_T */
+    }
 }
 #endif
 

Python/fileutils.c

@@ -470,15 +470,37 @@ decode_ascii(const char *arg, wchar_t **wstr, size_t *wlen,
 /* Fall back to mbstowcs()/wcstombs(). */
 #  define _Py_ICONV_FALLBACK (-9)
 
+/* Return non-zero if wchar_t is not Unicode in the current locale, i.e. if
+   the locale encoding is neither UTF-8 nor ASCII.  The result is memoized:
+   this is called before every conversion, and in the common case (a UTF-8
+   locale) it is the only overhead. */
+int
+_Py_LocaleNeedsWcharConversion(void)
+{
+    const char *codeset = nl_langinfo(CODESET);
+    if (codeset == NULL) {
+        return 0;
+    }
+    struct _Py_wchar_conversion_state *state = &_PyRuntime.fileutils.wchar;
+    if (strcmp(codeset, state->codeset) == 0) {
+        return state->needed;
+    }
+    int needed = (*codeset != '\0'
+                  && strcmp(codeset, "UTF-8") != 0
+                  && strcmp(codeset, "US-ASCII") != 0
+                  && strcmp(codeset, "646") != 0);
+    if (strlen(codeset) < sizeof(state->codeset)) {
+        memcpy(state->codeset, codeset, strlen(codeset) + 1);
+        state->needed = needed;
+    }
+    return needed;
+}
+
 static iconv_t
 locale_iconv_open(int decode)
 {
     const char *codeset = nl_langinfo(CODESET);
-    if (!codeset || !*codeset
-        || strcmp(codeset, "UTF-8") == 0
-        || strcmp(codeset, "US-ASCII") == 0
-        || strcmp(codeset, "646") == 0)
-    {
+    if (!_Py_LocaleNeedsWcharConversion()) {
         /* mbstowcs() and wcstombs() are correct for these encodings */
         return (iconv_t)-1;
     }
@@ -490,6 +512,10 @@ static int
 decode_locale_iconv(const char *arg, wchar_t **wstr, size_t *wlen,
                     const char **reason, _Py_error_handler errors)
 {
+    if (!_Py_LocaleNeedsWcharConversion()) {
+        return _Py_ICONV_FALLBACK;
+    }
+
     int surrogateescape;
     if (get_surrogateescape(errors, &surrogateescape) < 0) {
         return -3;
@@ -569,6 +595,10 @@ encode_locale_iconv(const wchar_t *text, char **str,
                     size_t *error_pos, const char **reason,
                     int raw_malloc, _Py_error_handler errors)
 {
+    if (!_Py_LocaleNeedsWcharConversion()) {
+        return _Py_ICONV_FALLBACK;
+    }
+
     int surrogateescape;
     if (get_surrogateescape(errors, &surrogateescape) < 0) {
         return -3;
@@ -687,6 +717,10 @@ encode_locale_iconv(const wchar_t *text, char **str,
 int
 _Py_UnicodeToLocaleWchar_InPlace(wchar_t *str, Py_ssize_t size)
 {
+    if (!_Py_LocaleNeedsWcharConversion()) {
+        return 0;
+    }
+
     iconv_t cd = (iconv_t)-1;
     int res = 0;
     for (Py_ssize_t i = 0; i < size; i++) {
@@ -763,6 +797,10 @@ _Py_UnicodeToLocaleWchar_InPlace(wchar_t *str, Py_ssize_t size)
 void
 _Py_LocaleWcharToUnicode_InPlace(wchar_t *str, Py_ssize_t size)
 {
+    if (!_Py_LocaleNeedsWcharConversion()) {
+        return;
+    }
+
     iconv_t cd = (iconv_t)-1;
     for (Py_ssize_t i = 0; i < size; i++) {
         wchar_t wc = str[i];