From b3a660f8ecab54d5bdeb26059ae506cbdfbc8d84 Mon Sep 17 00:00:00 2001
From: Nunuhara Cabbage <nunuhara@haniwa.technology>
Date: Sun, 20 Oct 2019 21:26:51 -0700
Subject: [PATCH] strings: copy-on-write

Only copy strings when they are mutated. This avoids copying strings in
cases where neither the copy nor the original would be modified during
the lifetime of the copy, e.g. when passing strings by-value.

This is purely an optimization.
---
 ain.c       |   7 +-
 string.c    | 237 +++++++++++++++++++++++++++++-----------------------
 system4.md  |  24 +++++-
 vm.c        |  54 ++++++------
 vm_string.h |  28 ++++---
 5 files changed, 203 insertions(+), 147 deletions(-)

diff --git a/ain.c b/ain.c
index 8e9ba4c..294ed9a 100644
--- a/ain.c
+++ b/ain.c
@@ -59,11 +59,8 @@ static int32_t read_int32(struct ain_reader *r)
 
 static uint8_t *read_code(struct ain_reader *r, size_t len)
 {
-	uint8_t *bytes = xmalloc(len + 6);
+	uint8_t *bytes = xmalloc(len);
 	memcpy(bytes, r->buf + r->index, len);
-	// XXX: main() returns to CALLSYS 0x0 (system.Exit)
-	LittleEndian_putW(bytes, len, CALLSYS);
-	LittleEndian_putDW(bytes, len + 2, 0);
 	r->index += len;
 	return bytes;
 }
@@ -91,7 +88,7 @@ static struct string *read_vm_string(struct ain_reader *r)
 {
 	size_t len = strlen((char*)r->buf + r->index);
 	struct string *s = make_string((char*)r->buf + r->index, len);
-	s->literal = true;
+	s->cow = true;
 	r->index += len + 1;
 	return s;
 }
diff --git a/string.c b/string.c
index 7dc5ab4..19e6965 100644
--- a/string.c
+++ b/string.c
@@ -22,127 +22,52 @@
 #include "vm_string.h"
 #include "utfsjis.h"
 
+int nr_strings = 0;
+
+static struct string *alloc_string(int size)
+{
+	nr_strings++;
+	return xmalloc(sizeof(struct string) + size + 1);
+}
+
+void free_string(struct string *str)
+{
+	if (!str->ref)
+		ERROR("Double free of string object");
+	if (!--str->ref) {
+		free(str);
+		nr_strings--;
+	}
+}
+
 struct string *make_string(const char *str, unsigned int len)
 {
-	struct string *s = xmalloc(sizeof(struct string) + len + 1);
+	struct string *s = alloc_string(len);
 	s->size = len;
-	s->literal = false;
+	s->ref = 1;
+	s->cow = 0;
 	memcpy(s->text, str, len);
 	s->text[len] = '\0';
 	return s;
 }
 
-void free_string(struct string *str)
+struct string *string_ref(struct string *s)
 {
-	if (!str->literal)
-		free(str);
+	s->cow = 1;
+	s->ref++;
+	return s;
 }
 
 struct string *string_dup(const struct string *in)
 {
-	struct string *out = xmalloc(sizeof(struct string) + in->size + 1);
+	struct string *out = alloc_string(in->size);
 	out->size = in->size;
-	out->literal = false;
+	out->ref = 1;
+	out->cow = 0;
 	memcpy(out->text, in->text, in->size + 1);
 	return out;
 }
 
