From b42da1e56fcce91268bb8241ba2f23b80d39dd72 Mon Sep 17 00:00:00 2001
From: Daniel Kolesa <daniel@octaforge.org>
Date: Fri, 19 Mar 2021 01:00:11 +0100
Subject: [PATCH] add custom buffer allocation into string manager

---
 src/cs_util.cc | 51 +++++++++++++++++++++++++++++++++++++-------------
 src/cs_util.hh | 10 ++++++++++
 2 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/src/cs_util.cc b/src/cs_util.cc
index 1b1e969..3f17370 100644
--- a/src/cs_util.cc
+++ b/src/cs_util.cc
@@ -205,21 +205,11 @@ char const *cs_strman::add(ostd::string_range str) {
     }
     /* not present: allocate brand new data */
     auto ss = str.size();
-    auto mem = cstate->alloc(nullptr, 0, ss + sizeof(cs_strref_state) + 1);
-    if (!mem) {
-        throw cs_internal_error{"allocation failed"};
-    }
-    /* write length and refcount */
-    auto *sst = static_cast<cs_strref_state *>(mem);
-    sst->length = ss;
-    sst->refcount = 1;
-    /* write string data */
-    auto *strp = reinterpret_cast<char *>(sst + 1);
+    auto strp = alloc_buf(ss);
+    /* write string data, it's already pre-terminated */
     memcpy(strp, str.data(), ss);
-    /* terminated for best compatibility */
-    strp[ss] = '\0';
     /* store it */
-    counts.emplace(ostd::string_range{strp, strp + ss}, sst);
+    counts.emplace(ostd::string_range{strp, strp + ss}, get_ref_state(strp));
     return strp;
 }
 
@@ -229,6 +219,25 @@ char const *cs_strman::ref(char const *ptr) {
     return ptr;
 }
 
+char const *cs_strman::steal(char *ptr) {
+    auto *ss = get_ref_state(ptr);
+    auto sr = ostd::string_range{ptr, ptr + ss->length};
+    /* much like add(), but we already have memory */
+    auto it = counts.find(sr);
+    if (it != counts.end()) {
+        auto *st = it->second;
+        if (st) {
+            ++st->refcount;
+            /* the buffer is superfluous now */
+            cstate->alloc(ss, ss->length + sizeof(cs_strref_state) + 1, 0);
+            return reinterpret_cast<char const *>(st + 1);
+        }
+    }
+    ss->refcount = 1;
+    counts.emplace(sr, ss);
+    return ptr;
+}
+
 void cs_strman::unref(char const *ptr) {
     auto *ss = get_ref_state(ptr);
     if (!--ss->refcount) {
@@ -261,6 +270,22 @@ ostd::string_range cs_strman::get(char const *ptr) const {
     return ostd::string_range{ptr, ptr + ss->length};
 }
 
+char *cs_strman::alloc_buf(std::size_t len) const {
+    auto mem = cstate->alloc(nullptr, 0, len + sizeof(cs_strref_state) + 1);
+    if (!mem) {
+        throw cs_internal_error{"allocation failed"};
+    }
+    /* write length and initial refcount */
+    auto *sst = static_cast<cs_strref_state *>(mem);
+    sst->length = len;
+    sst->refcount = 1;
+    /* pre-terminate */
+    auto *strp = reinterpret_cast<char *>(sst + 1);
+    strp[len] = '\0';
+    /* now the user can fill it */
+    return strp;
+};
+
 namespace util {
     OSTD_EXPORT ostd::string_range parse_string(
         cs_state &cs, ostd::string_range str, size_t &nlines
diff --git a/src/cs_util.hh b/src/cs_util.hh
index e773834..ded5c18 100644
--- a/src/cs_util.hh
+++ b/src/cs_util.hh
@@ -86,6 +86,11 @@ struct cs_strman {
      */
     char const *ref(char const *ptr);
 
+    /* this will use the provided memory, assuming it is a fresh string that
+     * is yet to be added; the memory must be allocated with alloc_buf()
+     */
+    char const *steal(char *ptr);
+
     /* decrements the reference count and removes it from the system if
      * that reaches zero; likewise, only safe with pointers that are managed
      */
@@ -99,6 +104,11 @@ struct cs_strman {
     /* a quick helper to make a proper ostd string range out of a ptr */
     ostd::string_range get(char const *ptr) const;
 
+    /* this will allocate a buffer of the given length (plus one for
+     * terminating zero) so you can fill it; use steal() to write it
+     */
+    char *alloc_buf(std::size_t len) const;
+
     cs_shared_state *cstate;
     /* FIXME: use main allocator */
     std::unordered_map<ostd::string_range, cs_strref_state *> counts{};