Added category dependency

//
//  XPathQuery.m
//
//  Created by Matt Gallagher on 4/08/08.
//  Modified by Matt Bauman on 12/01/09 to ignore attributes and simplify the
//  created data structure (depends upon the NSDictionary Category method
//  - (void)addObject:forKey: below).
//  Copyright 2008 Matt Gallagher. All rights reserved.
//
//  Permission is given to use this source code file, free of charge, in any
//  project, commercial or otherwise, entirely at your risk, with the condition
//  that any redistribution (in part or whole) of source code must retain
//  this copyright and permission notice. Attribution in compiled projects is
//  appreciated but not required.
//

#import "XPathQuery.h"
#import "NSMutableDictionary+MBAddObject.h"

#import <libxml/tree.h>
#import <libxml/parser.h>
#import <libxml/HTMLparser.h>
#import <libxml/xpath.h>
#import <libxml/xpathInternals.h>

id ElementForNode(xmlNodePtr currentNode, NSMutableDictionary *currentDictionary)
{
    if (!currentNode->name) return nil;
    NSString *currentNodeName = [NSString stringWithCString:(const char *)currentNode->name 
                                                   encoding:NSUTF8StringEncoding];

    if (currentNode->children && currentNode->children == currentNode->last &&
        (currentNode->children->type == XML_TEXT_NODE || currentNode->children->type == XML_CDATA_SECTION_NODE) ) {
        /* We have a simple element; add it to the existing dictionary */
        NSString *currentNodeText = [NSString stringWithCString:(const char *)currentNode->children->content 
                                                       encoding:NSUTF8StringEncoding];
        currentNodeText = [currentNodeText stringByTrimmingCharactersInSet:[NSCharacterSet whitespaceAndNewlineCharacterSet]];
        
        if ([currentNodeText length] > 0) {
            /* Note that currentDictionary may be null; in that case addObject:forKey:
             * falls flat and the returned object should suffice. */
            [currentDictionary addObject:currentNodeText forKey:currentNodeName];
            return currentNodeText;
        }
        return nil;
    }
    
    
    if (currentNode->type == XML_ELEMENT_NODE) {
        /* create a new dictionary for this node */
        NSMutableDictionary *nodeDictionary = [NSMutableDictionary dictionary];
        [currentDictionary addObject:nodeDictionary forKey:currentNodeName];
        
        /* And go through its children */
        xmlNodePtr childNode = currentNode->children;
        do {
            ElementForNode(childNode, nodeDictionary);
        } while (childNode = childNode->next);
        
        return nodeDictionary;
    }
        
    return nil;
}

NSArray *PerformXPathQuery(xmlDocPtr doc, NSString *query) {
    xmlXPathContextPtr xpathCtx; 
    xmlXPathObjectPtr xpathObj; 

    /* Create xpath evaluation context */
    xpathCtx = xmlXPathNewContext(doc);
    if(xpathCtx == NULL) {
        NSLog(@"Unable to create XPath context.");
        return nil;
    }
    
    /* Evaluate xpath expression */
    xpathObj = xmlXPathEvalExpression((xmlChar *)[query cStringUsingEncoding:NSUTF8StringEncoding], xpathCtx);
    if(xpathObj == NULL) {
        NSLog(@"Unable to evaluate XPath.");
        return nil;
    }
    
    xmlNodeSetPtr nodes = xpathObj->nodesetval;
    if (!nodes) {
        NSLog(@"Nodes was nil.");
        return nil;
    }
    
    NSMutableArray *resultNodes = [NSMutableArray array];
    for (NSInteger i = 0; i < nodes->nodeNr; i++) {
        id nodeElement = ElementForNode(nodes->nodeTab[i], nil);
        if (nodeElement) {
            [resultNodes addObject:nodeElement];
        }
    }

    /* Cleanup */
    xmlXPathFreeObject(xpathObj);
    xmlXPathFreeContext(xpathCtx); 
    
    return resultNodes;
}

NSArray *PerformHTMLXPathQuery(NSData *document, NSString *query) {
    xmlDocPtr doc;

    /* Load XML document */
    doc = htmlReadMemory([document bytes], [document length], "", NULL, HTML_PARSE_NOWARNING | HTML_PARSE_NOERROR);
    
    if (doc == NULL) {
        NSLog(@"Unable to parse.");
        return nil;
    }
    
    NSArray *result = PerformXPathQuery(doc, query);
    xmlFreeDoc(doc); 
    
    return result;
}

NSArray *PerformXMLXPathQuery(NSData *document, NSString *query) {
    xmlDocPtr doc;
    
    /* Load XML document */
    doc = xmlReadMemory([document bytes], [document length], "", NULL, XML_PARSE_RECOVER);
    
    if (doc == NULL) {
        NSLog(@"Unable to parse.");
        return nil;
    }
    
    NSArray *result = PerformXPathQuery(doc, query);
    xmlFreeDoc(doc); 
    
    return result;
}

// --
//
//  NSMutableDictionary+MBAddObject
//
//  Created by Matt Bauman on 12/1/09.
//  Copyright 2008 Matt Bauman. All rights reserved.
//
//  Permission is given to use this source code file, free of charge, in any
//  project, commercial or otherwise, entirely at your risk, with the condition
//  that any redistribution (in part or whole) of source code must retain
//  this copyright and permission notice. Attribution in compiled projects is
//  appreciated but not required.

#import <Cocoa/Cocoa.h>

@interface NSMutableDictionary (MBAddObject)

/* Adds the object for a given key.  If the key already exists, it adds the
 * object to an array at the given key location. */
- (void)addObject:(id)anObject forKey:(id)aKey;

@end

@implementation NSMutableDictionary (MBAddObject)

- (void)addObject:(id)anObject forKey:(id)aKey {
    id existingObject = [self objectForKey:aKey];
    if (!existingObject) {
        [self setObject:anObject forKey:aKey];
    } else {
        /* We already have a key for this element. If it's not an array, create
         * one, and add the existing element back into the new array. */
        if (![existingObject isKindOfClass:[NSMutableArray class]]) {
            NSMutableArray *newArray = [NSMutableArray arrayWithObject:existingObject];
            [self setObject:newArray forKey:aKey];
            existingObject = newArray;
        }
        [(NSMutableArray *)existingObject addObject:anObject];
    }
    
}

@end