-void string_append(struct string **_a, const struct string *b)
-{
-	struct string *a = *_a;
-	size_t a_size = a->size;
-
-	if (a->literal) {
-		a = xmalloc(sizeof(struct string) + a->size + b->size + 1);
-		a->literal = false;
-		memcpy(a->text, (*_a)->text, a_size+1);
-	} else {
-		a = xrealloc(a, sizeof(struct string) + a->size + b->size + 1);
-	}
-	a->size = a_size + b->size;
-	memcpy(a->text + a_size, b->text, b->size + 1);
-	*_a = a;
-}
-
-struct string *string_concatenate(const struct string *a, const struct string *b)
-{
-	struct string *s = xmalloc(sizeof(struct string) + a->size + b->size + 1);
-	s->size = a->size + b->size;
-	s->literal = false;
-	memcpy(s->text, a->text, a->size);
-	memcpy(s->text + a->size, b->text, b->size + 1);
-	return s;
-}
-
-struct string *string_copy(const struct string *s, int index, int len)
-{
-	if (index < 0)
-		index = 0;
-	if (len <= 0)
-		return make_string("", 0);
-	if ((index = sjis_index(s->text, index)) < 0)
-		return make_string("", 0);
-
-	if ((len = sjis_index(s->text + index, len)) < 0)
-		len = s->size - index;
-
-	return make_string(s->text + index, len);
-}
-
-int string_find(const struct string *haystack, const struct string *needle)
-{
-	int c = 0;
-	for (int i = 0; i < haystack->size; i++, c++) {
-		if (!strncmp(haystack->text+i, needle->text, needle->size))
-			return c;
-		if (SJIS_2BYTE(haystack->text[i])) {
-			i++;
-		}
-	}
-	return -1;
-}
-
-void string_push_back(struct string **s, int c)
-{
-	int bytes = SJIS_2BYTE(c) ? 2 : 1;
-	*s = xrealloc(*s, sizeof(struct string) + (*s)->size + bytes + 1);
-	(*s)->text[(*s)->size++] = c & 0xFF;
-	if (bytes == 2) {
-		(*s)->text[(*s)->size++] = c >> 8;
-	}
-	(*s)->text[(*s)->size] = '\0';
-}
-
-void string_pop_back(struct string *s)
-{
-	// get index of last character
-	int c = 0;
-	for (int i = 0; i < s->size; i++) {
-		c = i;
-		if (SJIS_2BYTE(s->text[i])) {
-			i++;
-		}
-	}
-	s->text[c] = '\0';
-	s->size = c;
-}
-
-void string_erase(struct string *s, int index)
-{
-	int bytes;
-	if (index < 0)
-		index = 0;
-	if (index >= s->size)
-		return;
-	if ((index = sjis_index(s->text, index)) < 0)
-		return;
-	bytes = SJIS_2BYTE(s->text[index]) ? 2 : 1;
-
-	memcpy(s->text + index, s->text + index + bytes, s->size - index - bytes);
-	s->size -= bytes;
-	s->text[s->size] = '\0';
-}
-
 struct string *integer_to_string(int n)
 {
 	char buf[512];
@@ -163,3 +88,109 @@ struct string *float_to_string(float f, int precision)
 	len = snprintf(buf, 512, "%.*f", precision, f);
 	return make_string(buf, len);
 }
+
+struct string *string_concatenate(const struct string *a, const struct string *b)
+{
+	struct string *s = alloc_string(a->size + b->size);
+	s->size = a->size + b->size;
+	s->ref = 1;
+	s->cow = 0;
+	memcpy(s->text, a->text, a->size);
+	memcpy(s->text + a->size, b->text, b->size + 1);
+	return s;
+}
+
+struct string *string_copy(const struct string *s, int index, int len)
+{
+	if (index < 0)
+		index = 0;
+	if (len <= 0)
+		return make_string("", 0);
+	if ((index = sjis_index(s->text, index)) < 0)
+		return make_string("", 0);
+
+	if ((len = sjis_index(s->text + index, len)) < 0)
+		len = s->size - index;
+
+	return make_string(s->text + index, len);
+}
+
+static struct string *cow_check(struct string *s)
+{
+	if (s->cow && s->ref > 1) {
+		struct string *out = string_dup(s);
+		free_string(s);
+		return out;
+	}
+	if (s->cow)
+		s->cow = 0;
+	return s;
+}
+
+void string_append(struct string **_a, const struct string *b)
+{
+	struct string *a = *_a = cow_check(*_a);
+	size_t a_size = a->size;
+
+	a = xrealloc(a, sizeof(struct string) + a->size + b->size + 1);
+	a->size = a_size + b->size;
+	memcpy(a->text + a_size, b->text, b->size + 1);
+	*_a = a;
+}
+
+void string_push_back(struct string **s, int c)
+{
+	int bytes = SJIS_2BYTE(c) ? 2 : 1;
+	*s = cow_check(*s);
+	*s = xrealloc(*s, sizeof(struct string) + (*s)->size + bytes + 1);
+	(*s)->text[(*s)->size++] = c & 0xFF;
+	if (bytes == 2) {
+		(*s)->text[(*s)->size++] = c >> 8;
+	}
+	(*s)->text[(*s)->size] = '\0';
+}
+
+void string_pop_back(struct string **s)
+{
+	*s = cow_check(*s);
+	// get index of last character
+	int c = 0;
+	for (int i = 0; i < (*s)->size; i++) {
+		c = i;
+		if (SJIS_2BYTE((*s)->text[i])) {
+			i++;
+		}
+	}
+	(*s)->text[c] = '\0';
+	(*s)->size = c;
+}
+
+void string_erase(struct string **s, int index)
+{
+	int bytes;
+	if (index < 0)
+		index = 0;
+	if (index >= (*s)->size)
+		return;
+	if ((index = sjis_index((*s)->text, index)) < 0)
+		return;
+	bytes = SJIS_2BYTE((*s)->text[index]) ? 2 : 1;
+
+	*s = cow_check(*s);
+	memcpy((*s)->text + index, (*s)->text + index + bytes, (*s)->size - index - bytes);
+	(*s)->size -= bytes;
+	(*s)->text[(*s)->size] = '\0';
+}
+
+int string_find(const struct string *haystack, const struct string *needle)
+{
+	int c = 0;
+	for (int i = 0; i < haystack->size; i++, c++) {
+		if (!strncmp(haystack->text+i, needle->text, needle->size))
+			return c;
+		if (SJIS_2BYTE(haystack->text[i])) {
+			i++;
+		}
+	}
+	return -1;
+}
diff --git a/system4.md b/system4.md
index 8c3b2e4..82c0478 100644
--- a/system4.md
+++ b/system4.md
@@ -13,7 +13,8 @@ from the stack and stored in a "page" by the CALLFUNC instruction.
 ### Pages
 
 Non-pointer variables are stored in "pages". There is a global page for global
