YapRTreeTest.swift

//
//  YapRTreeTest.swift
//  YapDatabaseRTreeTest
//
//  Created by Maël Primet on 06/19/15.
//  Copyright (c) 2015 Snips. All rights reserved.
//

import Foundation

import YapDatabaseFramework

@objc
public class YapDB {
    public static let sharedInstance: YapDB = YapDB()

    public let db : YapDatabase
    public let writeConnection: YapDatabaseConnection
    public let readConnection: YapDatabaseConnection
    
    public init() {
        var paths = NSSearchPathForDirectoriesInDomains(NSSearchPathDirectory.DocumentDirectory, NSSearchPathDomainMask.UserDomainMask, true)
        let dbFilename = paths[0].stringByAppendingPathComponent("database.ydb")
        NSLog("Initialized database in path : %@", dbFilename)
        self.db = YapDatabase(path: dbFilename)

        self.writeConnection = self.db.newConnection()
        self.readConnection = self.db.newConnection()
    }
}

class LonlatObject: NSObject, NSCoding {
    let name: String
    let lon: Double
    let lat: Double
    
    init(name: String, lon: Double, lat: Double) {
        self.name = name
        self.lon = lon
        self.lat = lat
    }
    
    required init(coder:NSCoder) {
        self.name = coder.decodeObjectForKey("name") as! String
        self.lon = coder.decodeDoubleForKey("lon")
        self.lat = coder.decodeDoubleForKey("lat")
    }
    
    
    /* NSCoding */
    func encodeWithCoder(coder: NSCoder) {
        coder.encodeObject(name, forKey: "name")
        coder.encodeDouble(lon, forKey: "lon")
        coder.encodeDouble(lat, forKey: "lat")
    }
}

class RTreeTest {
    static let indexName = "rTreeIndex"
    
    static func run() {
        RTreeTest.createIndex()
        RTreeTest.writeObjects()
        let items = RTreeTest.queryObjectsInBounds(minLon: 8, maxLon: 2, minLat: 4, maxLat: 8)
        println(items.map({[$0.name, $0.lon, $0.lat]}))
    }
    
    private static func createIndex() {
        let yapDb = YapDB.sharedInstance

        let setup = YapDatabaseRTreeIndexSetup()
        setup.setColumns(["minLon", "maxLon", "minLat", "maxLat"])
        let handler = YapDatabaseRTreeIndexHandler.withObjectBlock({ (d: NSMutableDictionary!, collection: String!, key: String!, object: AnyObject!) -> Void in
            if let item = object as? LonlatObject {
                d["minLon"] = item.lon
                d["maxLon"] = item.lon
                d["minLat"] = item.lat
                d["maxLat"] = item.lat
            }
        })

        let ext = YapDatabaseRTreeIndex(setup: setup, handler: handler, versionTag: "1.0")
        yapDb.db.registerExtension(ext, withName: RTreeTest.indexName)
        
        self.nukeDatabase()
    }
    
    private static func nukeDatabase() {
        let yapDb = YapDB.sharedInstance
        
        var items = []
        yapDb.readConnection.readWriteWithBlock({ (transaction: YapDatabaseReadWriteTransaction) in
            transaction.removeAllObjectsInCollection("lonlatObjects")
        })
    }
    
    private static func writeObjects() {
        let yapDb = YapDB.sharedInstance

        yapDb.writeConnection.readWriteWithBlock({ (transaction: YapDatabaseReadWriteTransaction) in
            for i in (0..<10) {
                for j in (0..<10) {
                    let item = LonlatObject(name: "grid-\(i)-\(j)", lon: Double(i), lat: Double(j))
                    transaction.setObject(item, forKey: item.name, inCollection: "lonlatObjects")
                }
            }
        })
    }

    private static func queryObjectsInBounds(#minLon: Double, maxLon: Double, minLat: Double, maxLat: Double) -> [LonlatObject] {
        let yapDb = YapDB.sharedInstance

        var arguments: [CVarArgType] = [minLon, maxLon, minLat, maxLat].map({NSNumber(double: $0)}) // must be var to get a pointer, must be a NSObject to be formatted
        let queryString = "WHERE minLon >= ? AND maxLon <= ? AND minLat >= ? AND maxLat <= ?"
        let query = withVaList(arguments, { (pointer: CVaListPointer) -> YapDatabaseQuery in
            return YapDatabaseQuery(format: queryString, arguments: pointer)
        })
        var items: [LonlatObject] = []

        yapDb.readConnection.readWithBlock({ (transaction: YapDatabaseReadTransaction) in
            if let index = transaction.ext("rTreeIndex") as? YapDatabaseRTreeIndexTransaction {
                let enumerateBlock: (String!, String!, AnyObject!, UnsafeMutablePointer<ObjCBool>) -> Void = {
                    (collection: String!, key: String!, object: AnyObject!, stop: UnsafeMutablePointer<ObjCBool>) in
                    if let item = object as? LonlatObject {
                        items.append(item)
                    }
                }
                index.enumerateKeysAndObjectsMatchingQuery(query, usingBlock: enumerateBlock)
            }
        })
        
        return items
    }
    
}