http://bsonspec.org

BSON.swift

import Foundation

public struct BSONError : LocalizedError, CustomStringConvertible {
    public var message: String
    public init(_ message: String) { self.message = message }
    public var description: String { return message }
    public var errorDescription: String? { description }
}

public enum BSONBinarySubtypes : UInt8 {
    case generic = 0x00
    case function = 0x01
    case oldBinary = 0x02
    case oldUUID = 0x03
    case uuid = 0x04
    case md5 = 0x05
    case encryptedBSONValue = 0x06
    case userDefined = 0x80
}

public final class BSONDocumentWriter {
    private let writer: BSONBaseWriter
    private let beginPosition: UInt64
    
    public convenience init(fileHandle: FileHandle) {
        let writer = BSONBaseWriter(fileHandle: fileHandle)
        self.init(writer: writer)
    }
    
    public init(writer: BSONBaseWriter) {
        let position = writer.position()
        self.writer = writer
        self.beginPosition = position
        
        writer.writeRawInt32(0)
    }
    
    public func end() {
        writer.writeRawUInt8(0)
        let endPosition = writer.position()
        let size = endPosition - beginPosition
        writer.seek(to: beginPosition)
        writer.writeRawInt32(Int32(size))
        writer.seek(to: endPosition)
    }
    
    public func writeDouble(name: String, value: Double) {
        writer.writeRawUInt8(0x01)
        writer.writeCString(name)
        writer.writeRawDouble(value)
    }
    
    public func writeString(name: String, value: String) {
        writer.writeRawUInt8(0x02)
        writer.writeCString(name)
        writer.writeString(value)
    }
    
    public func beginDocument(name: String) -> BSONDocumentWriter {
        writer.writeRawUInt8(0x03)
        writer.writeCString(name)
        return BSONDocumentWriter(writer: writer)
    }
    
    public func writeDocument(name: String, _ f: (BSONDocumentWriter) throws -> Void) rethrows {
        let dw = beginDocument(name: name)
        try f(dw)
        dw.end()
    }
    
    public func beginArray(name: String) -> BSONArrayWriter {
        writer.writeRawUInt8(0x04)
        writer.writeCString(name)
        return BSONArrayWriter(writer: writer)
    }
    
    public func writeArray(name: String, _ f: (BSONArrayWriter) throws -> Void) rethrows {
        let aw = beginArray(name: name)
        try f(aw)
        aw.end()
    }
    
    public func writeBinary(name: String, value: Data) {
        writeBinary(name: name, subtype: .generic, value: value)
    }
    
    public func writeBinary(name: String, subtype: BSONBinarySubtypes, value: Data) {
        writeBinary(name: name, subtype: subtype.rawValue, value: value)
    }
    
    public func writeBinary(name: String, subtype: UInt8, value: Data) {
        writer.writeRawUInt8(0x05)
        writer.writeCString(name)
        writer.writeRawInt32(Int32(value.count))
        writer.writeRawUInt8(subtype)
        writer.writeRawData(value)
    }
    
    public func writeBool(name: String, value: Bool) {
        writer.writeRawUInt8(0x08)
        writer.writeCString(name)
        writer.writeRawUInt8(value ? 0x00 : 0x01)
    }
    
    public func writeInt32(name: String, value: Int32) {
        writer.writeRawUInt8(0x10)
        writer.writeCString(name)
        writer.writeRawInt32(value)
    }
    
    public func writeInt64(name: String, value: Int64) {
        writer.writeRawUInt8(0x12)
        writer.writeCString(name)
        writer.writeRawInt64(value)
    }
}

public final class BSONArrayWriter {
    private let writer: BSONDocumentWriter
    private var index: Int
    
    private var indexString: String { return "\(index)" }
    
    public convenience init(fileHandle: FileHandle) {
        let writer = BSONBaseWriter(fileHandle: fileHandle)
        self.init(writer: writer)
    }
    
