From 9b5ce15bf73af7225331b30b134edd248152289b Mon Sep 17 00:00:00 2001
From: "K. Lange" <klange@toaruos.org>
Date: Sun, 10 Jul 2022 13:13:27 +0900
Subject: [PATCH] Implement Python's identifier mangling

---
 src/compiler.c                   | 34 +++++++++++++++++++++++++---
 test/testNameMangling.krk        | 39 ++++++++++++++++++++++++++++++++
 test/testNameMangling.krk.expect |  1 +
 3 files changed, 71 insertions(+), 3 deletions(-)
 create mode 100644 test/testNameMangling.krk
 create mode 100644 test/testNameMangling.krk.expect

diff --git a/src/compiler.c b/src/compiler.c
index 88f823f..6088a70 100644
--- a/src/compiler.c
+++ b/src/compiler.c
@@ -598,7 +598,31 @@ static size_t emitConstant(KrkValue value) {
 	return krk_writeConstant(currentChunk(), value, parser.previous.line);
 }
 
+static int isMangleable(const char * name, size_t len) {
+	return (len > 2 && name[0] == '_' && name[1] == '_' && name[len-1] != '_' && (len < 4 || name[len-2] != '_'));
+}
+
 static ssize_t identifierConstant(KrkToken * name) {
+	if (currentClass && isMangleable(name->start, name->length)) {
+		/* Mangle it */
+		const char * className = currentClass->name.start;
+		size_t classLength = currentClass->name.length;
+		while (classLength && *className == '_') {
+			classLength--;
+			className++;
+		}
+
+		size_t total = name->length + 2 + classLength;
+		char * mangled = malloc(total);
+		snprintf(mangled, total, "_%.*s%.*s",
+			(int)classLength, className,
+			(int)name->length, name->start);
+		return krk_addConstant(currentChunk(), OBJECT_VAL(krk_takeString(mangled, total-1)));
+	}
+	return krk_addConstant(currentChunk(), OBJECT_VAL(krk_copyString(name->start, name->length)));
+}
+
+static ssize_t nonidentifierTokenConstant(KrkToken * name) {
 	return krk_addConstant(currentChunk(), OBJECT_VAL(krk_copyString(name->start, name->length)));
 }
 
@@ -1545,7 +1569,7 @@ static void classBody(size_t blockWidth) {
 
 #define ATTACH_PROPERTY(propName,how,propValue) do { \
 	KrkToken val_tok = syntheticToken(propValue); \
-	size_t val_ind = identifierConstant(&val_tok); \
+	size_t val_ind = nonidentifierTokenConstant(&val_tok); \
 	EMIT_OPERAND_OP(how, val_ind); \
 	KrkToken name_tok = syntheticToken(propName); \
 	size_t name_ind = identifierConstant(&name_tok); \
@@ -1563,9 +1587,13 @@ static KrkToken classDeclaration(void) {
 
 	beginScope();
 
-	size_t constInd = identifierConstant(&parser.previous);
+	/* We want to expose the mangled name within the class definition, which then
+	 * becomes available as a nonlocal, but we want to hand the non-mangled name
+	 * to the CLASS instruction so that __name__ is right. */
+	size_t nameInd = nonidentifierTokenConstant(&parser.previous);
+	identifierConstant(&parser.previous);
 	declareVariable();
-	EMIT_OPERAND_OP(OP_CLASS, constInd);
+	EMIT_OPERAND_OP(OP_CLASS, nameInd);
 	markInitialized();
 
 	ClassCompiler classCompiler;
diff --git a/test/testNameMangling.krk b/test/testNameMangling.krk
new file mode 100644
index 0000000..48033f9
--- /dev/null
+++ b/test/testNameMangling.krk
@@ -0,0 +1,39 @@
+class Foo():
+    def __test(self):
+        pass
+
+    def __test__(self):
+        pass
+
+    def ___(self):
+        pass
+
+let f = Foo()
+
+let dirResults = dir(f)
+
+# Straightforard name-mangling and non-mangling casses
+assert '_Foo__test' in dirResults
+assert '__test__' in dirResults
+assert '_Foo__test__' not in dirResults
+assert '___' in dirResults
+assert '_Foo____' not in dirResults
+
+class Baz():
+    def __init__(self):
+        class __Inner():
+            def __test(self):
+                pass
+        self.__inner = __Inner
+
+let b = Baz()
+
+assert '_Baz__inner' in dir(b)
+
+# Inner class name should not be mangled on its own
+assert b._Baz__inner.__name__ == '__Inner'
+
+# Mangled identifiers should reduce leading underscores from class name to one
+assert '_Inner__test' in dir(b._Baz__inner)
+
+print("okay")
diff --git a/test/testNameMangling.krk.expect b/test/testNameMangling.krk.expect
new file mode 100644
index 0000000..dcf02b2
--- /dev/null
+++ b/test/testNameMangling.krk.expect
@@ -0,0 +1 @@
+okay