diff --git a/HTMLKit/HTMLDocument.m b/HTMLKit/HTMLDocument.m
index c910f50..1837124 100644
--- a/HTMLKit/HTMLDocument.m
+++ b/HTMLKit/HTMLDocument.m
@@ -18,6 +18,7 @@
@interface HTMLDocument ()
{
HTMLDocument *_inertTemplateDocument;
+ NSMutableArray *_nodeIterators;
}
@property (nonatomic, assign) HTMLDocumentReadyState readyState;
@end
@@ -37,6 +38,7 @@
self = [super initWithName:@"#document" type:HTMLNodeDocument];
if (self) {
_readyState = HTMLDocumentLoading;
+ _nodeIterators = [NSMutableArray new];
}
return self;
}
@@ -64,6 +66,18 @@
}
}
+#pragma mark - Node Iterators
+
+- (void)attachNodeIterator:(HTMLNodeIterator *)iterator
+{
+ [_nodeIterators addObject:iterator];
+}
+
+- (void)detachNodeIterator:(HTMLNodeIterator *)iterator
+{
+ [_nodeIterators removeObject:iterator];
+}
+
#pragma mark - Mutation Algorithms
- (HTMLNode *)adoptNode:(HTMLNode *)node
diff --git a/HTMLKit/HTMLNodeIterator.m b/HTMLKit/HTMLNodeIterator.m
index 370155e..7b8314e 100644
--- a/HTMLKit/HTMLNodeIterator.m
+++ b/HTMLKit/HTMLNodeIterator.m
@@ -7,6 +7,7 @@
//
#import "HTMLNodeIterator.h"
+#import "HTMLDocument.h"
#import "HTMLNode.h"
typedef NS_ENUM(short, TraverseDirection)
@@ -15,6 +16,11 @@ typedef NS_ENUM(short, TraverseDirection)
TraverseDirectionPrevious
};
+@interface HTMLDocument (Private)
+- (void)attachNodeIterator:(HTMLNodeIterator *)iterator;
+- (void)detachNodeIterator:(HTMLNodeIterator *)iterator;
+@end
+
@interface HTMLNodeIterator ()
{
HTMLNode *_root;
@@ -23,6 +29,8 @@ typedef NS_ENUM(short, TraverseDirection)
@implementation HTMLNodeIterator
+#pragma mark - Lifecycle
+
- (instancetype)initWithNode:(HTMLNode *)node
{
return [self initWithNode:node filter:nil];
@@ -45,10 +53,18 @@ typedef NS_ENUM(short, TraverseDirection)
_whatToShow = showOptions;
_referenceNode = _root;
_pointerBeforeReferenceNode = YES;
+ [_root.ownerDocument attachNodeIterator:self];
}
return self;
}
+- (void)dealloc
+{
+ [_root.ownerDocument detachNodeIterator:self];
+}
+
+#pragma mark - Traversal
+
- (HTMLNode *)traverseInDirection:(TraverseDirection)direction
{
HTMLNode *node = self.referenceNode;
@@ -129,11 +145,6 @@ NS_INLINE HTMLNodeFilterValue FilterNode(HTMLNodeIterator *iterator, HTMLNode *n
return [iterator.filter acceptNode:node];
}
-- (id)nextObject
-{
- return self.nextNode;
-}
-
- (HTMLNode *)nextNode
{
return [self traverseInDirection:TraverseDirectionNext];
@@ -144,4 +155,11 @@ NS_INLINE HTMLNodeFilterValue FilterNode(HTMLNodeIterator *iterator, HTMLNode *n
return [self traverseInDirection:TraverseDirectionPrevious];
}
+#pragma mark - NSEnumerator
+
+- (id)nextObject
+{
+ return self.nextNode;
+}
+
@end