    public init(writer: BSONBaseWriter) {
        let dw = BSONDocumentWriter(writer: writer)
        self.writer = dw
        self.index = 0
    }

    public func end() {
        writer.end()
    }
    
    public func writeDouble(_ value: Double) {
        writer.writeDouble(name: indexString, value: value)
        index += 1
    }
    
    public func writeString(_ value: String) {
        writer.writeString(name: indexString, value: value)
        index += 1
    }

    public func beginDocument() -> BSONDocumentWriter {
        let dw = writer.beginDocument(name: indexString)
        index += 1
        return dw
    }
    
    public func writeDocument(_ f: (BSONDocumentWriter) throws -> Void) rethrows {
        let dw = beginDocument()
        try f(dw)
        dw.end()
    }
    
    public func beginArray() -> BSONArrayWriter {
        let aw = writer.beginArray(name: indexString)
        index += 1
        return aw
    }
    
    public func writeArray(_ f: (BSONArrayWriter) throws -> Void) rethrows {
        let aw = beginArray()
        try f(aw)
        aw.end()
    }
    
    public func writeBinary(_ value: Data) {
        writeBinary(subtype: .generic, value: value)
    }
    
    public func writeBinary(subtype: BSONBinarySubtypes, value: Data) {
        writeBinary(subtype: subtype.rawValue, value: value)
    }
    
    public func writeBinary(subtype: UInt8, value: Data) {
        writer.writeBinary(name: indexString, subtype: subtype, value: value)
        index += 1
    }
    
    public func writeBool(_ value: Bool) {
        writer.writeBool(name: indexString, value: value)
        index += 1
    }
    
    public func writeInt32(_ value: Int32) {
        writer.writeInt32(name: indexString, value: value)
        index += 1
    }
    
    public func writeInt64(_ value: Int64) {
        writer.writeInt64(name: indexString, value: value)
        index += 1
    }
}

public final class BSONBaseWriter {
    private let fileHandle: FileHandle
    
    public init(fileHandle: FileHandle) {
        self.fileHandle = fileHandle
    }
    
    public func position() -> UInt64 {
        return fileHandle.offsetInFile
    }
    
    public func seek(to position: UInt64) {
        fileHandle.seek(toFileOffset: position)
    }
    
    public func writeRawData(_ data: Data) {
        fileHandle.write(data)
    }
    
    public func writeRawUInt8(_ value: UInt8) {
        writeRawPrimitive(value) }
    
    public func writeRawInt32(_ value: Int32) {
        writeRawPrimitive(value) }
    
    public func writeRawInt64(_ value: Int64) {
        writeRawPrimitive(value) }
    
    public func writeRawDouble(_ value: Double) {
        writeRawPrimitive(value) }
    
    public func writeRawPrimitive<T>(_ value: T) {
        var value = value
        let pointer = UnsafeMutablePointer(&value)
        let data = Data(bytesNoCopy: pointer,
                        count: MemoryLayout<T>.size,
                        deallocator: .none)
        writeRawData(data)
    }
    
    public func writeCString(_ value: String) {
        var value = value
        value.withUTF8 { (buffer) in
            let pointer = UnsafeMutablePointer(mutating: buffer.baseAddress!)
            let data = Data(bytesNoCopy: pointer,
                            count: buffer.count,
                            deallocator: .none)
            writeRawData(data)
            writeRawUInt8(0)
        }
    }
    
    public func writeString(_ value: String) {
        var value = value
        value.withUTF8 { (buffer) in
            writeRawInt32(Int32(buffer.count + 1))
            let pointer = UnsafeMutablePointer(mutating: buffer.baseAddress!)
            let data = Data(bytesNoCopy: pointer,
                            count: buffer.count,
                            deallocator: .none)
            writeRawData(data)
            writeRawUInt8(0)
        }
    }
}

public enum BSONReadResult {
    public struct Document {
        public var value: [(name: String, value: BSONReadResult)]
        
