1 files changed, 32 insertions, 23 deletions

diff --git a/src/nvim/os/wstream.c b/src/nvim/os/wstream.c
index 22524e2430..57afdd0e8f 100644
--- a/src/nvim/os/wstream.c
+++ b/src/nvim/os/wstream.c
@@ -20,14 +20,14 @@ struct wstream {
   bool freed;
 };
 
+struct wbuffer {
+  size_t refcount, size;
+  char *data;
+};
+
 typedef struct {
   WStream *wstream;
-  // Buffer containing data to be written
-  char *buffer;
-  // Size of the buffer
-  size_t length;
-  // If it's our responsibility to free the buffer
-  bool free;
+  WBuffer *buffer;
 } WriteData;
 
 static void write_cb(uv_write_t *req, int status);
@@ -59,51 +59,60 @@ void wstream_set_stream(WStream *wstream, uv_stream_t *stream)
   wstream->stream = stream;
 }
 
-bool wstream_write(WStream *wstream, char *buffer, size_t length, bool free)
+bool wstream_write(WStream *wstream, WBuffer *buffer)
 {
   WriteData *data;
   uv_buf_t uvbuf;
   uv_write_t *req;
 
-  if (wstream->freed) {
-    // Don't accept write requests after the WStream instance was freed
-    return false;
-  }
+  // This should not be called after a wstream was freed
+  assert(!wstream->freed);
 
-  if (wstream->curmem + length > wstream->maxmem) {
+  if (wstream->curmem + buffer->size > wstream->maxmem) {
     return false;
   }
 
-  if (free) {
-    // We should only account for buffers that are ours to free
-    wstream->curmem += length;
-  }
-
+  buffer->refcount++;
+  wstream->curmem += buffer->size;
   data = xmalloc(sizeof(WriteData));
   data->wstream = wstream;
   data->buffer = buffer;
-  data->length = length;
-  data->free = free;
   req = xmalloc(sizeof(uv_write_t));
   req->data = data;
-  uvbuf.base = buffer;
-  uvbuf.len = length;
+  uvbuf.base = buffer->data;
+  uvbuf.len = buffer->size;
   wstream->pending_reqs++;
   uv_write(req, wstream->stream, &uvbuf, 1, write_cb);
 
   return true;
 }
 
+WBuffer *wstream_new_buffer(char *data, size_t size, bool copy)
+{
+  WBuffer *rv = xmalloc(sizeof(WBuffer));
+  rv->size = size;
+  rv->refcount = 0;
+
+  if (copy) {
+    rv->data = xmemdup(data, size);
+  } else {
+    rv->data = data;
+  }
+
+  return rv;
+}
+
 static void write_cb(uv_write_t *req, int status)
 {
   WriteData *data = req->data;
 
   free(req);
+  data->wstream->curmem -= data->buffer->size;
 
-  if (data->free) {
+  if (!--data->buffer->refcount) {
     // Free the data written to the stream
+    free(data->buffer->data);
     free(data->buffer);
-    data->wstream->curmem -= data->length;
   }
 
   data->wstream->pending_reqs--;