--- xnu/xnu-6153.41.3/libkern/c++/OSCollectionIterator.cpp
+++ xnu/xnu-8792.61.2/libkern/c++/OSCollectionIterator.cpp
@@ -27,10 +27,13 @@
  */
 /* IOArray.h created by rsulack on Thu 11-Sep-1997 */
 
+#define IOKIT_ENABLE_SHARED_PTR
+
+#include <libkern/c++/OSArray.h>
+#include <libkern/c++/OSCollection.h>
 #include <libkern/c++/OSCollectionIterator.h>
-#include <libkern/c++/OSCollection.h>
-#include <libkern/c++/OSArray.h>
 #include <libkern/c++/OSLib.h>
+#include <libkern/c++/OSSharedPtr.h>
 
 #define super OSIterator
 
@@ -43,8 +46,7 @@
 		return false;
 	}
 
-	inColl->retain();
-	collection = inColl;
+	collection.reset(inColl, OSRetain);
 	collIterator = NULL;
 	initialUpdateStamp = 0;
 	valid = false;
@@ -52,14 +54,13 @@
 	return true;
 }
 
-OSCollectionIterator *
+OSSharedPtr<OSCollectionIterator>
 OSCollectionIterator::withCollection(const OSCollection *inColl)
 {
-	OSCollectionIterator *me = new OSCollectionIterator;
+	OSSharedPtr<OSCollectionIterator> me = OSMakeShared<OSCollectionIterator>();
 
 	if (me && !me->initWithCollection(inColl)) {
-		me->release();
-		return NULL;
+		return nullptr;
 	}
 
 	return me;
@@ -68,57 +69,247 @@
 void
 OSCollectionIterator::free()
 {
-	if (collIterator) {
-		kfree(collIterator, collection->iteratorSize());
+	freeIteratorStorage();
+
+	collection.reset();
+
+	super::free();
+}
+
+void
+OSCollectionIterator::reset()
+{
+	valid = false;
+	bool initialized = initializeIteratorStorage();
+
+	if (!initialized) {
+		// reusing existing storage
+		void * storage = getIteratorStorage();
+		bzero(storage, collection->iteratorSize());
+
+		if (!collection->initIterator(storage)) {
+			return;
+		}
+
+		initialUpdateStamp = collection->updateStamp;
+		valid = true;
+	}
+}
+
+bool
+OSCollectionIterator::isValid()
+{
+	initializeIteratorStorage();
+
+	if (!valid || collection->updateStamp != initialUpdateStamp) {
+		return false;
+	}
+
+	return true;
+}
+
+bool
+OSCollectionIterator::initializeIteratorStorage()
+{
+	void * result = NULL;
+	bool initialized = false;
+
+#if __LP64__
+	OSCollectionIteratorStorageType storageType = getStorageType();
+	switch (storageType) {
+	case OSCollectionIteratorStorageUnallocated:
+		if (collection->iteratorSize() > sizeof(inlineStorage) || isSubclassed()) {
+			collIterator = (void *)kalloc_data(collection->iteratorSize(), Z_WAITOK);
+			OSCONTAINER_ACCUMSIZE(collection->iteratorSize());
+			if (!collection->initIterator(collIterator)) {
+				kfree_data(collIterator, collection->iteratorSize());
+				OSCONTAINER_ACCUMSIZE(-((size_t) collection->iteratorSize()));
+				collIterator = NULL;
+				initialized = false;
+				setStorageType(OSCollectionIteratorStorageUnallocated);
+			} else {
+				setStorageType(OSCollectionIteratorStoragePointer);
+				result = collIterator;
+				initialized = true;
+			}
+		} else {
+			bzero(&inlineStorage[0], collection->iteratorSize());
+			if (!collection->initIterator(&inlineStorage[0])) {
+				bzero(&inlineStorage[0], collection->iteratorSize());
+				initialized = false;
+				setStorageType(OSCollectionIteratorStorageUnallocated);
+			} else {
+				setStorageType(OSCollectionIteratorStorageInline);
+				result = &inlineStorage[0];
+				initialized = true;
+			}
+		}
+		break;
+	case OSCollectionIteratorStoragePointer:
+		// already initialized
+		initialized = false;
+		break;
+	case OSCollectionIteratorStorageInline:
+		// already initialized
+		initialized = false;
+		break;
+	default:
+		panic("unexpected storage type %u", storageType);
+	}
+#else
+	if (!collIterator) {
+		collIterator = (void *)kalloc_data(collection->iteratorSize(), Z_WAITOK);
+		OSCONTAINER_ACCUMSIZE(collection->iteratorSize());
+		if (!collection->initIterator(collIterator)) {
+			kfree_data(collIterator, collection->iteratorSize());
+			OSCONTAINER_ACCUMSIZE(-((size_t) collection->iteratorSize()));
+			collIterator = NULL;
+			initialized = false;
+			setStorageType(OSCollectionIteratorStorageUnallocated);
+		} else {
+			setStorageType(OSCollectionIteratorStoragePointer);
+			result = collIterator;
+			initialized = true;
+		}
+	}
+#endif /* __LP64__ */
+
+	if (initialized) {
+		valid = true;
+		initialUpdateStamp = collection->updateStamp;
+	}
+
+	return initialized;
+}
+
+void *
+OSCollectionIterator::getIteratorStorage()
+{
+	void * result = NULL;
+
+#if __LP64__
+	OSCollectionIteratorStorageType storageType = getStorageType();
+
+	switch (storageType) {
+	case OSCollectionIteratorStorageUnallocated:
+		result = NULL;
+		break;
+	case OSCollectionIteratorStoragePointer:
+		result = collIterator;
+		break;
+	case OSCollectionIteratorStorageInline:
+		result = &inlineStorage[0];
+		break;
+	default:
+		panic("unexpected storage type %u", storageType);
+	}
+#else
+	OSCollectionIteratorStorageType storageType __assert_only = getStorageType();
+	assert(storageType == OSCollectionIteratorStoragePointer || storageType == OSCollectionIteratorStorageUnallocated);
+	result = collIterator;
+#endif /* __LP64__ */
+
+	return result;
+}
+
+void
+OSCollectionIterator::freeIteratorStorage()
+{
+#if __LP64__
+	OSCollectionIteratorStorageType storageType = getStorageType();
+
+	switch (storageType) {
+	case OSCollectionIteratorStorageUnallocated:
+		break;
+	case OSCollectionIteratorStoragePointer:
+		kfree_data(collIterator, collection->iteratorSize());
 		OSCONTAINER_ACCUMSIZE(-((size_t) collection->iteratorSize()));
 		collIterator = NULL;
-	}
-
-	if (collection) {
-		collection->release();
-		collection = NULL;
-	}
-
-	super::free();
+		setStorageType(OSCollectionIteratorStorageUnallocated);
+		break;
+	case OSCollectionIteratorStorageInline:
+		bzero(&inlineStorage[0], collection->iteratorSize());
+		setStorageType(OSCollectionIteratorStorageUnallocated);
+		break;
+	default:
+		panic("unexpected storage type %u", storageType);
+	}
+#else
+	if (collIterator != NULL) {
+		assert(getStorageType() == OSCollectionIteratorStoragePointer);
+		kfree_data(collIterator, collection->iteratorSize());
+		OSCONTAINER_ACCUMSIZE(-((size_t) collection->iteratorSize()));
+		collIterator = NULL;
+		setStorageType(OSCollectionIteratorStorageUnallocated);
+	} else {
+		assert(getStorageType() == OSCollectionIteratorStorageUnallocated);
+	}
+#endif /* __LP64__ */
+}
+
+bool
+OSCollectionIterator::isSubclassed()
+{
+	return getMetaClass() != OSCollectionIterator::metaClass;
+}
+
+OSCollectionIteratorStorageType
+OSCollectionIterator::getStorageType()
+{
+#if __LP64__
+	// Storage type is in the most significant 2 bits of collIterator
+	return (OSCollectionIteratorStorageType)((uintptr_t)(collIterator) >> 62);
+#else
+	if (collIterator != NULL) {
+		return OSCollectionIteratorStoragePointer;
+	} else {
+		return OSCollectionIteratorStorageUnallocated;
+	}
+#endif /* __LP64__ */
 }
 
 void
-OSCollectionIterator::reset()
-{
-	valid = false;
-
-	if (!collIterator) {
-		collIterator = (void *)kalloc_container(collection->iteratorSize());
-		OSCONTAINER_ACCUMSIZE(collection->iteratorSize());
-		if (!collIterator) {
-			return;
+OSCollectionIterator::setStorageType(OSCollectionIteratorStorageType storageType)
+{
+#if __LP64__
+	switch (storageType) {
+	case OSCollectionIteratorStorageUnallocated:
+		if (collIterator != NULL) {
+			assert(getStorageType() == OSCollectionIteratorStorageInline);
+			collIterator = NULL;
 		}
-	}
-
-	if (!collection->initIterator(collIterator)) {
-		return;
-	}
-
-	initialUpdateStamp = collection->updateStamp;
-	valid = true;
-}
-
-bool
-OSCollectionIterator::isValid()
-{
-	if (!collIterator) {
-		collIterator = (void *)kalloc_container(collection->iteratorSize());
-		OSCONTAINER_ACCUMSIZE(collection->iteratorSize());
-		if (!collection->initIterator(collIterator)) {
-			return false;
-		}
-		initialUpdateStamp = collection->updateStamp;
-		valid = true;
-	} else if (!valid || collection->updateStamp != initialUpdateStamp) {
-		return false;
-	}
-
-	return true;
+		break;
+	case OSCollectionIteratorStoragePointer:
+		// Should already be set
+		assert(collIterator != NULL);
+		assert(getStorageType() == OSCollectionIteratorStoragePointer);
+		break;
+	case OSCollectionIteratorStorageInline:
+		// Set the two most sigificant bits of collIterator to 10b
+		collIterator = (void *)(((uintptr_t)collIterator & ~0xC000000000000000) | ((uintptr_t)OSCollectionIteratorStorageInline << 62));
+		break;
+	default:
+		panic("unexpected storage type %u", storageType);
+	}
+#else
+	switch (storageType) {
+	case OSCollectionIteratorStorageUnallocated:
+		// Should already be set
+		assert(collIterator == NULL);
+		assert(getStorageType() == OSCollectionIteratorStorageUnallocated);
+		break;
+	case OSCollectionIteratorStoragePointer:
+		// Should already be set
+		assert(collIterator != NULL);
+		assert(getStorageType() == OSCollectionIteratorStoragePointer);
+		break;
+	case OSCollectionIteratorStorageInline:
+		panic("cannot use inline storage on LP32");
+		break;
+	default:
+		panic("unexpected storage type %u", storageType);
+	}
+#endif /* __LP64__ */
 }
 
 OSObject *
@@ -126,11 +317,15 @@
 {
 	OSObject *retObj;
 	bool retVal;
+	void * storage;
 
 	if (!isValid()) {
 		return NULL;
 	}
 
-	retVal = collection->getNextObjectForIterator(collIterator, &retObj);
+	storage = getIteratorStorage();
+	assert(storage != NULL);
+
+	retVal = collection->getNextObjectForIterator(storage, &retObj);
 	return (retVal)? retObj : NULL;
 }