From ffbc7e9adea4739466dbc629979172270ac53bf4 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Thu, 23 Jan 2025 20:02:55 +0800
Subject: [PATCH] intern ascii literals

---
 include/pocketpy/interpreter/vm.h |  2 ++
 src/interpreter/vm.c              | 18 ++++++++++++++----
 src/public/modules.c              |  3 +--
 src/public/py_str.c               | 11 +++++++++++
 tests/04_str.py                   |  2 ++
 5 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/include/pocketpy/interpreter/vm.h b/include/pocketpy/interpreter/vm.h
index ddd2da90..c97a5aa8 100644
--- a/include/pocketpy/interpreter/vm.h
+++ b/include/pocketpy/interpreter/vm.h
@@ -27,6 +27,8 @@ typedef struct VM {
 
     py_Callbacks callbacks;
 
+    py_TValue ascii_literals[128+1];
+
     py_TValue last_retval;
     py_TValue curr_exception;
     volatile bool is_signal_interrupted;
diff --git a/src/interpreter/vm.c b/src/interpreter/vm.c
index 645089bd..d0da35ed 100644
--- a/src/interpreter/vm.c
+++ b/src/interpreter/vm.c
@@ -83,6 +83,12 @@ void VM__ctor(VM* self) {
     ValueStack__ctor(&self->stack);
 
     /* Init Builtin Types */
+    for(int i = 0; i < 128; i++) {
+        char* p = py_newstrn(&self->ascii_literals[i], 1);
+        *p = i;
+    }
+    py_newstrn(&self->ascii_literals[128], 0);
+
     // 0: unused
     void* placeholder = TypeList__emplace(&self->types);
     memset(placeholder, 0, sizeof(py_TypeInfo));
@@ -153,7 +159,7 @@ void VM__ctor(VM* self) {
 
     validate(tp_StopIteration, pk_StopIteration__register());
     py_setdict(&self->builtins, py_name("StopIteration"), py_tpobject(tp_StopIteration));
-    
+
     INJECT_BUILTIN_EXC(SyntaxError, tp_Exception);
     INJECT_BUILTIN_EXC(StackOverflowError, tp_Exception);
     INJECT_BUILTIN_EXC(OSError, tp_Exception);
@@ -224,8 +230,8 @@ void VM__ctor(VM* self) {
     pk__add_module_importlib();
 
     pk__add_module_conio();
-    pk__add_module_lz4();       // optional
-    pk__add_module_libhv();     // optional
+    pk__add_module_lz4();    // optional
+    pk__add_module_libhv();  // optional
     pk__add_module_pkpy();
 
     // add python builtins
@@ -509,7 +515,7 @@ FrameResult VM__vectorcall(VM* self, uint16_t argc, uint16_t kwargc, bool opcall
                 memcpy(argv, self->__vectorcall_buffer, co->nlocals * sizeof(py_TValue));
                 Frame* frame = Frame__new(co, &fn->module, p0, argv, true);
                 pk_newgenerator(py_retval(), frame, p0, self->stack.sp);
-                self->stack.sp = p0;    // reset the stack
+                self->stack.sp = p0;  // reset the stack
                 return RES_RETURN;
             }
             default: c11__unreachable();
@@ -640,6 +646,10 @@ void ManagedHeap__mark(ManagedHeap* self) {
     for(py_TValue* p = vm->stack.begin; p != vm->stack.end; p++) {
         pk__mark_value(p);
     }
+    // mark ascii literals
+    for(int i = 0; i < c11__count_array(vm->ascii_literals); i++) {
+        pk__mark_value(&vm->ascii_literals[i]);
+    }
     // mark modules
     ModuleDict__apply_mark(&vm->modules, mark_object);
     // mark types
diff --git a/src/public/modules.c b/src/public/modules.c
index 69228eb9..06bac088 100644
--- a/src/public/modules.c
+++ b/src/public/modules.c
@@ -456,8 +456,7 @@ static bool builtins_chr(int argc, py_Ref argv) {
     PY_CHECK_ARG_TYPE(0, tp_int);
     py_i64 val = py_toint(py_arg(0));
     if(val < 0 || val > 128) { return ValueError("chr() arg not in range(128)"); }
-    char* data = py_newstrn(py_retval(), 1);
-    data[0] = (char)val;
+    py_assign(py_retval(), &pk_current_vm->ascii_literals[val]);
     return true;
 }
 
diff --git a/src/public/py_str.c b/src/public/py_str.c
index 4b6bd43d..c5be746f 100644
--- a/src/public/py_str.c
+++ b/src/public/py_str.c
@@ -21,6 +21,17 @@ char* py_newstrn(py_Ref out, int size) {
 }
 
 void py_newstrv(py_OutRef out, c11_sv sv) {
+    if(sv.size == 0) {
+        *out = pk_current_vm->ascii_literals[128];
+        return;
+    }
+    if(sv.size == 1) {
+        int c = sv.data[0];
+        if(c >= 0 && c < 128) {
+            *out = pk_current_vm->ascii_literals[c];
+            return;
+        }
+    }
     char* data = py_newstrn(out, sv.size);
     memcpy(data, sv.data, sv.size);
 }
diff --git a/tests/04_str.py b/tests/04_str.py
index fd4af45a..f6437b6a 100644
--- a/tests/04_str.py
+++ b/tests/04_str.py
@@ -190,6 +190,8 @@ assert (1 != '1') is True
 assert (1 == '1') is False
 assert 1 == 1.0
 
+assert chr(97) is 'a'
+
 exit()
 
 # test format()