strict reference counting for bytecode references

src/cs_val.cc

@@ -71,10 +71,7 @@ static inline void csv_cleanup(cs_value_type tv, T &stor) {
             reinterpret_cast<cs_strref *>(&stor)->~cs_strref();
             break;
         case cs_value_type::CODE: {
-            uint32_t *bcode = csv_get<uint32_t *>(stor);
-            if (bcode[-1] == CS_CODE_START) {
-                delete[] &bcode[-1];
-            }
+            bcode_unref(csv_get<uint32_t *>(stor));
             break;
         }
         default:
@@ -119,7 +116,7 @@ cs_value &cs_value::operator=(cs_value const &v) {
             };
             break;
         case cs_value_type::CODE:
-            set_code(cs_copy_code(v.get_code()));
+            set_code(v.get_code());
             break;
         default:
             break;
@@ -169,6 +166,7 @@ void cs_value::set_none() {
 void cs_value::set_code(cs_bcode *val) {
     csv_cleanup(p_type, p_stor);
     p_type = cs_value_type::CODE;
+    bcode_ref(reinterpret_cast<uint32_t *>(val));
     csv_get<cs_bcode *>(p_stor) = val;
 }
 

src/cs_vm.cc

@@ -135,7 +135,18 @@ cs_stack_state cs_error::save_stack(cs_state &cs) {
     return cs_stack_state(cs, ret, total > dalias->get_value());
 }
 
-static void bcode_ref(uint32_t *code) {
+static inline void bcode_incr(uint32_t *bc) {
+    *bc += 0x100;
+}
+
+static inline void bcode_decr(uint32_t *bc) {
+    *bc -= 0x100;
+    if (std::int32_t(*bc) < 0x100) {
+        delete[] bc;
+    }
+}
+
+void bcode_ref(uint32_t *code) {
     if (!code) {
         return;
     }
@@ -154,7 +165,7 @@ static void bcode_ref(uint32_t *code) {
     }
 }
 
-static void bcode_unref(uint32_t *code) {
+void bcode_unref(uint32_t *code) {
     if (!code) {
         return;
     }
@@ -198,6 +209,28 @@ cs_bcode_ref &cs_bcode_ref::operator=(cs_bcode_ref &&v) {
     return *this;
 }
 
+void cs_alias_impl::clean_code() {
+    uint32_t *bcode = reinterpret_cast<uint32_t *>(p_acode);
+    if (bcode) {
+        bcode_decr(bcode);
+        p_acode = nullptr;
+    }
+}
+
+cs_bcode *cs_alias_impl::compile_code(cs_state &cs) {
+    if (!p_acode) {
+        cs_gen_state gs(cs);
+        gs.code.reserve(64);
+        gs.gen_main(get_value().get_str());
+        /* i wish i could steal the memory somehow */
+        uint32_t *code = new uint32_t[gs.code.size()];
+        memcpy(code, gs.code.data(), gs.code.size() * sizeof(uint32_t));
+        bcode_incr(code);
+        p_acode = reinterpret_cast<cs_bcode *>(code);
+    }
+    return p_acode;
+}
+
 static inline uint32_t *forcecode(cs_state &cs, cs_value &v) {
     uint32_t *code = reinterpret_cast<uint32_t *>(v.get_code());
     if (!code) {
@@ -254,52 +287,6 @@ static inline void force_arg(cs_value &v, int type) {
     }
 }
 
-static uint32_t *skipcode(uint32_t *code) {
-    int depth = 0;
-    for (;;) {
-        uint32_t op = *code++;
-        switch (op & 0xFF) {
-            case CS_CODE_VAL | CS_RET_STRING: {
-                uint32_t len = op >> 8;
-                code += len / sizeof(uint32_t) + 1;
-                continue;
-            }
-            case CS_CODE_BLOCK:
-            case CS_CODE_JUMP:
-            case CS_CODE_JUMP_B | CS_CODE_FLAG_TRUE:
-            case CS_CODE_JUMP_B | CS_CODE_FLAG_FALSE:
-            case CS_CODE_JUMP_RESULT | CS_CODE_FLAG_TRUE:
-            case CS_CODE_JUMP_RESULT | CS_CODE_FLAG_FALSE: {
-                uint32_t len = op >> 8;
-                code += len;
-                continue;
-            }
-            case CS_CODE_ENTER:
-            case CS_CODE_ENTER_RESULT:
-                ++depth;
-                continue;
-            case CS_CODE_EXIT | CS_RET_NULL:
-            case CS_CODE_EXIT | CS_RET_STRING:
-            case CS_CODE_EXIT | CS_RET_INT:
-            case CS_CODE_EXIT | CS_RET_FLOAT:
-                if (depth <= 0) {
-                    return code;
-                }
-                --depth;
-                continue;
-        }
-    }
-}
-
-cs_bcode *cs_copy_code(cs_bcode *c) {
-    uint32_t *bcode = reinterpret_cast<uint32_t *>(c);
-    uint32_t *end = skipcode(bcode);
-    uint32_t *dst = new uint32_t[end - bcode + 1];
-    *dst++ = CS_CODE_START;
-    memcpy(dst, bcode, (end - bcode) * sizeof(uint32_t));
-    return reinterpret_cast<cs_bcode *>(dst);
-}
-
 static inline void callcommand(
     cs_state &cs, cs_command_impl *id, cs_value *args, cs_value &res,
     int numargs, bool lookup = false
@@ -1574,10 +1561,9 @@ static void cs_run(
     gs.done();
     uint32_t *cbuf = new uint32_t[gs.code.size()];
     memcpy(cbuf, gs.code.data(), gs.code.size() * sizeof(uint32_t));
+    bcode_ref(cbuf);
     runcode(cs, cbuf + 1, ret);
-    if (int(cbuf[0]) < 0x100) {
-        delete[] cbuf;
-    }
+    bcode_unref(cbuf);
 }
 
 void cs_state::run(std::string_view code, cs_value &ret) {

src/cs_vm.hh

@@ -421,38 +421,8 @@ struct cs_gen_state {
 
 bool cs_check_num(std::string_view s);
 
-static inline void bcode_incr(uint32_t *bc) {
-    *bc += 0x100;
-}
-
-static inline void bcode_decr(uint32_t *bc) {
-    *bc -= 0x100;
-    if (std::int32_t(*bc) < 0x100) {
-        delete[] bc;
-    }
-}
-
-inline void cs_alias_impl::clean_code() {
-    uint32_t *bcode = reinterpret_cast<uint32_t *>(p_acode);
-    if (bcode) {
-        bcode_decr(bcode);
-        p_acode = nullptr;
-    }
-}
-
-inline cs_bcode *cs_alias_impl::compile_code(cs_state &cs) {
-    if (!p_acode) {
-        cs_gen_state gs(cs);
-        gs.code.reserve(64);
-        gs.gen_main(get_value().get_str());
-        /* i wish i could steal the memory somehow */
-        uint32_t *code = new uint32_t[gs.code.size()];
-        memcpy(code, gs.code.data(), gs.code.size() * sizeof(uint32_t));
-        bcode_incr(code);
-        p_acode = reinterpret_cast<cs_bcode *>(code);
-    }
-    return p_acode;
-}
+void bcode_ref(uint32_t *code);
+void bcode_unref(uint32_t *code);
 
 template<typename F>
 static void cs_do_args(cs_state &cs, F body) {
@@ -492,8 +462,6 @@ static void cs_do_args(cs_state &cs, F body) {
     });
 }
 
-cs_bcode *cs_copy_code(cs_bcode *c);
-
 } /* namespace cscript */
 
 #endif /* LIBCUBESCRIPT_CS_VM_HH */