-variables, and a local page for each function call.
+variables, and a local page for each function call. Structures and arrays are
+also implemented as pages.
 
 The SH_LOCAL* family of instructions implicitly operate on the current page.
 
@@ -50,7 +51,28 @@ memory leaks by creating circular references. E.g. the following code leaks:
         // they each hold a reference to the other, even after they're both out
         // of scope
     }
+
+Another point to note is that reference counting operates at the page level. So
+if you create a reference to a local variable which escapes the normal scope of
+that variable, *every other variable in the same page* remains referenced until
+that reference is deleted. E.g.
+
+    struct a {
+        int a;
+        ~a() { system.Output("destructor called\n"); }
+    };
     
+    ref int global_ref;
+    
+    void leak(void)
+    {
+        int i;
+        a local_a;
+        global_ref <- i;
+        // the destructor for local_a is never called, because it lives in the
+        // same page as i, which is referenced by a global
+    }
+
 Calling Convention
 ------------------
 
diff --git a/vm.c b/vm.c
index d6b524c..77feac8 100644
--- a/vm.c
+++ b/vm.c
@@ -222,11 +222,6 @@ static struct string *stack_peek_string(int n)
 	return heap[stack_peek(n).i].s;
 }
 
-static struct string *stack_pop_string(void)
-{
-	return heap[stack[--stack_ptr].i].s;
-}
-
 /*
  * System 4 calling convention:
  *   - caller pushes arguments, in order
@@ -352,7 +347,8 @@ static int struct_copy(int no, int src_slot)
 }
 
 static struct string EMPTY_STRING = {
-	.literal = true,
+	.cow = true,
+	.ref = 1,
 	.size = 0,
 	.text = ""
 };
@@ -366,7 +362,7 @@ static void create_struct(int no, union vm_value *var)
 		switch (s->members[i].data_type) {
 		case AIN_STRING:
 			memb = heap_alloc_slot(VM_STRING);
-			heap[memb].s = &EMPTY_STRING;
+			heap[memb].s = string_ref(&EMPTY_STRING);
 			heap[slot].page[i].i = memb;
 			break;
 		case AIN_STRUCT:
@@ -404,13 +400,6 @@ static void execute_instruction(int16_t opcode)
 	case F_PUSH:
 		stack_push(get_argument_float(0));
 		break;
-	case S_PUSH:
-		stack_push_string(ain->strings[get_argument(0)]);
-		break;
-	case S_POP:
-		index = stack_pop().i;
-		heap_unref(index);
-		break;
 	case REF:
 		// Dereference a reference to a value.
 		index = stack_pop_var()[0].i;
@@ -423,11 +412,6 @@ static void execute_instruction(int16_t opcode)
 		stack_push(ref[0].i);
 		stack_push(ref[1].i);
 		break;
-	case S_REF:
-		// Dereference a reference to a string
-		index = stack_pop_var()->i;
-		stack_push_string(string_dup(heap[index].s));
-		break;
 	case DUP:
 		// A -> AA
 		stack_push(stack_peek(0).i);
@@ -505,7 +489,7 @@ static void execute_instruction(int16_t opcode)
 		create_struct(get_argument(1), local_ptr(get_argument(0)));
 		break;
 	//
-	// --- Function Calls ---
+	// --- Control Flow ---
 	//
 	case CALLFUNC:
 		function_call(get_argument(0), instr_ptr + instruction_width(CALLFUNC));
@@ -519,9 +503,6 @@ static void execute_instruction(int16_t opcode)
 	case CALLSYS:
 		system_call(get_argument(0));
 		break;
-	//
-	// --- Control Flow ---
-	//
 	case JUMP: // ADDR
 		instr_ptr = get_argument(0);
 		break;
@@ -726,29 +707,42 @@ static void execute_instruction(int16_t opcode)
 	//
 	// --- Strings ---
 	//
+	case S_PUSH:
+		stack_push_string(string_ref(ain->strings[get_argument(0)]));
+		break;
+	case S_POP:
+		index = stack_pop().i;
+		heap_unref(index);
+		break;
+	case S_REF:
+		// Dereference a reference to a string
+		index = stack_pop_var()->i;
+		stack_push_string(string_ref(heap[index].s));
+		break;
 	case S_ASSIGN: // A = B
 		b = stack_peek(0).i;
 		a = stack_peek(1).i;
 		if (heap[a].s) {
 			free_string(heap[a].s);
 		}
-		heap[a].s = string_dup(heap[b].s);
+		heap[a].s = string_ref(heap[b].s);
 		// remove A from the stack, but leave B
 		stack_set(1, b);
 		stack_pop();
 		break;
 	case S_PLUSA2:
-		b = stack_peek(0).i;
 		a = stack_peek(1).i;
+		b = stack_peek(0).i;
 		string_append(&heap[a].s, heap[b].s);
 		heap_unref(b);
 		stack_pop();
 		stack_pop();
-		stack_push_string(string_dup(heap[a].s));
+		stack_push_string(string_ref(heap[a].s));
 		break;
 	case S_ADD:
 		b = stack_pop().i;
 		a = stack_pop().i;
+		// TODO: can use string_append here?
 		stack_push_string(string_concatenate(heap[a].s, heap[b].s));
 		heap_unref(a);
 		heap_unref(b);
@@ -822,13 +816,15 @@ static void execute_instruction(int16_t opcode)
 		break;
 	//case S_POPBACK: // ???
 	case S_POPBACK2:
-		string_pop_back(stack_pop_string());
+		index = stack_pop().i;
+		string_pop_back(&heap[index].s);
 		break;
 	//case S_ERASE: // ???
 	case S_ERASE2:
 		b = stack_pop().i; // ???
 		a = stack_pop().i; // index
-		string_erase(stack_pop_string(), a);
+		index = stack_pop().i;
+		string_erase(&heap[index].s, a);
 		break;
 	case I_STRING:
 		stack_push_string(integer_to_string(stack_pop().i));
@@ -872,6 +868,7 @@ static void vm_execute(void)
 	}
 }
 
+
 void vm_execute_ain(struct ain *program)
 {
 	// initialize VM state
@@ -922,3 +919,4 @@ void vm_execute_ain(struct ain *program)
 
 	vm_call(ain->main, -1);
 }
+
diff --git a/vm_string.h b/vm_string.h
index d963857..4578f2c 100644
--- a/vm_string.h
+++ b/vm_string.h
@@ -18,25 +18,33 @@
 #define SYSTEM4_STRING_H
 
 #include <stdbool.h>
-#include <stddef.h>
+#include <stdlib.h>
 
 struct string {
 	int size;
-	bool literal;
+	unsigned int ref : 24;
+	unsigned int cow : 1;
 	char text[];
 };
 
-struct string *make_string(const char *str, unsigned int len);
 void free_string(struct string *str);
+
+// constructors
+struct string *make_string(const char *str, unsigned int len);
+struct string *string_ref(struct string *s);
 struct string *string_dup(const struct string *in);
-void string_append(struct string **a, const struct string *b);
-struct string *string_concatenate(const struct string *a, const struct string *b);
-struct string *string_copy(const struct string *s, int index, int len);
-int string_find(const struct string *haystack, const struct string *needle);
-void string_push_back(struct string **s, int c);
-void string_pop_back(struct string *s);
-void string_erase(struct string *s, int index);
 struct string *integer_to_string(int n);
 struct string *float_to_string(float f, int precision);
+struct string *string_concatenate(const struct string *a, const struct string *b);
+struct string *string_copy(const struct string *s, int index, int len);
+
+// mutators
+void string_append(struct string **a, const struct string *b);
+void string_push_back(struct string **s, int c);
+void string_pop_back(struct string **s);
+void string_erase(struct string **s, int index);
+
+// queries
+int string_find(const struct string *haystack, const struct string *needle);
 
 #endif