Invalidate caches of subclass method pointers when they change in a superclass

src/kuroko/object.h

@@ -207,6 +207,8 @@ typedef struct KrkClass {
 	KrkObj * _descset;        /**< @brief @c %__set__      Called when a descriptor object is assigned to as a property */
 	KrkObj * _classgetitem;   /**< @brief @c %__class_getitem__ Class method called when a type object is subscripted; used for type hints */
 	KrkObj * _hash;           /**< @brief @c %__hash__     Called when an instance is a key in a dict or an entry in a set */
+
+	KrkTable subclasses;
 } KrkClass;
 
 /**

src/memory.c

@@ -233,6 +233,10 @@ static void freeObject(KrkObj * object) {
 		case KRK_OBJ_CLASS: {
 			KrkClass * _class = (KrkClass*)object;
 			krk_freeTable(&_class->methods);
+			krk_freeTable(&_class->subclasses);
+			if (_class->base) {
+				krk_tableDelete(&_class->base->subclasses, OBJECT_VAL(object));
+			}
 			FREE(KrkClass, object);
 			break;
 		}

src/obj_base.c

@@ -59,6 +59,23 @@ static KrkValue _class_to_str(int argc, KrkValue argv[], int hasKw) {
 	return OBJECT_VAL(out);
 }
 
+static KrkValue _class_subclasses(int argc, KrkValue argv[], int hasKw) {
+	if (!IS_CLASS(argv[0])) return krk_runtimeError(vm.exceptions->typeError, "expected class");
+
+	KrkClass * _class = AS_CLASS(argv[0]);
+
+	KrkValue myList = krk_list_of(0,NULL,0);
+	krk_push(myList);
+
+	for (size_t i = 0; i < _class->subclasses.capacity; ++i) {
+		KrkTableEntry * entry = &_class->subclasses.entries[i];
+		if (IS_KWARGS(entry->key)) continue;
+		krk_writeValueArray(AS_LIST(myList), entry->key);
+	}
+
+	return krk_pop();
+}
+
 _noexport
 void _createAndBind_type(void) {
 	ADD_BASE_CLASS(vm.baseClasses->typeClass, "type", vm.baseClasses->objectClass);
@@ -69,6 +86,7 @@ void _createAndBind_type(void) {
 	krk_defineNative(&vm.baseClasses->typeClass->methods, "__init__", _type_init);
 	krk_defineNative(&vm.baseClasses->typeClass->methods, "__str__", _class_to_str);
 	krk_defineNative(&vm.baseClasses->typeClass->methods, "__repr__", _class_to_str);
+	krk_defineNative(&vm.baseClasses->typeClass->methods, "__subclasses__", _class_subclasses);
 	krk_finalizeClass(vm.baseClasses->typeClass);
 	KRK_DOC(vm.baseClasses->typeClass, "Obtain the object representation of the class of an object.");
 }

src/object.c

@@ -312,12 +312,15 @@ KrkClass * krk_newClass(KrkString * name, KrkClass * baseClass) {
 	_class->name = name;
 	_class->allocSize = sizeof(KrkInstance);
 	krk_initTable(&_class->methods);
+	krk_initTable(&_class->subclasses);
 
 	if (baseClass) {
 		_class->base = baseClass;
 		_class->allocSize = baseClass->allocSize;
 		_class->_ongcscan = baseClass->_ongcscan;
 		_class->_ongcsweep = baseClass->_ongcsweep;
+
+		krk_tableSet(&baseClass->subclasses, OBJECT_VAL(_class), NONE_VAL());
 	}
 
 	return _class;

src/table.c

@@ -49,6 +49,9 @@ inline int krk_hashValue(KrkValue value, uint32_t *hashOut) {
 		*hashOut = (uint32_t)AS_INTEGER(result);
 		return 0;
 	}
+	if (IS_CLASS(value)) {
+		return INTEGER_VAL((int)(intptr_t)AS_OBJECT(value));
+	}
 _unhashable:
 	if (IS_NONE(krk_currentThread.currentException))
 		krk_runtimeError(vm.exceptions->typeError, "unhashable type: '%s'", krk_typeName(value));

src/vm.c

@@ -439,6 +439,7 @@ void krk_finalizeClass(KrkClass * _class) {
 	};
 
 	for (struct TypeMap * entry = specials; entry->method; ++entry) {
+		*entry->method = NULL;
 		KrkClass * _base = _class;
 		while (_base) {
 			if (krk_tableGet(&_base->methods, vm.specialMethodNames[entry->index], &tmp)) break;
@@ -454,6 +455,12 @@ void krk_finalizeClass(KrkClass * _class) {
 			_class->_hash = NULL;
 		}
 	}
+
+	for (size_t i = 0; i < _class->subclasses.capacity; ++i) {
+		KrkTableEntry * entry = &_class->subclasses.entries[i];
+		if (IS_KWARGS(entry->key)) continue;
+		krk_finalizeClass(AS_CLASS(entry->key));
+	}
 }
 
 /**
@@ -2109,6 +2116,9 @@ static int valueDelProperty(KrkString * name) {
 		if (!krk_tableDelete(&_class->methods, OBJECT_VAL(name))) {
 			return 0;
 		}
+		if (name->length && name->chars[0] == '_') {
+			krk_finalizeClass(_class);
+		}
 		krk_pop(); /* the original value */
 		return 1;
 	}
@@ -2455,10 +2465,14 @@ _finishReturn: (void)0;
 					goto _finishException;
 				}
 				KrkClass * subclass = AS_CLASS(krk_peek(1));
+				if (subclass->base) {
+					krk_tableDelete(&subclass->base->subclasses, krk_peek(1));
+				}
 				subclass->base = AS_CLASS(superclass);
 				subclass->allocSize = AS_CLASS(superclass)->allocSize;
 				subclass->_ongcsweep = AS_CLASS(superclass)->_ongcsweep;
 				subclass->_ongcscan = AS_CLASS(superclass)->_ongcscan;
+				krk_tableSet(&AS_CLASS(superclass)->subclasses, krk_peek(1), NONE_VAL());
 				krk_pop(); /* Super class */
 				break;
 			}

test/testSubclassCachedMethods.krk

@@ -0,0 +1,36 @@
+class A():
+    def __iter__(self):
+        def _():
+            return _
+        return _
+class B(A):
+    pass
+
+let a = A()
+let b = B()
+
+def __iter__(self):
+    let o = 4
+    def _():
+        if o:
+            return (o -= 1)
+        return _
+    return _
+
+A.__iter__ = __iter__
+
+for i in a:
+    print(i)
+
+for i in b:
+    print(i)
+
+# now what if we remove it
+del A.__iter__
+
+try:
+    for i in b:
+        print(i)
+    print('did not raise attribute error')
+except AttributeError:
+    print('raised attribute error')

test/testSubclassCachedMethods.krk.expect

@@ -0,0 +1,9 @@
+3
+2
+1
+0
+3
+2
+1
+0
+raised attribute error