Free a hidden class and stop it copying a dispatch table

Test/AssociatedObjectCleanup.m

@@ -0,0 +1,38 @@
+#include "Test.h"
+
+// A hidden class carries uninstalled_dtable until the object is next
+// messaged, and a method is recorded in the class as it is installed into a
+// dtable, so the .cxx_destruct that carries the cleanup has to be recorded
+// when the hidden class is made.  An object destroyed without being messaged
+// again after gaining an association exercises that.
+
+static BOOL deallocCalled = NO;
+static const char *key = "AssociatedObjectCleanupKey";
+
+@interface Associated : Test
+@end
+
+@implementation Associated
+- (void)dealloc
+{
+	deallocCalled = YES;
+	[super dealloc];
+}
+@end
+
+int main(void)
+{
+	Associated *object = [Associated new];
+	Test *holder = [Test new];
+
+	objc_setAssociatedObject(holder, &key, object, OBJC_ASSOCIATION_RETAIN);
+	// The association holds the only reference.
+	[object release];
+	assert(!deallocCalled);
+
+	// Destroy the holder without sending it another message first.
+	object_dispose(holder);
+
+	assert(deallocCalled);
+	return 0;
+}

Test/CMakeLists.txt

@@ -17,6 +17,7 @@ set(TESTS
 	AllocatePair.m
 	AssociatedObject.m
 	AssociatedObject2.m
+	AssociatedObjectCleanup.m
 	BlockTest_arc.m
 	ConstantString.m
 	Category.m

associate.mm

@@ -35,7 +35,7 @@
 	uintptr_t policy;
 };
 
-#define REFERENCE_LIST_SIZE 10
+#define REFERENCE_LIST_SIZE 4
 
 /**
  * Linked list of references associated with an object.  We assume that there
@@ -45,8 +45,8 @@
 struct reference_list
 {
 	/**
-	 * Next group of references.  This is only ever used if we have more than
-	 * 10 references associated with an object, which seems highly unlikely.
+	 * Next group of references.  This is only ever used if we have more
+	 * references associated with an object than a single block holds.
 	 */
 	struct reference_list *next;
 	/**
@@ -213,7 +213,23 @@ static Class allocateHiddenClass(Class superclass)
 	newClass->info = objc_class_flag_resolved | objc_class_flag_user_created |
 		objc_class_flag_hidden_class | objc_class_flag_assoc_class;
 	newClass->super_class = superclass;
-	newClass->dtable = uninstalled_dtable;
+	// A hidden class adds no methods of its own until one is added through
+	// object_addMethod_np, so it can answer sends from its superclass's
+	// dispatch table rather than building a copy it would never change.
+	dtable_t superDtable = dtable_for_class((struct objc_class *)superclass);
+	if (superDtable != uninstalled_dtable)
+	{
+		newClass->dtable = superDtable;
+		newClass->info |= objc_class_flag_shared_dtable;
+		// The class answers the same methods as its superclass, so it has the
+		// same answers to the questions create_dtable_for_class would ask.
+		newClass->info |= (superclass->info &
+			(objc_class_flag_fast_arc | objc_class_flag_fast_alloc_init));
+	}
+	else
+	{
+		newClass->dtable = uninstalled_dtable;
+	}
 	newClass->instance_size = superclass->instance_size;
 
 	LOCK_RUNTIME_FOR_SCOPE();
@@ -227,14 +243,10 @@ static inline Class initHiddenClassForObject(id obj)
 {
 	Class hiddenClass = allocateHiddenClass(obj->isa); 
 	assert(!class_isMetaClass(obj->isa));
-	static SEL cxx_destruct;
-	if (NULL == cxx_destruct)
-	{
-		cxx_destruct = sel_registerName(".cxx_destruct");
-	}
-	const char *types = sizeof(void*) == 4 ? "v8@0:4" : "v16@0:8";
-	class_addMethod(hiddenClass, cxx_destruct,
-		(IMP)deallocHiddenClass, types);
+	// Recorded in the class rather than added as a method: the class may be
+	// answering sends from its superclass's dispatch table, which must not
+	// gain this.  call_cxx_destruct reads the field.
+	hiddenClass->cxx_destruct = (IMP)deallocHiddenClass;
 	obj->isa = hiddenClass;
 	return hiddenClass;
 }
@@ -250,7 +262,10 @@ static void deallocHiddenClass(id obj, SEL _cmd)
 	cleanupReferenceList(list);
 	freeReferenceList(list->next);
 	//fprintf(stderr, "Deallocating dtable %p\n", hiddenClass->dtable);
-	free_dtable(hiddenClass->dtable);
+	if (!objc_test_class_flag(hiddenClass, objc_class_flag_shared_dtable))
+	{
+		free_dtable(hiddenClass->dtable);
+	}
 	// We shouldn't have any subclasses left at this point
 	assert(hiddenClass->subclass_list == 0);
 	// Remove the class from the subclass list of its superclass
@@ -426,14 +441,30 @@ static Class hiddenClassForObject(id object)
 	return hiddenClass;
 }
 
+/**
+ * Gives a hidden class a dispatch table of its own, so that a method added to
+ * it does not reach the table its superclass is using.
+ */
+static Class hiddenClassForMethodAdd(id object)
+{
+	Class hiddenClass = hiddenClassForObject(object);
+	LOCK_RUNTIME_FOR_SCOPE();
+	if (objc_test_class_flag(hiddenClass, objc_class_flag_shared_dtable))
+	{
+		objc_clear_class_flag(hiddenClass, objc_class_flag_shared_dtable);
+		hiddenClass->dtable = create_dtable_for_class(hiddenClass, uninstalled_dtable);
+	}
+	return hiddenClass;
+}
+
 BOOL object_addMethod_np(id object, SEL name, IMP imp, const char *types)
 {
-	return class_addMethod(hiddenClassForObject(object), name, imp, types);
+	return class_addMethod(hiddenClassForMethodAdd(object), name, imp, types);
 }
 
 IMP object_replaceMethod_np(id object, SEL name, IMP imp, const char *types)
 {
-	return class_replaceMethod(hiddenClassForObject(object), name, imp, types);
+	return class_replaceMethod(hiddenClassForMethodAdd(object), name, imp, types);
 }
 static char prototypeKey;
 

class.h

@@ -343,6 +343,11 @@ enum objc_class_flags
 	 * On a class, guarantees that `+init` is trivial.
 	 */
 	objc_class_flag_fast_alloc_init = (1<<15),
+	/**
+	 * This class's dtable belongs to its superclass and must not be modified
+	 * or freed through this class.
+	 */
+	objc_class_flag_shared_dtable = (1<<19),
 	/**
 	 * The class is a block class. Reference count management must be done by
 	 * the underlying blocks runtime.

dtable.h

@@ -119,6 +119,7 @@ static inline int classHasDtable(struct objc_class *cls)
  * Updates the dtable for a class and its subclasses.  Must be called after
  * modifying a class's method list.
  */
+dtable_t create_dtable_for_class(Class, dtable_t);
 void objc_update_dtable_for_class(Class);
 /**
  * Updates the dtable for a class and its subclasses.  Must be called after