        public init(_ value: [(name: String, value: BSONReadResult)]) {
            self.value = value
        }
        
        public func getIfPresent(name: String) -> BSONReadResult?
        {
            return value.first { $0.name == name }?.value
        }
        
        public func get(name: String) throws -> BSONReadResult {
            guard let x = getIfPresent(name: name) else {
                throw BSONError("not exists: \(name)")
            }
            return x
        }

    }
    public struct Array {
        public var value: [BSONReadResult]
        public init(_ value: [BSONReadResult]) {
            self.value = value
        }
        
        public func getIfPresent(at index: Int) -> BSONReadResult? {
            guard 0 <= index, index < value.count else {
                return nil
            }
            return value[index]
        }
        
        public func get(at index: Int) throws -> BSONReadResult {
            guard let x = getIfPresent(at: index) else {
                throw BSONError("invalid index: \(index), count=\(value.count)")
            }
            return x
        }
    }
    
    case double(Double)
    case stringPosition(UInt64, size: Int)
    case string(String)
    case documentPosition(UInt64, size: Int)
    case document(Document)
    case arrayPosition(UInt64, size: Int)
    case array(Array)
    case binaryPosition(UInt64, subtype: UInt8, size: Int)
    case binary(subtype: UInt8, data: Data)
    case boolean(Bool)
    case int32(Int32)
    case int64(Int64)
    
    public func asDouble() throws -> Double {
        guard case .double(let x) = self else {
            throw BSONError("not double")
        }
        return x
    }
    
    public func asString() throws -> String {
        guard case .string(let x) = self else {
            throw BSONError("not string")
        }
        return x
    }
    
    public func asDocument() throws -> Document {
        guard case .document(let x) = self else {
            throw BSONError("not document")
        }
        return x
    }
    
    public func asArray() throws -> Array {
        guard case .array(let x) = self else {
            throw BSONError("not array")
        }
        return x
    }
    
    public func asBinary() throws -> (subtype: UInt8, data: Data) {
        guard case .binary(subtype: let subtype, data: let data) = self else {
            throw BSONError("not binary")
        }
        return (subtype: subtype, data: data)
    }
    
    public func asBoolean() throws -> Bool {
        guard case .boolean(let x) = self else {
            throw BSONError("not boolean")
        }
        return x
    }
    
    public func asInt32() throws -> Int32 {
        guard case .int32(let x) = self else {
            throw BSONError("not int32")
        }
        return x
    }
    
    public func asInt64() throws -> Int64 {
        guard case .int64(let x) = self else {
            throw BSONError("not int64")
        }
        return x
    }
}

public enum BSONCompoundReadMode {
    case skip
    case read
}

public final class BSONDocumentReader {
    private let reader: BSONBaseReader
    private let beginPosition: UInt64
    private let size: Int
    
    public convenience init(fileHandle: FileHandle) throws
    {
        let reader = BSONBaseReader(fileHandle: fileHandle)
        try self.init(reader: reader)
    }
    
    public init(reader: BSONBaseReader) throws
    {
        self.reader = reader
        self.beginPosition = reader.position()
        self.size = Int(try reader.readRawInt32())
    }
    
