<commit_before>import Foundation
import NIO
import MongoKitten

final class FileWriter {
    
    static let allocator = ByteBufferAllocator()
    static let encoder = BSONEncoder()
    
    let fs: GridFS
    let fileId: Primitive
    let chunkSize: Int
    var buffer: ByteBuffer
    var nextChunkNumber = 0
    var length = 0
    var finalized = false
    
    internal init(fs: GridFS, fileId: Primitive, chunkSize: Int, buffer: ByteBuffer? = nil) {
        self.fs = fs
        self.fileId = fileId
        self.chunkSize = chunkSize
        self.buffer = buffer ?? FileWriter.allocator.buffer(capacity: chunkSize)
    }
    
    public func write(data: ByteBuffer) -> EventLoopFuture<Void> {
        assert(self.finalized == false, "Writing to a finalized writer is an error")
        
        self.length += data.readableBytes
        var source = data
        buffer.write(buffer: &source)
        
        guard buffer.readableBytes >= chunkSize else {
            return fs.eventLoop.newSucceededFuture(result: ())
        }
        
        return self.flush()
    }
    
    public func finalize() -> EventLoopFuture<Void> {
        assert(self.finalized == false, "Finalizing a finalized writer is an error")
        
        self.finalized = true
        
        return self.flush(finalize: true)
            .then {
                let file = File(id: self.fileId,
                                length: self.length,
                                chunkSize: self.chunkSize,
                                metadata: nil, // TODO
                                filename: nil, // TODO
                                fs: self.fs)
                
                do {
                    let encoded = try FileWriter.encoder.encode(file)
                    
                    return self.fs.filesCollection.insert(encoded).map { _ in }
                } catch {
                    return self.fs.eventLoop.newFailedFuture(error: error)
                }
        }
    }
    
    private func flush(finalize: Bool = false) -> EventLoopFuture<Void> {
        guard buffer.readableBytes > 0, finalize || buffer.readableBytes >= chunkSize else {
            return fs.eventLoop.newSucceededFuture(result: ())
        }
        
        guard let slice = buffer.readSlice(length: buffer.readableBytes >= chunkSize ? chunkSize : buffer.readableBytes) else {
            // TODO: Replace with error future
            fatalError()
        }
        
        do {
            let chunk = Chunk(filesId: fileId, sequenceNumber: nextChunkNumber, data: .init(buffer: slice))
            let encoded = try FileWriter.encoder.encode(chunk)
            
            return fs.chunksCollection.insert(encoded).map { _ in
                self.flush(finalize: finalize)
            }
        } catch {
            return fs.eventLoop.newFailedFuture(error: error)
        }
    }
    
}<commit_msg>Sanity check<commit_after>import Foundation
import NIO
import MongoKitten

final class FileWriter {
    
    static let allocator = ByteBufferAllocator()
    static let encoder = BSONEncoder()
    
    let fs: GridFS
    let fileId: Primitive
    let chunkSize: Int
    var buffer: ByteBuffer
    var nextChunkNumber = 0
    var length = 0
    var finalized = false
    
    internal init(fs: GridFS, fileId: Primitive, chunkSize: Int, buffer: ByteBuffer? = nil) {
        self.fs = fs
        self.fileId = fileId
        self.chunkSize = chunkSize
        self.buffer = buffer ?? FileWriter.allocator.buffer(capacity: chunkSize)
    }
    
    public func write(data: ByteBuffer) -> EventLoopFuture<Void> {
        assert(self.finalized == false, "Writing to a finalized writer is an error")
        
        self.length += data.readableBytes
        var source = data
        buffer.write(buffer: &source)
        
        guard buffer.readableBytes >= chunkSize else {
            return fs.eventLoop.newSucceededFuture(result: ())
        }
        
        return self.flush()
    }
    
    public func finalize() -> EventLoopFuture<Void> {
        assert(self.finalized == false, "Finalizing a finalized writer is an error")
        
        self.finalized = true
        
        return self.flush(finalize: true)
            .then {
                let file = File(id: self.fileId,
                                length: self.length,
                                chunkSize: self.chunkSize,
                                metadata: nil, // TODO
                                filename: nil, // TODO
                                fs: self.fs)
                
                do {
                    let encoded = try FileWriter.encoder.encode(file)
                    
                    return self.fs.filesCollection.insert(encoded).map { _ in }
                } catch {
                    return self.fs.eventLoop.newFailedFuture(error: error)
                }
        }
    }
    
    private func flush(finalize: Bool = false) -> EventLoopFuture<Void> {
        guard buffer.readableBytes > 0, finalize || buffer.readableBytes >= chunkSize else {
            return fs.eventLoop.newSucceededFuture(result: ())
        }
        
        guard let slice = buffer.readSlice(length: buffer.readableBytes >= chunkSize ? chunkSize : buffer.readableBytes) else {
            // TODO: Replace with error future
            fatalError()
        }
        
        do {
            let chunk = Chunk(filesId: fileId, sequenceNumber: nextChunkNumber, data: .init(buffer: slice))
            let encoded = try FileWriter.encoder.encode(chunk)
            
            return fs.chunksCollection.insert(encoded).map { _ in
                self.flush(finalize: finalize)
            }
        } catch {
            return fs.eventLoop.newFailedFuture(error: error)
        }
    }
    
    deinit {
        assert(finalized == true || length == 0, "A GridFS FileWriter was deinitialized, while the writing has not been finalized. This will cause orphan chunks in the chunks collection in GridFS.")
    }
    
}
{'fraction_non_alphanumeric': 0.0822826808228268, 'fraction_numerical': 0.001161247511612475, 'mean_word_length': 1.60995670995671, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5807582', 'n_tokens_mistral': 1594, 'n_tokens_neox': 1537, 'n_words': 405}