<commit_before>import Foundation
import NIO

extension Document: ExpressibleByArrayLiteral {
    /// Gets all top level values in this Document
    public var values: [Primitive] {
        var values = [Primitive]()
        values.reserveCapacity(32)

        var index = 4

        while index < storage.readableBytes {
            guard
                let typeNum = storage.getInteger(at: index, as: UInt8.self),
                let type = TypeIdentifier(rawValue: typeNum)
            else {
                return values
            }

            index += 1
            guard skipKey(at: &index) else {
                return values
            }

            guard let value = value(forType: type, at: index) else {
                return values
            }

            _ = skipValue(ofType: type, at: &index)

            values.append(value)
        }

        return values
    }
    
    public subscript(index: Int) -> Primitive {
        get {
            var offset = 4
            for _ in 0..<index {
                guard skipKeyValuePair(at: &offset) else {
                    fatalError("Index \(index) out of range")
                }
            }

            guard
                let typeId = storage.getInteger(at: offset, as: UInt8.self),
                let type = TypeIdentifier(rawValue: typeId)
            else {
                fatalError("Index \(index) out of range")
            }

            guard skipKey(at: &offset), let value = self.value(forType: type, at: offset) else {
                fatalError("Index \(index) out of range")
            }

            return value
        }
        set {
            var offset = 4
            for _ in 0..<index {
                guard skipKeyValuePair(at: &offset) else {
                    fatalError("Index \(index) out of range")
                }
            }
            
            let typeOffset = offset
            
            guard
                let typeByte = storage.getInteger(at: typeOffset, as: UInt8.self),
                let type = TypeIdentifier(rawValue: typeByte)
            else {
                fatalError("Index \(index) out of range")
            }
            
            offset += 1
            
            guard skipKey(at: &offset) else {
                fatalError("Index \(index) out of range")
            }
            
            let valueOffset = offset
            
            guard let oldValueLength = valueLength(forType: type, at: offset) else {
                fatalError("Index \(index) out of range")
            }
            
            let valueEnd = valueOffset + oldValueLength
            
            func reserveRoom(_ type: TypeIdentifier, _ newValueLength: Int) {
                storage.setInteger(type.rawValue, at: typeOffset)
                
                let newValueEnd = valueOffset + newValueLength
                let diff = oldValueLength - newValueLength
                storage.moveWriterIndex(to: storage.readableBytes - diff)
                self.usedCapacity -= Int32(diff)
                
                moveBytes(
                    from: valueEnd,
                    to: newValueEnd,
                    length: numericCast(self.usedCapacity) - valueEnd - diff
                )
            }
            
            switch newValue {
            case let double as Double: // 0x01
                reserveRoom(.double, 8)
                storage.setInteger(double.bitPattern, at: valueOffset, endianness: .little)
            case let string as String: // 0x0
                let stringUtf8Length = string.utf8.count
                reserveRoom(.string, 4 + stringUtf8Length + 1)
                storage.setInteger(Int32(stringUtf8Length), at: valueOffset, endianness: .little)
                storage.setString(string, at: valueOffset + 4)
                storage.setInteger(0, at: valueOffset + 4 + stringUtf8Length, endianness: .little, as: UInt8.self)
            case let document as Document: // 0x03 (embedded document) or 0x04 (array)
                reserveRoom(document.isArray ? .array : .document, document.storage.writerIndex)
                storage.setBuffer(document.storage, at: valueOffset)
            case let binary as Binary: // 0x05
                reserveRoom(.binary, 4 + 1 + binary.count)
                storage.setInteger(Int32(binary.count), at: valueOffset, endianness: .little)
                storage.setInteger(binary.subType.identifier, at: valueOffset + 4, endianness: .little)
                storage.setBuffer(binary.storage, at: valueOffset + 4 + 1)<commit_msg>Fix strings overwriting the last byte<commit_after>import Foundation
import NIO

extension Document: ExpressibleByArrayLiteral {
    /// Gets all top level values in this Document
    public var values: [Primitive] {
        var values = [Primitive]()
        values.reserveCapacity(32)

        var index = 4

        while index < storage.readableBytes {
            guard
                let typeNum = storage.getInteger(at: index, as: UInt8.self),
                let type = TypeIdentifier(rawValue: typeNum)
            else {
                return values
            }

            index += 1
            guard skipKey(at: &index) else {
                return values
            }

            guard let value = value(forType: type, at: index) else {
                return values
            }

            _ = skipValue(ofType: type, at: &index)

            values.append(value)
        }

        return values
    }
    
    public subscript(index: Int) -> Primitive {
        get {
            var offset = 4
            for _ in 0..<index {
                guard skipKeyValuePair(at: &offset) else {
                    fatalError("Index \(index) out of range")
                }
            }

            guard
                let typeId = storage.getInteger(at: offset, as: UInt8.self),
                let type = TypeIdentifier(rawValue: typeId)
            else {
                fatalError("Index \(index) out of range")
            }

            guard skipKey(at: &offset), let value = self.value(forType: type, at: offset) else {
                fatalError("Index \(index) out of range")
            }

            return value
        }
        set {
            var offset = 4
            for _ in 0..<index {
                guard skipKeyValuePair(at: &offset) else {
                    fatalError("Index \(index) out of range")
                }
            }
            
            let typeOffset = offset
            
            guard
                let typeByte = storage.getInteger(at: typeOffset, as: UInt8.self),
                let type = TypeIdentifier(rawValue: typeByte)
            else {
                fatalError("Index \(index) out of range")
            }
            
            offset += 1
            
            guard skipKey(at: &offset) else {
                fatalError("Index \(index) out of range")
            }
            
            let valueOffset = offset
            
            guard let oldValueLength = valueLength(forType: type, at: offset) else {
                fatalError("Index \(index) out of range")
            }
            
            let valueEnd = valueOffset + oldValueLength
            
            func reserveRoom(_ type: TypeIdentifier, _ newValueLength: Int) {
                storage.setInteger(type.rawValue, at: typeOffset)
                
                let newValueEnd = valueOffset + newValueLength
                let diff = oldValueLength - newValueLength
                storage.moveWriterIndex(to: storage.readableBytes - diff)
                self.usedCapacity -= Int32(diff)
                
                moveBytes(
                    from: valueEnd,
                    to: newValueEnd,
                    length: numericCast(self.usedCapacity) - valueEnd - diff
                )
            }
            
            switch newValue {
            case let double as Double: // 0x01
                reserveRoom(.double, 8)
                storage.setInteger(double.bitPattern, at: valueOffset, endianness: .little)
            case let string as String: // 0x0
                let stringUtf8Length = string.utf8.count
                reserveRoom(.string, 4 + stringUtf8Length + 1)
                storage.setInteger(Int32(stringUtf8Length + 1), at: valueOffset, endianness: .little)
                storage.setString(string, at: valueOffset + 4)
                storage.setInteger(0, at: valueOffset + 4 + stringUtf8Length, endianness: .little, as: UInt8.self)
            case let document as Document: // 0x03 (embedded document) or 0x04 (array)
                reserveRoom(document.isArray ? .array : .document, document.storage.writerIndex)
                storage.setBuffer(document.storage, at: valueOffset)
            case let binary as Binary: // 0x05
                reserveRoom(.binary, 4 + 1 + binary.count)
                storage.setInteger(Int32(binary.count), at: valueOffset, endianness: .little)
                storage.setInteger(binary.subType.identifier, at: valueOffset + 4, endianness: .little)
                storage.setBuffer(binary.storage, at: valueOffset + 4 + 1)
{'fraction_non_alphanumeric': 0.08155487804878049, 'fraction_numerical': 0.010779616724738676, 'mean_word_length': 1.3919270833333333, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '120826', 'n_tokens_mistral': 2414, 'n_tokens_neox': 2233, 'n_words': 625}