    public func readItem(compound: BSONCompoundReadMode) throws -> (name: String, value: BSONReadResult)?
    {
        let tag = try reader.readRawUInt8()
        if tag == 0x00 {
            // document has terminator zero
            return nil
        }
        
        let name = try reader.readCString()
        let position = reader.position()
        switch tag {
        case 0x01:
            let value = try reader.readRawDouble()
            return (name: name, value: .double(value))
        case 0x02:
            switch compound {
            case .skip:
                let size = Int(try reader.readRawInt32())
                return (name: name, value: .stringPosition(position, size: size))
            case .read:
                let value = try reader.readString()
                return (name: name, value: .string(value))
            }
        case 0x03:
            switch compound {
            case .skip:
                let size = Int(try reader.readRawInt32())
                return (name: name, value: .documentPosition(position, size: size))
            case .read:
                let value = try reader.readDocument(compound: .read)
                return (name: name, value: .document(value))
            }
        case 0x04:
            switch compound {
            case .skip:
                let size = Int(try reader.readRawInt32())
                return (name: name, value: .arrayPosition(position, size: size))
            case .read:
                let value = try reader.readArray(compound: .read)
                return (name: name, value: .array(value))
            }
        case 0x05:
            switch compound {
            case .skip:
                let size = Int(try reader.readRawInt32())
                let subtype = try reader.readRawUInt8()
                return (name: name, value: .binaryPosition(position, subtype: subtype, size: size))
            case .read:
                let value = try reader.readBinary()
                return (name: name, value: .binary(subtype: value.subtype, data: value.data))
            }
        case 0x08:
            let value = try reader.readRawUInt8()
            return (name: name, value: .boolean(value != 0 ? true : false))
        case 0x10:
            let value = try reader.readRawInt32()
            return (name: name, value: .int32(value))
        case 0x12:
            let value = try reader.readRawInt64()
            return (name: name, value: .int64(value))
        default:
            let tagStr = String(format: "%02x", tag)
            throw BSONError("unknown tag: \(tagStr)")
        }
    }
    
    public func read(compound: BSONCompoundReadMode) throws -> BSONReadResult.Document {
        var result: BSONReadResult.Document = .init([])
        while true {
            guard let item = try readItem(compound: compound) else {
                break
            }
            result.value.append(item)
        }
        return result
    }
    
    public func seekToEnd() {
        reader.seek(to: beginPosition + UInt64(size))
    }
    
    public func readSkipped(_ item: BSONReadResult,
                            compound: BSONCompoundReadMode) throws -> BSONReadResult {
        return try reader.readSkipped(item, compound: compound)
    }
}

public final class BSONArrayReader {
    private let reader: BSONDocumentReader
    private var index: Int
    
    public convenience init(fileHandle: FileHandle) throws
    {
        let reader = BSONBaseReader(fileHandle: fileHandle)
        try self.init(reader: reader)
    }
    
    public init(reader: BSONBaseReader) throws
    {
        self.reader = try BSONDocumentReader(reader: reader)
        self.index = 0
    }
    
    public func readItem(compound: BSONCompoundReadMode) throws -> BSONReadResult?
    {
        guard let item = try reader.readItem(compound: compound) else {
            return nil
        }
        
        let indexString = "\(index)"
        guard indexString == item.name else {
            throw BSONError("invalid array index: expected=\(indexString), actual=\(item.name)")
        }
        
        index += 1

        return item.value
    }
    
    public func read(compound: BSONCompoundReadMode) throws -> BSONReadResult.Array {
        var result: BSONReadResult.Array = .init([])
        while true {
            guard let item = try readItem(compound: compound) else {
                break
            }
            result.value.append(item)
        }
        return result
    }
    
    public func seekToEnd() {
        reader.seekToEnd()
    }
    
    public func readSkipped(_ item: BSONReadResult,
                            compound: BSONCompoundReadMode) throws -> BSONReadResult {
        return try reader.readSkipped(item, compound: compound)
    }
}

public final class BSONBaseReader {
    private let fileHandle: FileHandle
    
    public init(fileHandle: FileHandle) {
        self.fileHandle = fileHandle
    }
    
    public func position() -> UInt64 {
        return fileHandle.offsetInFile
    }
    
    public func seek(to position: UInt64) {
        fileHandle.seek(toFileOffset: position)
    }
    
    public func readRawUInt8() throws -> UInt8 {
        return try readRawPrimitive(type: UInt8.self)
    }
    
    public func readRawInt32() throws -> Int32 {
        return try readRawPrimitive(type: Int32.self)
    }
    
    public func readRawInt64() throws -> Int64 {
        return try readRawPrimitive(type: Int64.self)
    }
    
    public func readRawDouble() throws -> Double {
        return try readRawPrimitive(type: Double.self)
    }

