Allow static upb_strings.

This can allow strings to reference static data, and
reduced the memory footprint of test_def by about 10%
(3k).

6 files changed, 103 insertions, 15 deletions

diff --git a/core/upb_def.c b/core/upb_def.c
index b9402c5..c0d72db 100644
--- a/core/upb_def.c
+++ b/core/upb_def.c
@@ -1018,12 +1018,10 @@ static upb_src *upb_baredecoder_src(upb_baredecoder *d)
 void upb_symtab_add_descriptorproto(upb_symtab *symtab)
 {
   // TODO: allow upb_strings to be static or on the stack.
-  upb_string *descriptor = upb_strduplen(descriptor_pb, descriptor_pb_len);
-  upb_baredecoder *decoder = upb_baredecoder_new(descriptor);
+  upb_baredecoder *decoder = upb_baredecoder_new(&descriptor_str);
   upb_status status = UPB_STATUS_INIT;
   upb_symtab_addfds(symtab, upb_baredecoder_src(decoder), &status);
   upb_baredecoder_free(decoder);
-  upb_string_unref(descriptor);
 
   if(!upb_ok(&status)) {
     // upb itself is corrupt.
diff --git a/core/upb_string.c b/core/upb_string.c
index 3563c9e..ca3c669 100644
--- a/core/upb_string.c
+++ b/core/upb_string.c
@@ -61,12 +61,12 @@ void _upb_string_free(upb_string *str) {
 }
 
 upb_string *upb_string_tryrecycle(upb_string *str) {
-  if(str == NULL || upb_atomic_read(&str->refcount) > 1) {
-    return upb_string_new();
-  } else {
+  if(str && upb_atomic_read(&str->refcount) == 1) {
     str->ptr = NULL;
     upb_string_release(str);
     return str;
+  } else {
+    return upb_string_new();
   }
 }
 
@@ -125,3 +125,9 @@ upb_string *upb_string_asprintf(const char *format, ...) {
   va_end(args);
   return str;
 }
+
+upb_string *upb_strdup(upb_string *s) {
+  upb_string *str = upb_string_new();
+  upb_strcpy(str, s);
+  return str;
+}
diff --git a/core/upb_string.h b/core/upb_string.h
index 5cc0eaf..65ba404 100644
--- a/core/upb_string.h
+++ b/core/upb_string.h
@@ -63,6 +63,17 @@ struct _upb_string {
   struct _upb_string *src;
 };
 
+// Internal-only initializer for upb_string instances.
+#ifdef UPB_HAVE_MSIZE
+#define _UPB_STRING_INIT(str, len, refcount) {(char*)str, NULL, len, {refcount}, NULL}
+#else
+#define _UPB_STRING_INIT(str, len, refcount) {(char*)str, NULL, len, 0, {refcount}, NULL}
+#endif
+
+// Special pseudo-refcounts for static/stack-allocated strings, respectively.
+#define _UPB_STRING_REFCOUNT_STATIC -1
+#define _UPB_STRING_REFCOUNT_STACK -2
+
 // Returns a newly-created, empty, non-finalized string.  When the string is no
 // longer needed, it should be unref'd, never freed directly.
 upb_string *upb_string_new();
@@ -72,15 +83,21 @@ void _upb_string_free(upb_string *str);
 // Releases a ref on the given string, which may free the memory.  "str"
 // can be NULL, in which case this is a no-op.
 INLINE void upb_string_unref(upb_string *str) {
-  if (str && upb_atomic_unref(&str->refcount)) _upb_string_free(str);
+  if (str && upb_atomic_read(&str->refcount) > 0 &&
+      upb_atomic_unref(&str->refcount)) {
+    _upb_string_free(str);
+  }
 }
 
+upb_string *upb_strdup(upb_string *s);  // Forward-declare.
+
 // Returns a string with the same contents as "str".  The caller owns a ref on
 // the returned string, which may or may not be the same object as "str.
 INLINE upb_string *upb_string_getref(upb_string *str) {
-  // If/when we support stack-allocated strings, this will have to allocate
-  // a new string if the given string is on the stack.
-  upb_atomic_ref(&str->refcount);
+  int refcount = upb_atomic_read(&str->refcount);
+  if (refcount == _UPB_STRING_REFCOUNT_STACK) return upb_strdup(str);
+  // We don't ref the special <0 refcount for static strings.
+  if (refcount > 0) upb_atomic_ref(&str->refcount);
   return str;
 }
 
@@ -151,6 +168,38 @@ void upb_string_substr(upb_string *str, upb_string *target_str,
 #define UPB_STRARG(str) upb_string_len(str), upb_string_getrobuf(str)
 #define UPB_STRFMT "%.*s"
 
+// Macros for constructing upb_string objects statically or on the stack.  These
+// can be used like:
+//
+// upb_string static_str = UPB_STATIC_STRING("Foo");
+//
+// int main() {
+//   upb_string stack_str = UPB_STACK_STRING("Foo");
+//   // Now:
+//   //   upb_streql(&static_str, &stack_str) == true
+//   //   upb_streql(&static_str, UPB_STRLIT("Foo")) == true
+// }
+//
+// You can also use UPB_STACK_STRING or UPB_STATIC_STRING with character arrays,
+// but you must not change the underlying data once you've passed the string on:
+//
+// void foo() {
+//   char data[] = "ABC123";
+//   upb_string stack_str = UPB_STACK_STR(data);
+//   bar(&stack_str);
+//   data[0] = "B";  // NOT ALLOWED!!
+// }
+//
+// TODO: should the stack business just be like attach/detach?  The latter seems
+// more flexible, though it does require a stack allocation.  Maybe put this off
+// until there is a clear use case.
+#define UPB_STATIC_STRING(str) \
+    _UPB_STRING_INIT(str, sizeof(str)-1, _UPB_STRING_REFCOUNT_STATIC)
+#define UPB_STATIC_STRING_LEN(str, len) \
+    _UPB_STRING_INIT(str, len, _UPB_STRING_REFCOUNT_STATIC)
+#define UPB_STACK_STRING(str) _UPB_STRING_INIT(str, _UPB_STRING_REFCOUNT_STACK)
+#define UPB_STRLIT(str) &(upb_string)UPB_STATIC_STRING(str)
+
 /* upb_string library functions ***********************************************/
 
 // Named like their <string.h> counterparts, these are all safe against buffer
diff --git a/descriptor/descriptor.c b/descriptor/descriptor.c
index cd50a16..ee6b25b 100644
--- a/descriptor/descriptor.c
+++ b/descriptor/descriptor.c
@@ -1,4 +1,6 @@
-unsigned char descriptor_pb[] = {
+#include "descriptor.h"
+
+static unsigned char descriptor_pb[] = {
   0x0a, 0x9b, 0x1b, 0x0a, 0x1b, 0x64, 0x65, 0x73, 0x63, 0x72, 0x69, 0x70,
   0x74, 0x6f, 0x72, 0x2f, 0x64, 0x65, 0x73, 0x63, 0x72, 0x69, 0x70, 0x74,
   0x6f, 0x72, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x0f, 0x67, 0x6f,
@@ -291,4 +293,5 @@ unsigned char descriptor_pb[] = {
   0x44, 0x65, 0x73, 0x63, 0x72, 0x69, 0x70, 0x74, 0x6f, 0x72, 0x50, 0x72,
   0x6f, 0x74, 0x6f, 0x73, 0x48, 0x01
 };
-unsigned int descriptor_pb_len = 3486;
+static const unsigned int descriptor_pb_len = 3486;
+upb_string descriptor_str = UPB_STATIC_STRING(descriptor_pb);
diff --git a/descriptor/descriptor.h b/descriptor/descriptor.h
index b598a9a..f6d3ca3 100644
--- a/descriptor/descriptor.h
+++ b/descriptor/descriptor.h
@@ -11,12 +11,13 @@
 #ifndef UPB_DESCRIPTOR_H_
 #define UPB_DESCRIPTOR_H_
 
+#include "upb_string.h"
+
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-extern unsigned char descriptor_pb[];
-extern unsigned int descriptor_pb_len;
+extern upb_string descriptor_str;
 
 #ifdef __cplusplus
 }  /* extern "C" */
diff --git a/tests/test_string.c b/tests/test_string.c
index 46f35b9..7c9ed02 100644
--- a/tests/test_string.c
+++ b/tests/test_string.c
@@ -3,8 +3,33 @@
 #include "upb_string.h"
 
 char static_str[] = "Static string.";
+upb_string static_upbstr = UPB_STATIC_STRING(static_str);
 
-int main() {
+static void test_static() {
+  // Static string is initialized appropriately.
+  assert(upb_streql(&static_upbstr, UPB_STRLIT("Static string.")));
+
+  // Taking a ref on a static string returns the same string, and repeated
+  // refs don't get the string in a confused state.
+  assert(upb_string_getref(&static_upbstr) == &static_upbstr);
+  assert(upb_string_getref(&static_upbstr) == &static_upbstr);
+  assert(upb_string_getref(&static_upbstr) == &static_upbstr);
+
+  // Unreffing a static string does nothing (is not harmful).
+  upb_string_unref(&static_upbstr);
+  upb_string_unref(&static_upbstr);
+  upb_string_unref(&static_upbstr);
+  upb_string_unref(&static_upbstr);
+  upb_string_unref(&static_upbstr);
+
+  // Recycling a static string returns a new string (that can be modified).
+  upb_string *str = upb_string_tryrecycle(&static_upbstr);
+  assert(str != &static_upbstr);
+
+  upb_string_unref(str);
+}
+
+static void test_dynamic() {
   upb_string *str = upb_string_new();
   assert(str != NULL);
   upb_string_unref(str);
@@ -29,6 +54,7 @@ int main() {
   const char *robuf2 = upb_string_getrobuf(str);
   assert(robuf2 == robuf);
   assert(upb_streqlc(str, "XX"));
+  assert(upb_streql(str, UPB_STRLIT("XX")));
 
   // Make string alias part of another string.
   str2 = upb_strdupc("WXYZ");
@@ -79,3 +105,8 @@ int main() {
   // Unref of NULL is harmless.
   upb_string_unref(NULL);
 }
+
+int main() {
+  test_static();
+  test_dynamic();
+}