From 13fe4849e9925b2eec851cb03c64938f5a3b8ae0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C5=8Dshin?= <git@wholezero.org>
Date: Thu, 6 Jun 2024 21:16:14 -0700
Subject: [PATCH] Fix some memory issues with ctl::string

There were a few errors in how capacity and memory was being handled for
small strings. The capacity errors meant that small strings would become
big strings too soon, and the memory error introduced undefined behavior
that was caught by CheckMemoryLeaks in our test file but only sometimes.

The crucial change is in reserve: we only copy n bytes into p2, and then
we manually set the null terminator instead of expecting it to have been
there already. (E.g. it might not be there for an empty small string.)

We also fix one other doozy in append when we were exactly at the small-
to-big string boundary: we set the last byte (i.e., the remainder field)
to 0, then decremented it, giving us size_t max. Whoops. We boneheadedly
fix this by setting the 0 byte after we've fixed up the remainder, so it
is at worst a no-op.

Otherwise, capacity now works the same for small strings as it does with
big strings: it's the amount of space available including the null byte.

We test all of this with a new test that only gets included if our class
under test is not std::string (presumably meaning it's ctl::string.) The
test manually verifies that the small string optimization behaves how we
expect.

Since this test checks against std::string, we go ahead and include that
other header from the STL.
---
 ctl/string.cc           |  7 ++++---
 ctl/string.h            |  2 +-
 test/ctl/string_test.cc | 18 +++++++++++++++++-
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/ctl/string.cc b/ctl/string.cc
index ff864c306..2a2b22eb1 100644
--- a/ctl/string.cc
+++ b/ctl/string.cc
@@ -99,7 +99,8 @@ string::reserve(size_t c2) noexcept
     if (!isbig()) {
         if (!(p2 = (char*)malloc(c2)))
             __builtin_trap();
-        memcpy(p2, data(), size() + 1);
+        memcpy(p2, data(), size());
+        p2[size()] = 0;
     } else {
         if (!(p2 = (char*)realloc(big()->p, c2)))
             __builtin_trap();
@@ -134,18 +135,18 @@ string::append(char ch) noexcept
     if (ckd_add(&n2, size(), 2))
         __builtin_trap();
     if (n2 > capacity()) {
-        size_t c2 = capacity() + 2;
+        size_t c2 = capacity();
         if (ckd_add(&c2, c2, c2 >> 1))
             __builtin_trap();
         reserve(c2);
     }
     data()[size()] = ch;
-    data()[size() + 1] = 0;
     if (isbig()) {
         ++big()->n;
     } else {
         --small()->rem;
     }
+    data()[size()] = 0;
 }
 
 void
diff --git a/ctl/string.h b/ctl/string.h
index 858b7a224..c830fdf12 100644
--- a/ctl/string.h
+++ b/ctl/string.h
@@ -151,7 +151,7 @@ class string
         if (isbig() && big()->c <= __::sso_max)
             __builtin_trap();
 #endif
-        return isbig() ? __::big_mask & big()->c : __::sso_max;
+        return isbig() ? __::big_mask & big()->c : __::string_size;
     }
 
     iterator begin() noexcept
diff --git a/test/ctl/string_test.cc b/test/ctl/string_test.cc
index 1a35cfda3..be664b22b 100644
--- a/test/ctl/string_test.cc
+++ b/test/ctl/string_test.cc
@@ -23,7 +23,7 @@
 #include "libc/runtime/runtime.h"
 #include "libc/str/str.h"
 
-// #include <string>
+#include <string>
 // #define ctl std
 
 int
@@ -353,6 +353,22 @@ main()
             return 78;
     }
 
+    if constexpr (!std::is_same_v<ctl::string, std::string>) {
+        // tests the small-string optimization on ctl::string
+        ctl::string s;
+        char *d = s.data();
+        for (int i = 0; i < 23; ++i) {
+            s.append('a');
+            if (s.data() != d) {
+                return 79 + i;
+            }
+        }
+        s.append('a');
+        if (s.data() == d) {
+            return 103;
+        }
+    }
+
     CheckForMemoryLeaks();
     return 0;
 }