    public func readRawPrimitive<T>(type: T.Type) throws -> T {
        let data = try readRawData(size: MemoryLayout<T>.size)
        let value = data.withUnsafeBytes { (buffer) in
            buffer.bindMemory(to: T.self).baseAddress!.pointee
        }
        return value
    }
    
    public func readRawData(size: Int) throws -> Data {
        let data = fileHandle.readData(ofLength: size)
        if data.count < size {
            throw BSONError("reached to end")
        }
        return data
    }
    
    public func readCString() throws -> String {
        let beginPosition = fileHandle.offsetInFile
        var data: Data = Data()
        while true {
            let chunk: Data = fileHandle.readData(ofLength: 128)
            if chunk.count == 0 {
                throw BSONError("reached to end")
            }
            // search null terminator
            guard let index = (chunk.firstIndex(of: 0)) else {
                data.append(chunk)
                continue
            }
            // trim null
            data.append(chunk[..<index])
            break
        }
        // +1 means null
        fileHandle.seek(toFileOffset: beginPosition + UInt64(data.count) + 1)
        guard let str = String(data: data, encoding: .utf8) else {
            throw BSONError("invalid UTF-8")
        }
        return str
    }
    
    public func readString() throws -> String {
        let size = Int(try readRawInt32())
        guard size >= 1 else {
            throw BSONError("invalid strig size: \(size)")
        }
        // trim null
        let data = try readRawData(size: size - 1)
        // skip null
        _ = try readRawUInt8()
        guard let str = String(data: data, encoding: .utf8) else {
            throw BSONError("invalid UTF-8")
        }
        return str
    }
    
    public func readBinary() throws -> (subtype: UInt8, data: Data) {
        let size = Int(try readRawInt32())
        let subtype = try readRawUInt8()
        let data = try readRawData(size: size)
        return (subtype: subtype, data: data)
    }
    
    public func readDocument(compound: BSONCompoundReadMode) throws -> BSONReadResult.Document
    {
        let dr = try BSONDocumentReader(reader: self)
        return try dr.read(compound: .read)
    }
    
    public func readArray(compound: BSONCompoundReadMode) throws -> BSONReadResult.Array {
        let ar = try BSONArrayReader(reader: self)
        return try ar.read(compound: .read)
    }
    
    public func readSkipped(_ item: BSONReadResult,
                            compound: BSONCompoundReadMode) throws -> BSONReadResult {
        switch item {
        case .stringPosition(let position, size: _):
            seek(to: position)
            let value = try readString()
            return .string(value)
        case .documentPosition(let position, size: _):
            seek(to: position)
            let value = try readDocument(compound: compound)
            return .document(value)
        case .arrayPosition(let position, size: _):
            seek(to: position)
            let value = try readArray(compound: compound)
            return .array(value)
        case .binaryPosition(let position, subtype: _, size: _):
            seek(to: position)
            let value = try readBinary()
            return .binary(subtype: value.subtype, data: value.data)
        case .double,
             .string,
             .document,
             .array,
             .binary,
             .boolean,
             .int32,
             .int64:
            throw BSONError("invalid kind: \(item)")
        }
    }
}

public protocol BSONDocumentWritable {
    func write(to writer: BSONDocumentWriter)
}

extension BSONDocumentWriter {
    public func write<X: BSONDocumentWritable>(name: String, value: X) {
        writeDocument(name: name) { (w) in
            value.write(to: w)
        }
    }
}

extension BSONArrayWriter {
    public func write<X: BSONDocumentWritable>(_ value: X) {
        writeDocument { (w) in
            value.write(to: w)
        }
    }
}

public protocol BSONDocumentReadable {
    init(from doc: BSONReadResult.Document) throws
}

extension BSONReadResult {
    public func `as`<X: BSONDocumentReadable>(type: X.Type) throws -> X {
        let doc = try self.asDocument()
        return try X(from: doc)
    }
}