1 files changed, 87 insertions, 13 deletions

diff --git a/src/lj_buf.h b/src/lj_buf.h
index 268fce5e..1fb70146 100644
--- a/src/lj_buf.h
+++ b/src/lj_buf.h
@@ -10,13 +10,53 @@
 #include "lj_gc.h"
 #include "lj_str.h"
 
-/* Resizable string buffers. SBuf struct definition in lj_obj.h. */
-#define sbufsz(sb)      ((MSize)((sb)->e - (sb)->b))
-#define sbuflen(sb)     ((MSize)((sb)->w - (sb)->b))
-#define sbufleft(sb)    ((MSize)((sb)->e - (sb)->w))
+/* Resizable string buffers. */
 
-#define sbufL(sb)       (mref((sb)->L, lua_State))
-#define setsbufL(sb, l) (setmref((sb)->L, (l)))
+/* The SBuf struct definition is in lj_obj.h:
+**   char *w;   Write pointer.
+**   char *e;   End pointer.
+**   char *b;   Base pointer.
+**   MRef L;    lua_State, used for buffer resizing. Extension bits in 3 LSB.
+*/
+
+/* Extended string buffer. */
+typedef struct SBufExt {
+  SBufHeader;
+  union {
+    GCRef cowref;       /* Copy-on-write object reference. */
+    MRef bsb;           /* Borrowed string buffer. */
+  };
+  char *r;              /* Read pointer. */
+  int depth;            /* Remaining recursion depth. */
+} SBufExt;
+
+#define sbufsz(sb)              ((MSize)((sb)->e - (sb)->b))
+#define sbuflen(sb)             ((MSize)((sb)->w - (sb)->b))
+#define sbufleft(sb)            ((MSize)((sb)->e - (sb)->w))
+#define sbufxlen(sbx)           ((MSize)((sbx)->w - (sbx)->r))
+#define sbufxslack(sbx)         ((MSize)((sbx)->r - (sbx)->b))
+
+#define SBUF_MASK_FLAG          (7)
+#define SBUF_MASK_L             (~(GCSize)SBUF_MASK_FLAG)
+#define SBUF_FLAG_EXT           1       /* Extended string buffer. */
+#define SBUF_FLAG_COW           2       /* Copy-on-write buffer. */
+#define SBUF_FLAG_BORROW        4       /* Borrowed string buffer. */
+
+#define sbufL(sb) \
+  ((lua_State *)(void *)(uintptr_t)(mrefu((sb)->L) & SBUF_MASK_L))
+#define setsbufL(sb, l)         (setmref((sb)->L, (l)))
+#define setsbufXL(sb, l, flag) \
+  (setmrefu((sb)->L, (GCSize)(uintptr_t)(void *)(l) + (flag)))
+#define setsbufXL_(sb, l) \
+  (setmrefu((sb)->L, (GCSize)(uintptr_t)(void *)(l) | (mrefu((sb)->L) & SBUF_MASK_FLAG)))
+
+#define sbufflag(sb)            (mrefu((sb)->L))
+#define sbufisext(sb)           (sbufflag((sb)) & SBUF_FLAG_EXT)
+#define sbufiscow(sb)           (sbufflag((sb)) & SBUF_FLAG_COW)
+#define sbufisborrow(sb)        (sbufflag((sb)) & SBUF_FLAG_BORROW)
+#define sbufX(sb) \
+  (lj_assertG_(G(sbufL(sb)), sbufisext(sb), "not an SBufExt"), (SBufExt *)(sb))
+#define setsbufflag(sb, flag)   (setmrefu((sb)->L, (flag)))
 
 /* Buffer management */
 LJ_FUNC char *LJ_FASTCALL lj_buf_need2(SBuf *sb, MSize sz);
@@ -45,6 +85,7 @@ static LJ_AINLINE SBuf *lj_buf_tmp_(lua_State *L)
 
 static LJ_AINLINE void lj_buf_free(global_State *g, SBuf *sb)
 {
+  lj_assertG(!sbufisext(sb), "bad free of SBufExt");
   lj_mem_free(g, sb->b, sbufsz(sb));
 }
 
@@ -62,6 +103,46 @@ static LJ_AINLINE char *lj_buf_more(SBuf *sb, MSize sz)
   return sb->w;
 }
 
+/* Extended buffer management */
+static LJ_AINLINE void lj_bufx_init(lua_State *L, SBufExt *sbx)
+{
+  memset(sbx, 0, sizeof(SBufExt));
+  setsbufXL(sbx, L, SBUF_FLAG_EXT);
+}
+
+static LJ_AINLINE void lj_bufx_init_borrow(lua_State *L, SBufExt *sbx, SBuf *sb)
+{
+  memset(sbx, 0, sizeof(SBufExt));
+  setsbufXL(sbx, L, SBUF_FLAG_EXT | SBUF_FLAG_BORROW);
+  setmref(sbx->bsb, sb);
+  sbx->r = sbx->w = sbx->b = sb->b;
+  sbx->e = sb->e;
+}
+
+static LJ_AINLINE void lj_bufx_init_cow(lua_State *L, SBufExt *sbx,
+                                        const char *p, MSize len)
+{
+  memset(sbx, 0, sizeof(SBufExt));
+  setsbufXL(sbx, L, SBUF_FLAG_EXT | SBUF_FLAG_COW);
+  sbx->r = sbx->b = (char *)p;
+  sbx->w = sbx->e = (char *)p + len;
+}
+
+static LJ_AINLINE void lj_bufx_reset(SBufExt *sbx)
+{
+  if (sbufiscow(sbx)) {
+    setmrefu(sbx->L, (mrefu(sbx->L) & ~(GCSize)SBUF_FLAG_COW));
+    setgcrefnull(sbx->cowref);
+    sbx->b = sbx->e = NULL;
+  }
+  sbx->r = sbx->w = sbx->b;
+}
+
+static LJ_AINLINE void lj_bufx_free(global_State *g, SBufExt *sbx)
+{
+  if (!sbufiscow(sbx)) lj_mem_free(g, sbx->b, sbufsz(sbx));
+}
+
 /* Low-level buffer put operations */
 LJ_FUNC SBuf *lj_buf_putmem(SBuf *sb, const void *q, MSize len);
 LJ_FUNC SBuf * LJ_FASTCALL lj_buf_putchar(SBuf *sb, int c);
@@ -97,11 +178,4 @@ static LJ_AINLINE GCstr *lj_buf_str(lua_State *L, SBuf *sb)
   return lj_str_new(L, sb->b, sbuflen(sb));
 }
 
-/* Interim user-accessible string buffer. */
-typedef struct StrBuf {
-  SBuf *sb;             /* Pointer to system buffer. */
-  char *r;              /* String buffer read pointer. */
-  int depth;            /* Remaining recursion depth. */
-} StrBuf;
-
 #endif