1 files changed, 67 insertions, 15 deletions

diff --git a/pw.c b/pw.c
index 5ef006d..e9026b3 100644
--- a/pw.c
+++ b/pw.c
@@ -1,3 +1,4 @@
+#include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <ctype.h>
@@ -17,6 +18,14 @@ enum status_flags {
   stat_susp = 4
 };
 
+typedef struct dstr {
+  int refs;
+  size_t len;
+  char str[];
+} dstr;
+
+#define dstr_of(str) ((dstr *) ((str) - sizeof (dstr)))
+
 static void panic(const char *fmt, ...)
 {
   va_list vl;
@@ -33,16 +42,60 @@ static void error(const char *fmt, ...)
   exit(EXIT_FAILURE);
 }
 
-static char *addch(char *line, size_t len, int ch)
+static char *dsref(char *str)
+{
+  dstr *ds = dstr_of(str);
+  ds->refs++;
+  return str;
+}
+
+static void dsdrop(char *str)
 {
-  if (len + 2 > len) {
-    char *nline = realloc(line, len + 2);
+  if (str) {
+    dstr *ds = dstr_of(str);
+    assert (ds->refs > 0);
+    if (!--ds->refs)
+      free(ds);
+  }
+}
+
+static size_t dslen(const char *str)
+{
+  const dstr *ds = dstr_of(str);
+  return ds->len;
+}
+
+static char *dsgrow(char *str, size_t len)
+{
+  dstr *ds = str ? dstr_of(str) : 0;
+  size_t size = sizeof *ds + len + 1;
+
+  assert (ds == 0 || ds->refs == 1);
+
+  if (size < len)
+    panic("string size overflow");
+  ds = realloc(ds, size);
+  if (ds == 0)
+    panic("out of memory");
+
+  ds->refs = 1;
+  ds->len = len;
+  ds->str[len] = 0;
+
+  return ds->str;
+}
+
+static char *addch(char *line, int ch)
+{
+  size_t len = line ? dslen(line) : 0;
+
+  if (len + 1 > len) {
+    char *nline = dsgrow(line, len + 1);
 
     if (nline == 0)
       panic("out of memory");
 
     nline[len] = ch;
-    nline[len + 1] = 0;
     return nline;
   }
 
@@ -53,7 +106,6 @@ static char *addch(char *line, size_t len, int ch)
 static char *getln(FILE *stream)
 {
   char *line = 0;
-  size_t len = 0;
 
   for (;;) {
     int ch = getc(stream);
@@ -64,17 +116,17 @@ static char *getln(FILE *stream)
     if (ch == '\n') {
       if (line)
         return line;
-      return addch(line, len, 0);
+      return addch(line, 0);
     }
 
     if (ch == 127) {
-      line = addch(line, len++, '^');
-      line = addch(line, len++, '?');
+      line = addch(line, '^');
+      line = addch(line, '?');
     } else if (ch < 32) {
-      line = addch(line, len++, '^');
-      line = addch(line, len++, ch + 64);
+      line = addch(line, '^');
+      line = addch(line, ch + 64);
     } else {
-      line = addch(line, len++, ch);
+      line = addch(line, ch);
     }
   }
 }
@@ -94,7 +146,7 @@ static void usage(const char *name)
 
 static void drawline(const char *line, int hpos, int columns)
 {
-  size_t len = strlen(line);
+  size_t len = dslen(line);
 
   if ((size_t) hpos <= len) {
     if (hpos) {
@@ -261,7 +313,7 @@ int main(int argc, char **argv)
   if (isatty(fd)) {
     while ((line = getln(stdin))) {
       puts(line);
-      free(line);
+      dsdrop(line);
     }
 
     return 0;
@@ -328,9 +380,9 @@ int main(int argc, char **argv)
         if ((line = getln(stdin))) {
           if (nlines == maxlines) {
             if ((stat & stat_susp)) {
-              free(line);
+              dsdrop(line);
             } else {
-              free(circbuf[0]);
+              dsdrop(circbuf[0]);
               memmove(circbuf, circbuf + 1, (nlines - 1) * sizeof *circbuf);
               circbuf[nlines - 1] = line;
               stat |= stat_dirty;