Consider element size when materializing NativeCharSequence

Author: fangerer

graalpython/com.oracle.graal.python.cext/src/capi.c

@@ -480,13 +480,16 @@ const char* PyTruffle_StringToCstr(void* o, int32_t strLen) {
     return str;
 }
 
+/* Use this function to decode a C char array to a Java string using the source file encoding. */
 void* PyTruffle_CstrToString(void* o) {
-    if (polyglot_fits_in_i64(o)) {
-        return polyglot_from_string((const char*)polyglot_as_i64(o), SRC_CS);
-    }
     return polyglot_from_string(o, SRC_CS);
 }
 
+/* Use this function to decode a C ASCII string to a Java string. */
+void* PyTruffle_AsciiToString(void* ptr) {
+    return polyglot_from_string(ptr, "ascii");
+}
+
 /* To be used from Java code only.
  * This function is used if a native class inherits from a managed class but uses the 'object.__new__'.
  * This function roughly corresponds to CPython's 'object_new'. */

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/NativeCAPISymbols.java

@@ -55,6 +55,7 @@ public abstract class NativeCAPISymbols {
     public static final String FUN_PY_NONE_HANDLE = "PyNoneHandle";
     public static final String FUN_WHCAR_SIZE = "PyTruffle_Wchar_Size";
     public static final String FUN_PY_TRUFFLE_CSTR_TO_STRING = "PyTruffle_CstrToString";
+    public static final String FUN_PY_TRUFFLE_ASCII_TO_STRING = "PyTruffle_AsciiToString";
     public static final String FUN_PY_FLOAT_AS_DOUBLE = "truffle_read_ob_fval";
     public static final String FUN_GET_OB_TYPE = "get_ob_type";
     public static final String FUN_GET_OB_REFCNT = "get_ob_refcnt";

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/str/StringNodes.java

@@ -50,6 +50,7 @@
 import com.oracle.graal.python.builtins.objects.cext.CExtNodes.ToSulongNode;
 import com.oracle.graal.python.builtins.objects.cext.NativeCAPISymbols;
 import com.oracle.graal.python.builtins.objects.cext.PythonNativeObject;
+import com.oracle.graal.python.builtins.objects.cext.common.CExtCommonNodes.UnicodeFromWcharNode;
 import com.oracle.graal.python.builtins.objects.common.SequenceNodes;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
 import com.oracle.graal.python.builtins.objects.ints.PInt;
@@ -111,10 +112,10 @@ static String doMaterializedNative(PString x) {
 
         @Specialization(guards = {"isNativeCharSequence(x)"}, replaces = "doMaterializedNative")
         static String doNative(PString x,
-                        @Cached PCallCapiFunction callCStringToStringNode) {
+                        @Cached PCallCapiFunction callCStringToStringNode,
+                        @Cached UnicodeFromWcharNode fromWcharNode) {
             // cast guaranteed by the guard
-            NativeCharSequence nativeCharSequence = (NativeCharSequence) x.getCharSequence();
-            String materialized = (String) callCStringToStringNode.call(NativeCAPISymbols.FUN_PY_TRUFFLE_CSTR_TO_STRING, nativeCharSequence.getPtr());
+            String materialized = materializeNativeCharSequence((NativeCharSequence) x.getCharSequence(), callCStringToStringNode, fromWcharNode);
             x.setCharSequence(materialized);
             return materialized;
         }
@@ -132,6 +133,29 @@ static String doMaterialized(PString x) {
             // cast guaranteed by the guard
             return (String) x.getCharSequence();
         }
+
+        public static String materializeNativeCharSequence(NativeCharSequence nativeCharSequence,
+                        PCallCapiFunction callCStringToStringNode,
+                        UnicodeFromWcharNode fromWcharNode) {
+            // cast guaranteed by the guard
+            String materialized;
+            if (nativeCharSequence.isAsciiOnly()) {
+                materialized = (String) callCStringToStringNode.call(NativeCAPISymbols.FUN_PY_TRUFFLE_ASCII_TO_STRING, nativeCharSequence.getPtr());
+            } else {
+                switch (nativeCharSequence.getElementSize()) {
+                    case 1:
+                        materialized = (String) callCStringToStringNode.call(NativeCAPISymbols.FUN_PY_TRUFFLE_CSTR_TO_STRING, nativeCharSequence.getPtr());
+                        break;
+                    case 2:
+                    case 4:
+                        materialized = fromWcharNode.execute(nativeCharSequence.getPtr(), nativeCharSequence.getElementSize());
+                        break;
+                    default:
+                        throw CompilerDirectives.shouldNotReachHere("illegal element size");
+                }
+            }
+            return materialized;
+        }
     }
 
     @GenerateUncached