Fix bad encoding of large upvalue indexes

src/compiler.c

@@ -1223,12 +1223,13 @@ static void block(size_t indentation, const char * blockName) {
 static void doUpvalues(Compiler * compiler, KrkCodeObject * function) {
 	assert(!!function->upvalueCount == !!compiler->upvalues);
 	for (size_t i = 0; i < function->upvalueCount; ++i) {
-		emitByte(compiler->upvalues[i].isLocal ? 1 : 0);
+		size_t index = compiler->upvalues[i].index;
-		if (i > 255) {
+		emitByte((compiler->upvalues[i].isLocal ? 1 : 0) | ((index > 255) ? 2 : 0));
-			emitByte((compiler->upvalues[i].index >> 16) & 0xFF);
+		if (index > 255) {
-			emitByte((compiler->upvalues[i].index >> 8) & 0xFF);
+			emitByte((index >> 16) & 0xFF);
+			emitByte((index >> 8) & 0xFF);
 		}
-		emitByte((compiler->upvalues[i].index) & 0xFF);
+		emitByte(index & 0xFF);
 	}
 }
 

src/debug.c

@@ -103,7 +103,15 @@ static int isJumpTarget(KrkCodeObject * func, size_t startPoint) {
 	case opc ## _LONG: { size = 4; more; break; }
 #define JUMP(opc,sign) case opc: { uint16_t jump = (chunk->code[offset + 1] << 8) | (chunk->code[offset + 2]); \
 	if ((size_t)(offset + 3 sign jump) == startPoint) return 1; size = 3; break; }
-#define CLOSURE_MORE size += AS_codeobject(chunk->constants.values[constant])->upvalueCount * 2
+#define CLOSURE_MORE \
+	KrkCodeObject * function = AS_codeobject(chunk->constants.values[constant]); \
+	for (size_t j = 0; j < function->upvalueCount; ++j) { \
+		int isLocal = chunk->code[offset++ + size]; \
+		offset++; \
+		if (isLocal & 2) { \
+			offset += 2; \
+		} \
+	}
 #define EXPAND_ARGS_MORE
 #define LOCAL_MORE
 
@@ -130,22 +138,20 @@ static int isJumpTarget(KrkCodeObject * func, size_t startPoint) {
 #define CONSTANT(opc,more) case opc: { size_t constant = chunk->code[offset + 1]; \
 	fprintf(f, "%-16s %4d ", opcodeClean(#opc), (int)constant); \
 	krk_printValueSafe(f, chunk->constants.values[constant]); \
-	more; \
+	size = 2; more; break; } \
-	size = 2; break; } \
 	case opc ## _LONG: { size_t constant = (chunk->code[offset + 1] << 16) | \
 	(chunk->code[offset + 2] << 8) | (chunk->code[offset + 3]); \
 	fprintf(f, "%-16s %4d ", opcodeClean(#opc "_LONG"), (int)constant); \
 	krk_printValueSafe(f, chunk->constants.values[constant]); \
-	more; size = 4; break; }
+	size = 4; more; break; }
 #define OPERANDB(opc,more) case opc: { uint32_t operand = chunk->code[offset + 1]; \
 	fprintf(f, "%-16s %4d", opcodeClean(#opc), (int)operand); \
-	more; size = 2; break; }
+	size = 2; more; break; }
 #define OPERAND(opc,more) OPERANDB(opc,more) \
 	case opc ## _LONG: { uint32_t operand = (chunk->code[offset + 1] << 16) | \
 	(chunk->code[offset + 2] << 8) | (chunk->code[offset + 3]); \
 	fprintf(f, "%-16s %4d", opcodeClean(#opc "_LONG"), (int)operand); \
-	more; fprintf(f,"\n"); \
+	size = 4; more; fprintf(f,"\n"); break; }
-	size = 4; break; }
 #define JUMP(opc,sign) case opc: { uint16_t jump = (chunk->code[offset + 1] << 8) | \
 	(chunk->code[offset + 2]); \
 	fprintf(f, "%-16s %4d (to %d)", opcodeClean(#opc), (int)jump, (int)(offset + 3 sign jump)); \
@@ -155,9 +161,13 @@ static int isJumpTarget(KrkCodeObject * func, size_t startPoint) {
 	KrkCodeObject * function = AS_codeobject(chunk->constants.values[constant]); \
 	fprintf(f, " "); \
 	for (size_t j = 0; j < function->upvalueCount; ++j) { \
-		int isLocal = chunk->code[offset++ + 2]; \
+		int isLocal = chunk->code[offset++ + size]; \
-		int index = chunk->code[offset++ + 2]; \
+		int index = chunk->code[offset++ + size]; \
-		if (isLocal) { \
+		if (isLocal & 2) { \
+			index = (index << 16) | (chunk->code[offset + size] << 8) | chunk->code[offset + 1 + size]; \
+			offset += 2; \
+		} \
+		if (isLocal & 1) { \
 			for (size_t i = 0; i < func->localNameCount; ++i) { \
 				if (func->localNames[i].id == (size_t)index && func->localNames[i].birthday <= offset && func->localNames[i].deathday >= offset) { \
 					fprintf(f, "%s", func->localNames[i].name->chars); \
@@ -621,7 +631,15 @@ KRK_FUNC(build,{
 	(chunk->code[offset + 2] << 8) | (chunk->code[offset + 3]); size = 4; more; break; }
 #define JUMP(opc,sign) case opc: { jump = 0 sign ((chunk->code[offset + 1] << 8) | (chunk->code[offset + 2])); \
 	size = 3; break; }
-#define CLOSURE_MORE size += AS_codeobject(chunk->constants.values[constant])->upvalueCount * 2
+#define CLOSURE_MORE \
+	KrkCodeObject * function = AS_codeobject(chunk->constants.values[constant]); \
+	for (size_t j = 0; j < function->upvalueCount; ++j) { \
+		int isLocal = chunk->code[offset++ + size]; \
+		offset++; \
+		if (isLocal & 2) { \
+			offset += 2; \
+		} \
+	}
 #define EXPAND_ARGS_MORE
 #define LOCAL_MORE local = operand;
 static KrkValue _examineInternal(KrkCodeObject* func) {

src/vm.c

@@ -2807,13 +2807,12 @@ _finishReturn: (void)0;
 				krk_push(OBJECT_VAL(closure));
 				for (size_t i = 0; i < closure->upvalueCount; ++i) {
 					int isLocal = READ_BYTE();
-					int index = 0;
+					int index = READ_BYTE();
-					if (i > 255) {
+					if (isLocal & 2) {
-						index = (frame->ip[0] << 16) | (frame->ip[1] << 8);
+						index = (index << 16) | (frame->ip[0] << 8) | (frame->ip[1]);
 						frame->ip += 2;
 					}
-					index |= READ_BYTE();
+					if (isLocal & 1) {
-					if (isLocal) {
 						closure->upvalues[i] = captureUpvalue(frame->slots + index);
 					} else {
 						closure->upvalues[i] = frame->closure->upvalues[index];

test/testLotsOfUpvalues.krk.expect

@@ -0,0 +1 @@
+124750

tools/make_test_upvalues.krk

@@ -0,0 +1,17 @@
+import fileio
+
+if __name__ == '__main__':
+    with fileio.open('test/testLotsOfUpvalues.krk','w') as f:
+        f.write('if True:\n')
+        for i in range(20):
+            f.write(f'    let red_herring_{i} = {i}\n')
+        for i in range(500):
+            f.write(f'    let a_{i} = {i}\n')
+        f.write('    def inner():\n')
+        f.write('        let s = 0\n')
+        for i in range(20):
+            f.write(f'        let inner_red_herring{i} = {i}\n')
+        for i in range(500):
+            f.write(f'        s += a_{i}\n')
+        f.write('        return s\n')
+        f.write('    print(inner())\n')