<commit_before>/*++

Module Name:

    SAM.cpp

Abstract:

    Sequence Alignment Map (SAM) file writer and reader.

Environment:

    User mode service.

    SamWriter and SamReader (and their subclasses) aren't thread safe.

--*/

#include "stdafx.h"
#include "BigAlloc.h"
#include "Compat.h"
#include "Read.h"
#include "SAM.h"
#include "Tables.h"
#include "RangeSplitter.h"
#include "ParallelTask.h"

using std::max;
using std::min;


//
// You'd think this would be in the C library.
// Like strchr, but with a max length so it doesn't
// run over the end of the buffer.  Basically,
// strings suck in C.
//

    char *
strnchr(char *str, char charToFind, size_t maxLen)
{
    for (size_t i = 0; i < maxLen; i++) {
        if (str[i] == charToFind) {
            return str + i;
        }
        if (str[i] == 0) {
            return NULL;
        }
    }
    return NULL;
}

    const char *
strnchr(const char *str, char charToFind, size_t maxLen)
{
    for (size_t i = 0; i < maxLen; i++) {
        if (str[i] == charToFind) {
            return str + i;
        }
        if (str[i] == 0) {
            return NULL;
        }
    }
    return NULL;
}

SAMWriter:: ~SAMWriter()
{
}


SAMWriter* SAMWriter::create(const char *fileName, const Genome *genome)
{
    SimpleSAMWriter *writer = new SimpleSAMWriter();
    if (!writer->open(fileName, genome)) {
        delete writer;
        return NULL;
    } else {
        return writer;
    }
}

    bool
SAMWriter::generateHeader(const Genome *genome, char *header, size_t headerBufferSize, size_t *headerActualSize, bool sorted)
{
    
    size_t bytesConsumed = snprintf(header, headerBufferSize, "@HD\tVN:1.4\tSO:%s\n", sorted ? "coordinate" : "unsorted");
    if (bytesConsumed > headerBufferSize) {
        fprintf(stderr,"SAMWriter: header buffer too small\n");
        return false;
    }

    // Write an @SQ line for each chromosome / piece in the genome
    const Genome::Piece *pieces = genome->getPieces();
    int numPieces = genome->getNumPieces();
    unsigned genomeLen = genome->getCountOfBases();
    for (int i = 0; i < numPieces; i++) {
        unsigned start = pieces[i].beginningOffset;
        unsigned end = (i + 1 < numPieces) ? pieces[i+1].beginningOffset : genomeLen;
        bytesConsumed += snprintf(header + bytesConsumed, headerBufferSize - bytesConsumed, "@SQ\tSN:%s\tLN:%u\n", pieces[i].name, end - start);

        if (bytesConsumed > headerBufferSize) {
            fprintf(stderr,"SAMWriter: header buffer too small\n");
            return false;
        }
    }

    *headerActualSize = bytesConsumed;
    return true;
}

// Compute the CIGAR edit sequence string for a read against a given genome location.
// Returns this string if possible or "*" if we fail to compute it (which would likely
// be a bug due to lack of buffer space). The pointer returned may be to cigarBuf so it
// will only be valid until computeCigarString is called again.
    const char *
SAMWriter::computeCigarString(
    const Genome *              genome,
    LandauVishkinWithCigar *    lv,
    char *                      cigarBuf,
    int                         cigarBufLen,
    char *                      cigarBufWithClipping,
    int                         cigarBufWithClippingLen,
    const char *                data,
    unsigned                    dataLength,
    unsigned                    basesClippedBefore,
    unsigned                    basesClippedAfter,
    unsigned                    genomeLocation,
    bool                        isRC
)
{
    const char *reference = genome->getSubstring(genomeLocation, dataLength);
    int r;
    if (NULL != reference) {
        r = lv->computeEditDistance(
                       genome->getSubstring(genomeLocation, dataLength),
                        dataLength,
                        data,
                        dataLength,
                        MAX_K - 1,
                        cigarBuf,
                        cigarBufLen);
    } else {
        //
        // Fell off the end of the chromosome.
        //
        return "*";
    }

    if (r == -2) {
        fprintf(stderr, "WARNING: computeEditDistance returned -2; cigarBuf may be too small\n");
        return "*";
    } else if (r == -1) {
        static bool warningPrinted = false;
        if (!warningPrinted) {
            fprintf(stderr, "WARNING: computeEditDistance returned -1; this shouldn't happen\n");
            warningPrinted = true;
        }
        return "*";
    } else {
        // Add some CIGAR instructions for soft-clipping if we've ignored some bases in the read.
        char clipBefore[16] = {'\0'};
        char clipAfter[16] = {'\0'};
        if (basesClippedBefore > 0) {
            snprintf(clipBefore, sizeof(clipBefore), "%uS", basesClippedBefore);
        }
        if (basesClippedAfter > 0) {
            snprintf(clipAfter, sizeof(clipAfter), "%uS", basesClippedAfter);
        }
        snprintf(cigarBufWithClipping, cigarBufWithClippingLen, "%s%s%s", clipBefore, cigarBuf, clipAfter);
        return cigarBufWithClipping;
    }
}


    bool
SAMWriter::generateSAMText(
                Read *                      read, 
                AlignmentResult             result, 
                unsigned                    genomeLocation, 
                bool                        isRC, 
                bool                        hasMate, 
                bool                        firstInPair, 
                Read *                      mate, 
                AlignmentResult             mateResult, 
                unsigned                    mateLocation,
                bool                        mateIsRC, 
                const Genome *              genome, 
                LandauVishkinWithCigar *    lv, 
                char *                      buffer, 
                size_t                      bufferSpace, 
                size_t *                    spaceUsed)
{
    const int MAX_READ = 10000;
    const int cigarBufSize = MAX_READ * 2;
    char cigarBuf[cigarBufSize];

    const int cigarBufWithClippingSize = MAX_READ * 2 + 32;
    char cigarBufWithClipping[cigarBufWithClippingSize];

    int flags = 0;
    const char *pieceName = "*";
    unsigned positionInPiece = 0;
    int mapQuality = 0;
    const char *cigar = "*";
    const char *matePieceName = "*";
    unsigned matePositionInPiece = 0;
    _int64 templateLength = 0;

    //
    // If the aligner said it didn't find anything, treat it as such.  Sometimes it will emit the
    // best match that it found, even if it's not within the maximum edit distance limit (but will
    // then say NotFound).  Here, we force that to be SAM_UNMAPPED.
    //
    if (NotFound == result) {
        genomeLocation = 0xffffffff;
    }

    // Write the data and quality strings. If the read is reverse complemented, these need to
    // be backwards from the original read. Also, both need to be unclipped.
    char data[MAX_READ];
    char quality[MAX_READ];
    const char *clippedData;
    unsigned clippedLength = read->getDataLength();
    unsigned fullLength = read->getUnclippedLength();
    unsigned basesClippedBefore;
    unsigned basesClippedAfter;
    if (isRC) {
      for (unsigned i = 0; i < fullLength; i++) {
        data[fullLength - 1 - i] = COMPLEMENT[read->getUnclippedData()[i]];
        quality[fullLength - 1 - i] = read->getUnclippedQuality()[i];
      }
      clippedData = &data[fullLength - clippedLength - read->getFrontClippedLength()];
      basesClippedBefore = fullLength - clippedLength - read->getFrontClippedLength();
      basesClippedAfter = read->getFrontClippedLength();
    } else {
      memcpy(data, read->getUnclippedData(), read->getUnclippedLength());
      memcpy(quality, read->getUnclippedQuality(), read->getUnclippedLength());
      clippedData = read->getData();
      basesClippedBefore = read->getFrontClippedLength();
      basesClippedAfter = fullLength - clippedLength - basesClippedBefore;
    }

    if (genomeLocation != 0xFFFFFFFF) {
        // This could be either a single hit read or a multiple hit read where we just
        // returned one location, but either way, let's print that location. We will then
        // set the quality to 60 if it was single or 0 if it was multiple. These are the
        // values the SAM FAQ suggests for aligners that don't compute confidence scores.
        if (isRC) {
            flags |= SAM_REVERSE_COMPLEMENT;
        }
        const Genome::Piece *piece = genome->getPieceAtLocation(genomeLocation);
        pieceName = piece->name;
        positionInPiece = genomeLocation - piece->beginningOffset + 1; // SAM is 1-based
        cigar = computeCigarString(genome, lv, cigarBuf, cigarBufSize, cigarBufWithClipping, cigarBufWithClippingSize, 
                                   clippedData, clippedLength, basesClippedBefore, basesClippedAfter,
                                   genomeLocation, isRC);
        mapQuality = (result == SingleHit || result == CertainHit) ? 60 : 0;
    } else {
        flags |= SAM_UNMAPPED;
    }

    if (hasMate) {
        flags |= SAM_MULTI_SEGMENT;
        flags |= (firstInPair ? SAM_FIRST_SEGMENT : SAM_LAST_SEGMENT);
        if (isOneLocation(result) && isOneLocation(mateResult)) {
            flags |= SAM_ALL_ALIGNED;
            // Also compute the length of the whole paired-end string whose ends we saw. This is slightly
            // tricky because (a) we may have clipped some bases before/after each end and (b) we need to
            // give a signed result based on whether our read is first or second in the pair.
            _int64 myStart = genomeLocation - basesClippedBefore;
            _int64 myEnd = genomeLocation + clippedLength + basesClippedAfter;
            _int64 mateBasesClippedBefore = mate->getFrontClippedLength();
            _int64 mateBasesClippedAfter = mate->getUnclippedLength() - mate->getDataLength() - mateBasesClippedBefore;
            _int64 mateStart = mateLocation - (mateIsRC ? mateBasesClippedAfter : mateBasesClippedBefore);
            _int64 mateEnd = mateLocation + mate->getDataLength() + (!mateIsRC ? mateBasesClippedAfter : mateBasesClippedBefore);
            if (myStart < mateStart) {
                templateLength = mateEnd - myStart;
            } else {
                templateLength = -(myEnd - mateStart);
            }
        }
        if (mateIsRC) {
            flags |= SAM_NEXT_REVERSED;
        }
        if (mateLocation != 0xFFFFFFFF) {
            const Genome::Piece *piece = genome->getPieceAtLocation(mateLocation);
            matePieceName = piece->name;
            matePositionInPiece = mateLocation - piece->beginningOffset + 1;
        }
    }

    if (result == MultipleHits && genomeLocation == 0xFFFFFFFF) {
        // Read was MultipleHits but we didn't return even one potential location, which means that all
        // the seeds were too popular. Report the mapping quality as 1 to let users differentiate this
        // from a NotFound read.
        mapQuality = 1;
    }

    // Write the SAM entry, which requires the following fields:
    //
    // 1. QNAME: Query name of the read or the read pair
    // 2. FLAG: Bitwise flag (pairing, strand, mate strand, etc.)
    // 3. RNAME: Reference sequence name
    // 4. POS: 1-Based leftmost position of clipped alignment
    // 5. MAPQ: Mapping quality (Phred-scaled)
    // 6. CIGAR: Extended CIGAR string (operations: MIDNSHP)
    // 7. MRNM: Mate reference name (‘=’ if same as RNAME)
    // 8. MPOS: 1-based leftmost mate position
    // 9. ISIZE: Inferred insert size
    // 10. SEQQuery: Sequence on the same strand as the reference
    // 11. QUAL: Query quality (ASCII-33=Phred base quality)    

    //
    // Some FASTQ files have spaces in their ID strings, which is illegal in SAM.  Just truncate them at the space.
    //
    unsigned readLen;
    const char *firstSpace = strnchr(read->getId(),' ',read->getIdLength());
    if (NULL == firstSpace) {
        readLen = read->getIdLength();
    } else {
        readLen = (unsigned)(firstSpace - read->getId());
    }

    int charsInString = snprintf(buffer, bufferSpace, "%.*s\t%d\t%s\t%u\t%d\t%s\t%s\t%u\t%lld\t%.*s\t%.*s\n",
        readLen, read->getId(),
        flags,
        pieceName,
        positionInPiece,
        mapQuality,
        cigar,
        matePieceName,
        matePositionInPiece,
        templateLength,
        fullLength, data,
        fullLength, quality);

    if (charsInString > bufferSpace) {
        //
        // Out of buffer space.
        //
        return false;
    }

    if (NULL != spaceUsed) {
        *spaceUsed = charsInString;
    }
    return true;
}


SimpleSAMWriter::SimpleSAMWriter()
{
    file = NULL;
}


SimpleSAMWriter::~SimpleSAMWriter()
{
    if (file != NULL) {
        close();
    }
}


bool SimpleSAMWriter::open(const char* fileName, const Genome *genome)
{
    buffer = (char *) BigAlloc(BUFFER_SIZE);
    if (buffer == NULL) {
        fprintf(stderr, "allocating write buffer failed\n");
        return false;
    }
    file = fopen(fileName, "w");
    if (file == NULL) {
        fprintf(stderr, "fopen failed\n");
        return false;
    }
    setvbuf(file, buffer, _IOFBF, BUFFER_SIZE);
    this->genome = genome;

    // Write out SAM header
    char *headerBuffer = new char[HEADER_BUFFER_SIZE];
    size_t headerSize;
    if (!generateHeader(genome,headerBuffer,HEADER_BUFFER_SIZE,&headerSize, false)) {
        fprintf(stderr,"SimpleSAMWriter: unable to generate SAM header\n");
        return false;
    }
    fprintf(file, "%s", headerBuffer);
    delete[] headerBuffer;

    return true;
}


bool SimpleSAMWriter::write(Read *read, AlignmentResult result, unsigned genomeLocation, bool isRC)
{
    if (file == NULL) {
        return false;
    }

    write(read, result, genomeLocation, isRC, false, false, NULL, NotFound, 0, false);
    return true;
}


bool SimpleSAMWriter::writePair(Read *read0, Read *read1, PairedAlignmentResult *result)
{
    if (file == NULL) {
        return false;
    }

    write(read0, result->status[0], result->location[0], result->isRC[0], true, true,
          read1, result->status[1], result->location[1], result->isRC[1]);
    write(read1, result->status[1], result->location[1], result->isRC[1], true, false,
          read0, result->status[0], result->location[0], result->isRC[0]);
    return true;
}


void SimpleSAMWriter::write(
        Read *read,
        AlignmentResult result,
        unsigned genomeLocation,
        bool isRC,
        bool hasMate,
        bool firstInPair,
        Read *mate,
        AlignmentResult mateResult,
        unsigned mateLocation,
        bool mateIsRC)
{
    const unsigned maxLineLength = 25000;
    char outputBuffer[maxLineLength];
    size_t outputBufferUsed;

    if (!generateSAMText(read,result,genomeLocation,isRC,hasMate,firstInPair,mate,mateResult,
            mateLocation,mateIsRC, genome, &lv, outputBuffer, maxLineLength,&outputBufferUsed)) {
        fprintf(stderr,"SimpleSAMWriter: tried to generate too long of a SAM line (> %d)\n",maxLineLength);
        exit(1);
    }
    
    if (1 != fwrite(outputBuffer,outputBufferUsed,1,file)) {
        fprintf(stderr,"Unable to write to SAM file.\n");
        exit(1);
    }
}


bool SimpleSAMWriter::close()
{
    if (file == NULL) {
        return false;
    }
    fclose(file);
    BigDealloc(buffer);
    file = NULL;
    buffer = NULL;
    return true;
}

ThreadSAMWriter::ThreadSAMWriter(size_t i_bufferSize)
    : remainingBufferSpace(i_bufferSize), bufferBeingCreated(0), bufferSize(i_bufferSize)
{
    buffer[0] = NULL;
    buffer[1] = NULL;
}

    bool
ThreadSAMWriter::initialize(AsyncFile* file, const Genome *i_genome, volatile _int64 *i_nextWriteOffset)
{
    genome = i_genome;
    nextWriteOffset = i_nextWriteOffset;
    buffer[0] = (char *)BigAlloc(bufferSize);
    buffer[1] = (char *)BigAlloc(bufferSize);
    writer[0] = file->getWriter();
    writer[1] = file->getWriter();

    if (NULL == buffer[0] || NULL == buffer[1] || NULL == writer[0] || NULL == writer[1]) {
        fprintf(stderr,"ThreadSAMWriter: failed to initialize\n");
        return false;
    }
    return true;
}
    
ThreadSAMWriter::~ThreadSAMWriter()
{
    BigDealloc(buffer[0]);
    BigDealloc(buffer[1]);
    delete writer[0];
    delete writer[1];
}

    bool
ThreadSAMWriter::close()
{
    if (remainingBufferSpace != bufferSize) {
        if (!startIo()) {
            fprintf(stderr,"WindowsSAMWriter::close(): startIo failed\n");
            return false;
        }

        if (!waitForIoCompletion()) {
            fprintf(stderr,"WindowsSAMWriter::close(): waitForIoCompletion failed\n");
            return false;
        }
    }
    bool ok = writer[0]->close();
    ok &= writer[1]->close();
    if (! ok) {
        fprintf(stderr, "WindowsSAMWriter::close closing writer failed\n");
    }

    return true;
}
    
    bool
ThreadSAMWriter::write(Read *read, AlignmentResult result, unsigned genomeLocation, bool isRC)
{
    size_t sizeUsed;
    if (!generateSAMText(read,result,genomeLocation,isRC,false,true,NULL,UnknownAlignment,0,false,genome,&lv,
            buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace,remainingBufferSpace, &sizeUsed)) {

        if (!startIo()) {
            return false;
        }

        if (!generateSAMText(read,result,genomeLocation,isRC,false,true,NULL,UnknownAlignment,0,false,genome,&lv,
                buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace,remainingBufferSpace, &sizeUsed)) {

            fprintf(stderr,"WindowsSAMWriter: create SAM string into fresh buffer failed\n");
            return false;
        }
    }
    size_t bufferOffset = bufferSize - remainingBufferSpace;
    remainingBufferSpace -= sizeUsed;
    afterWrite(result != NotFound ? genomeLocation : UINT32_MAX, bufferOffset, sizeUsed);
    return true;
}

    bool
ThreadSAMWriter::writePair(Read *read0, Read *read1, PairedAlignmentResult *result)
{
    //
    // We need to write both halves of the pair into the same buffer, so that a write from
    // some other thread doesn't separate them.  So, try the writes and if either doesn't
    // work start IO and try again.
    //
    size_t sizeUsed[2];
    bool writesFit = generateSAMText(read0,result->status[0],result->location[0],result->isRC[0],true,true,
                                     read1,result->status[1],result->location[1],result->isRC[1],
                        genome, &lv, buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace,remainingBufferSpace,&sizeUsed[0]);
    if (writesFit) {
        writesFit = generateSAMText(read1,result->status[1],result->location[1],result->isRC[1],true,false,
                                    read0,result->status[0],result->location[0],result->isRC[0],
                        genome, &lv, buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace + sizeUsed[0],remainingBufferSpace-sizeUsed[0],&sizeUsed[1]);
    }

    if (!writesFit) {
        if (!startIo()) {
            return false;
        }

        if (!generateSAMText(read0,result->status[0],result->location[0],result->isRC[0],true,true,
                             read1,result->status[1],result->location[1],result->isRC[1],
                genome, &lv, buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace,remainingBufferSpace,&sizeUsed[0]) ||
            !generateSAMText(read1,result->status[1],result->location[1],result->isRC[1],true,false,
                             read0,result->status[0],result->location[0],result->isRC[0],
                genome, &lv, buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace + sizeUsed[0],remainingBufferSpace-sizeUsed[0],&sizeUsed[1])) {

                fprintf(stderr,"WindowsSAMWriter: create SAM string into fresh buffer failed\n");
                return false;
            }
    }

    size_t bufferOffset[2] = {bufferSize - sizeUsed[0] - sizeUsed[1], bufferSize - sizeUsed[1]};
    remainingBufferSpace -= (sizeUsed[0] + sizeUsed[1]);
    afterWrite(result->status[0] != NotFound ? result->location[0] : UINT32_MAX, bufferOffset[0], sizeUsed[0]);
    afterWrite(result->status[1] != NotFound ? result->location[1] : UINT32_MAX, bufferOffset[1], sizeUsed[1]);
    return true;
}

    ParallelSAMWriter *
ParallelSAMWriter::create(
    const char *fileName,
    const Genome *genome,
    unsigned nThreads,
    bool sort,
    size_t sortBufferMemory)
{
    ParallelSAMWriter *parallelWriter = sort
        ? new SortedParallelSAMWriter(sortBufferMemory)
        : new ParallelSAMWriter();
    if (!parallelWriter->initialize(fileName, genome, nThreads, sort)) {
        fprintf(stderr, "unable to initialize parallel SAM writer\n");
        delete parallelWriter;
        return NULL;
    }
    return parallelWriter;
}

    bool

ParallelSAMWriter::initialize(const char *fileName, const Genome *genome, unsigned i_nThreads, bool sorted)
{
    file = AsyncFile::open(fileName, true);
    if (NULL == file) {
        fprintf(stderr,"Unable to create SAM file '%s'\n",fileName);
        return false;
    }

    const size_t headerBufferSize = 20000;
    char headerBuffer[headerBufferSize];
    size_t headerActualSize;

    if (!SAMWriter::generateHeader(genome,headerBuffer,headerBufferSize,&headerActualSize, sorted)) {
        fprintf(stderr,"WindowsParallelSAMWriter: unable to generate SAM header.\n");
        return false;
    }

    AsyncFile::Writer* hwriter = file->getWriter();

    if (NULL == hwriter) {
        fprintf(stderr,"ParallelSAMWriter: unable to create writer\n");
        return false;
    }

    size_t bytesWritten;
    if (! hwriter->beginWrite(headerBuffer, headerActualSize, 0, &bytesWritten)) {
        fprintf(stderr,"ParallelSAMWriter: unable to write header to file\n");
        return false;
    }

    if (! hwriter->waitForCompletion()) {
        fprintf(stderr,"ParallelSAMWriter: failed to complete\n");
        return false;
    }
    if (! hwriter->close()) {
        fprintf(stderr, "ParallelSAMWriter: failed to close\n");
    }
    delete hwriter;

    nextWriteOffset = headerActualSize;

    nThreads = i_nThreads;
    writer = new ThreadSAMWriter *[nThreads];

    if (!createThreadWriters(genome)) {
        fprintf(stderr,"Unable to create SAM writer.\n");
        return false;
    }

    return true;
}

    bool
ParallelSAMWriter::createThreadWriters(const Genome* genome)
{
    bool worked = true;
    for (unsigned i = 0; i < nThreads; i++) {
        writer[i] = new ThreadSAMWriter(UnsortedBufferSize);
        worked &= writer[i]->initialize(file, genome, &nextWriteOffset);
    }
    return worked;
}

ParallelSAMWriter::~ParallelSAMWriter()
{
    delete [] writer;
}

    SAMWriter*
ParallelSAMWriter::getWriterForThread(
    unsigned whichThread)
{
    _ASSERT(whichThread < nThreads);
    return writer[whichThread];
}

    bool
ParallelSAMWriter::close()
{
    for (unsigned i = 0; i < nThreads; i++) {
        // writers were already closed by each thread so they could flush in parallel
        delete writer[i];
        writer[i] = NULL;
    }
    if (! file->close()) {
        fprintf(stderr, "ParallelSAMWriter::close file close failed\n");
    }
    delete file;
    file = NULL;
    return true;
}

    bool
ThreadSAMWriter::startIo()
{
    //
    // It didn't fit in the buffer.  Start writing it.
    //
    _int64 writeOffset = InterlockedAdd64AndReturnNewValue(nextWriteOffset,bufferSize - remainingBufferSpace) - (bufferSize - remainingBufferSpace);
    if (!beforeFlush(writeOffset, bufferSize - remainingBufferSpace)) {
        fprintf(stderr, "ThreadSAMWriter: beforeFlush failed\n");
        return false;
    }
    if (!writer[bufferBeingCreated]->beginWrite(buffer[bufferBeingCreated], bufferSize - remainingBufferSpace, writeOffset, NULL)) {
        fprintf(stderr,"ThreadSAMWriter: WriteFile failed\n");
        return false;
    }

    //
    // If necessary, wait for the other buffer to finish writing.
    //
    if (!waitForIoCompletion()) {
        fprintf(stderr,"ThreadSAMWriter: waitForIoCompletion failed\n");
        return false;
    }
    bufferBeingCreated = 1 - bufferBeingCreated;
    remainingBufferSpace = bufferSize;

    return true;
}

    bool
ThreadSAMWriter::waitForIoCompletion()
{
    return writer[1 - bufferBeingCreated]->waitForCompletion();
}

SortBlock::SortBlock()
    : entries(0), fileOffset(0), fileBytes(0), index(0), reader()
{
}

SortBlock::SortBlock(size_t capacity)
    : entries(capacity), fileOffset(0), fileBytes(0), index(0), reader()
{
}

SortBlock::SortBlock(
    SortBlock& other)
{
    entries = other.entries;
    fileOffset = other.fileOffset;
    fileBytes = other.fileBytes;
    index = other.index;
    reader = other.reader;
}

    void
SortBlock::operator=(
    SortBlock& other)
{
    entries = other.entries;
    fileOffset = other.fileOffset;
    fileBytes = other.fileBytes;
    index = other.index;
    index = other.index;
    reader = other.reader;
}

SortedThreadSAMWriter::SortedThreadSAMWriter(size_t i_bufferSize)
    : ThreadSAMWriter(i_bufferSize),
        parent(NULL),
        largest(1000),
        locations(1000)
{
}

SortedThreadSAMWriter::~SortedThreadSAMWriter()
{
}

    bool
SortedThreadSAMWriter::initialize(
    SortedParallelSAMWriter* i_parent,
    const Genome* i_genome)
{
    parent = i_parent;
    return ThreadSAMWriter::initialize(parent->file, i_genome, &parent->nextWriteOffset);
}

    void
SortedThreadSAMWriter::afterWrite(
    unsigned location,
    size_t bufferOffset,
    unsigned length)
{
    SortEntry entry(bufferOffset, length, location);
    locations.entries.push_back(entry);
}

    bool
SortedThreadSAMWriter::beforeFlush(_int64 fileOffset, _int64 length)
{
    // sort buffered reads by location for later merge sort
    std::sort(locations.entries.begin(), locations.entries.end(), SortEntry::comparator);
    
    // wait for IO of other buffer to finish since we're going to sort into there
    if (! waitForIoCompletion()) {
        fprintf(stderr, "SortedThreadSAMWriter waitForIoCompletion failed\n");
        return false;
    }
    
    // copy into other buffer in sorted order & switch buffers
    unsigned target = 0;
    for (VariableSizeVector<SortEntry>::iterator i = locations.entries.begin(); i != locations.entries.end(); i++) {
        memcpy(buffer[1 - bufferBeingCreated] + target, buffer[bufferBeingCreated] + i->offset, i->length);
        i->offset = fileOffset + target;
        target += i->length;
    }
    bufferBeingCreated = 1 - bufferBeingCreated;
    
    // remember offsets for full-file sort, get a new vector of appropriate size
    locations.fileOffset = fileOffset;
    locations.fileBytes = length;
    parent->addLocations(locations);
    if (locations.entries.size() > largest) {
        largest = (locations.entries.size() * 6) / 5; // grow by 20% if it exceeds prior max
    }
    locations.entries.reserve(largest);

    return true;
}

    bool
SortedParallelSAMWriter::initialize(
    const char *fileName,
    const Genome *genome,
    unsigned i_nThreads,
    bool sorted)
{
    InitializeExclusiveLock(&lock);
    sortedFile = fileName;
    tempFile = (char*) malloc(strlen(fileName) + 5);
    strcpy(tempFile, fileName);
    strcat(tempFile, ".tmp");
    bool ok = ParallelSAMWriter::initialize(tempFile, genome, i_nThreads, sorted);
    headerSize = nextWriteOffset;
    return ok;
}
    
    bool
SortedParallelSAMWriter::createThreadWriters(const Genome* genome)
{
    size_t bufferSize = totalMemory / nThreads / 2;
    bool worked = true;
    for (unsigned i = 0; i < nThreads; i++) {
        SortedThreadSAMWriter* w = new SortedThreadSAMWriter(bufferSize);
        writer[i] = w;
        worked &= w->initialize(this, genome);
    }
    return worked;
}

    bool
SortedParallelSAMWriter::close()
{
    if (! ParallelSAMWriter::close()) {
        return false;
    }
    DestroyExclusiveLock(&lock);

    // merge sort from temp file into sorted file
    printf("sorting...");
    _int64 start = timeInMillis();

    // first replace offset in entries with block index, and sort them all in one large array
    size_t total = 0;
    for (VariableSizeVector<SortBlock>::iterator i = locations.begin(); i != locations.end(); i++) {
        total += i->entries.size();
    }
    SortEntry* entries = (SortEntry*) BigAlloc(total * sizeof(SortEntry));
    size_t offset = 0;
    for (VariableSizeVector<SortBlock>::iterator i = locations.begin(); i != locations.end(); i++) {
        unsigned n = i->entries.size();
        unsigned blockIndex = i - locations.begin();
        memcpy(entries + offset, i->entries.begin(), n * (size_t) sizeof(SortEntry));
        for (unsigned j = 0; j < n; j++) {
            entries[offset + j].offset = blockIndex;
        }
        offset += i->entries.size();
    }
    std::sort(entries, entries + total, SortEntry::comparator);
    printf(" %%ld s\nwriting sorted reads...", (timeInMillis() - start) / 1000);
    start = timeInMillis();

    // setup - open all files, read first block, begin read for second
    AsyncFile* temp = AsyncFile::open(tempFile, false);
    const size_t ReadBufferSize = UnsortedBufferSize;
    char* buffers = (char*) BigAlloc(ReadBufferSize * 2 * locations.size());
    unsigned j = 0;
    if (timeInMillis() - start > 1000) {
        printf(" (allocated %lld Mb in %lld s)",
            ReadBufferSize * 2 * locations.size() / (2 << 20), (timeInMillis() - start) / 1000);
    }
    for (VariableSizeVector<SortBlock>::iterator i = locations.begin(); i != locations.end(); i++, j++) {
        i->reader.open(temp, i->fileOffset, i->fileBytes, ReadBufferSize, true,
            buffers + j * 2 * ReadBufferSize, buffers + (j * 2 + 1) * ReadBufferSize);
    }
    for (VariableSizeVector<SortBlock>::iterator i = locations.begin(); i != locations.end(); i++) {
        i->reader.endOpen();<commit_msg>Printf string<commit_after>/*++

Module Name:

    SAM.cpp

Abstract:

    Sequence Alignment Map (SAM) file writer and reader.

Environment:

    User mode service.

    SamWriter and SamReader (and their subclasses) aren't thread safe.

--*/

#include "stdafx.h"
#include "BigAlloc.h"
#include "Compat.h"
#include "Read.h"
#include "SAM.h"
#include "Tables.h"
#include "RangeSplitter.h"
#include "ParallelTask.h"

using std::max;
using std::min;


//
// You'd think this would be in the C library.
// Like strchr, but with a max length so it doesn't
// run over the end of the buffer.  Basically,
// strings suck in C.
//

    char *
strnchr(char *str, char charToFind, size_t maxLen)
{
    for (size_t i = 0; i < maxLen; i++) {
        if (str[i] == charToFind) {
            return str + i;
        }
        if (str[i] == 0) {
            return NULL;
        }
    }
    return NULL;
}

    const char *
strnchr(const char *str, char charToFind, size_t maxLen)
{
    for (size_t i = 0; i < maxLen; i++) {
        if (str[i] == charToFind) {
            return str + i;
        }
        if (str[i] == 0) {
            return NULL;
        }
    }
    return NULL;
}

SAMWriter:: ~SAMWriter()
{
}


SAMWriter* SAMWriter::create(const char *fileName, const Genome *genome)
{
    SimpleSAMWriter *writer = new SimpleSAMWriter();
    if (!writer->open(fileName, genome)) {
        delete writer;
        return NULL;
    } else {
        return writer;
    }
}

    bool
SAMWriter::generateHeader(const Genome *genome, char *header, size_t headerBufferSize, size_t *headerActualSize, bool sorted)
{
    
    size_t bytesConsumed = snprintf(header, headerBufferSize, "@HD\tVN:1.4\tSO:%s\n", sorted ? "coordinate" : "unsorted");
    if (bytesConsumed > headerBufferSize) {
        fprintf(stderr,"SAMWriter: header buffer too small\n");
        return false;
    }

    // Write an @SQ line for each chromosome / piece in the genome
    const Genome::Piece *pieces = genome->getPieces();
    int numPieces = genome->getNumPieces();
    unsigned genomeLen = genome->getCountOfBases();
    for (int i = 0; i < numPieces; i++) {
        unsigned start = pieces[i].beginningOffset;
        unsigned end = (i + 1 < numPieces) ? pieces[i+1].beginningOffset : genomeLen;
        bytesConsumed += snprintf(header + bytesConsumed, headerBufferSize - bytesConsumed, "@SQ\tSN:%s\tLN:%u\n", pieces[i].name, end - start);

        if (bytesConsumed > headerBufferSize) {
            fprintf(stderr,"SAMWriter: header buffer too small\n");
            return false;
        }
    }

    *headerActualSize = bytesConsumed;
    return true;
}

// Compute the CIGAR edit sequence string for a read against a given genome location.
// Returns this string if possible or "*" if we fail to compute it (which would likely
// be a bug due to lack of buffer space). The pointer returned may be to cigarBuf so it
// will only be valid until computeCigarString is called again.
    const char *
SAMWriter::computeCigarString(
    const Genome *              genome,
    LandauVishkinWithCigar *    lv,
    char *                      cigarBuf,
    int                         cigarBufLen,
    char *                      cigarBufWithClipping,
    int                         cigarBufWithClippingLen,
    const char *                data,
    unsigned                    dataLength,
    unsigned                    basesClippedBefore,
    unsigned                    basesClippedAfter,
    unsigned                    genomeLocation,
    bool                        isRC
)
{
    const char *reference = genome->getSubstring(genomeLocation, dataLength);
    int r;
    if (NULL != reference) {
        r = lv->computeEditDistance(
                       genome->getSubstring(genomeLocation, dataLength),
                        dataLength,
                        data,
                        dataLength,
                        MAX_K - 1,
                        cigarBuf,
                        cigarBufLen);
    } else {
        //
        // Fell off the end of the chromosome.
        //
        return "*";
    }

    if (r == -2) {
        fprintf(stderr, "WARNING: computeEditDistance returned -2; cigarBuf may be too small\n");
        return "*";
    } else if (r == -1) {
        static bool warningPrinted = false;
        if (!warningPrinted) {
            fprintf(stderr, "WARNING: computeEditDistance returned -1; this shouldn't happen\n");
            warningPrinted = true;
        }
        return "*";
    } else {
        // Add some CIGAR instructions for soft-clipping if we've ignored some bases in the read.
        char clipBefore[16] = {'\0'};
        char clipAfter[16] = {'\0'};
        if (basesClippedBefore > 0) {
            snprintf(clipBefore, sizeof(clipBefore), "%uS", basesClippedBefore);
        }
        if (basesClippedAfter > 0) {
            snprintf(clipAfter, sizeof(clipAfter), "%uS", basesClippedAfter);
        }
        snprintf(cigarBufWithClipping, cigarBufWithClippingLen, "%s%s%s", clipBefore, cigarBuf, clipAfter);
        return cigarBufWithClipping;
    }
}


    bool
SAMWriter::generateSAMText(
                Read *                      read, 
                AlignmentResult             result, 
                unsigned                    genomeLocation, 
                bool                        isRC, 
                bool                        hasMate, 
                bool                        firstInPair, 
                Read *                      mate, 
                AlignmentResult             mateResult, 
                unsigned                    mateLocation,
                bool                        mateIsRC, 
                const Genome *              genome, 
                LandauVishkinWithCigar *    lv, 
                char *                      buffer, 
                size_t                      bufferSpace, 
                size_t *                    spaceUsed)
{
    const int MAX_READ = 10000;
    const int cigarBufSize = MAX_READ * 2;
    char cigarBuf[cigarBufSize];

    const int cigarBufWithClippingSize = MAX_READ * 2 + 32;
    char cigarBufWithClipping[cigarBufWithClippingSize];

    int flags = 0;
    const char *pieceName = "*";
    unsigned positionInPiece = 0;
    int mapQuality = 0;
    const char *cigar = "*";
    const char *matePieceName = "*";
    unsigned matePositionInPiece = 0;
    _int64 templateLength = 0;

    //
    // If the aligner said it didn't find anything, treat it as such.  Sometimes it will emit the
    // best match that it found, even if it's not within the maximum edit distance limit (but will
    // then say NotFound).  Here, we force that to be SAM_UNMAPPED.
    //
    if (NotFound == result) {
        genomeLocation = 0xffffffff;
    }

    // Write the data and quality strings. If the read is reverse complemented, these need to
    // be backwards from the original read. Also, both need to be unclipped.
    char data[MAX_READ];
    char quality[MAX_READ];
    const char *clippedData;
    unsigned clippedLength = read->getDataLength();
    unsigned fullLength = read->getUnclippedLength();
    unsigned basesClippedBefore;
    unsigned basesClippedAfter;
    if (isRC) {
      for (unsigned i = 0; i < fullLength; i++) {
        data[fullLength - 1 - i] = COMPLEMENT[read->getUnclippedData()[i]];
        quality[fullLength - 1 - i] = read->getUnclippedQuality()[i];
      }
      clippedData = &data[fullLength - clippedLength - read->getFrontClippedLength()];
      basesClippedBefore = fullLength - clippedLength - read->getFrontClippedLength();
      basesClippedAfter = read->getFrontClippedLength();
    } else {
      memcpy(data, read->getUnclippedData(), read->getUnclippedLength());
      memcpy(quality, read->getUnclippedQuality(), read->getUnclippedLength());
      clippedData = read->getData();
      basesClippedBefore = read->getFrontClippedLength();
      basesClippedAfter = fullLength - clippedLength - basesClippedBefore;
    }

    if (genomeLocation != 0xFFFFFFFF) {
        // This could be either a single hit read or a multiple hit read where we just
        // returned one location, but either way, let's print that location. We will then
        // set the quality to 60 if it was single or 0 if it was multiple. These are the
        // values the SAM FAQ suggests for aligners that don't compute confidence scores.
        if (isRC) {
            flags |= SAM_REVERSE_COMPLEMENT;
        }
        const Genome::Piece *piece = genome->getPieceAtLocation(genomeLocation);
        pieceName = piece->name;
        positionInPiece = genomeLocation - piece->beginningOffset + 1; // SAM is 1-based
        cigar = computeCigarString(genome, lv, cigarBuf, cigarBufSize, cigarBufWithClipping, cigarBufWithClippingSize, 
                                   clippedData, clippedLength, basesClippedBefore, basesClippedAfter,
                                   genomeLocation, isRC);
        mapQuality = (result == SingleHit || result == CertainHit) ? 60 : 0;
    } else {
        flags |= SAM_UNMAPPED;
    }

    if (hasMate) {
        flags |= SAM_MULTI_SEGMENT;
        flags |= (firstInPair ? SAM_FIRST_SEGMENT : SAM_LAST_SEGMENT);
        if (isOneLocation(result) && isOneLocation(mateResult)) {
            flags |= SAM_ALL_ALIGNED;
            // Also compute the length of the whole paired-end string whose ends we saw. This is slightly
            // tricky because (a) we may have clipped some bases before/after each end and (b) we need to
            // give a signed result based on whether our read is first or second in the pair.
            _int64 myStart = genomeLocation - basesClippedBefore;
            _int64 myEnd = genomeLocation + clippedLength + basesClippedAfter;
            _int64 mateBasesClippedBefore = mate->getFrontClippedLength();
            _int64 mateBasesClippedAfter = mate->getUnclippedLength() - mate->getDataLength() - mateBasesClippedBefore;
            _int64 mateStart = mateLocation - (mateIsRC ? mateBasesClippedAfter : mateBasesClippedBefore);
            _int64 mateEnd = mateLocation + mate->getDataLength() + (!mateIsRC ? mateBasesClippedAfter : mateBasesClippedBefore);
            if (myStart < mateStart) {
                templateLength = mateEnd - myStart;
            } else {
                templateLength = -(myEnd - mateStart);
            }
        }
        if (mateIsRC) {
            flags |= SAM_NEXT_REVERSED;
        }
        if (mateLocation != 0xFFFFFFFF) {
            const Genome::Piece *piece = genome->getPieceAtLocation(mateLocation);
            matePieceName = piece->name;
            matePositionInPiece = mateLocation - piece->beginningOffset + 1;
        }
    }

    if (result == MultipleHits && genomeLocation == 0xFFFFFFFF) {
        // Read was MultipleHits but we didn't return even one potential location, which means that all
        // the seeds were too popular. Report the mapping quality as 1 to let users differentiate this
        // from a NotFound read.
        mapQuality = 1;
    }

    // Write the SAM entry, which requires the following fields:
    //
    // 1. QNAME: Query name of the read or the read pair
    // 2. FLAG: Bitwise flag (pairing, strand, mate strand, etc.)
    // 3. RNAME: Reference sequence name
    // 4. POS: 1-Based leftmost position of clipped alignment
    // 5. MAPQ: Mapping quality (Phred-scaled)
    // 6. CIGAR: Extended CIGAR string (operations: MIDNSHP)
    // 7. MRNM: Mate reference name (‘=’ if same as RNAME)
    // 8. MPOS: 1-based leftmost mate position
    // 9. ISIZE: Inferred insert size
    // 10. SEQQuery: Sequence on the same strand as the reference
    // 11. QUAL: Query quality (ASCII-33=Phred base quality)    

    //
    // Some FASTQ files have spaces in their ID strings, which is illegal in SAM.  Just truncate them at the space.
    //
    unsigned readLen;
    const char *firstSpace = strnchr(read->getId(),' ',read->getIdLength());
    if (NULL == firstSpace) {
        readLen = read->getIdLength();
    } else {
        readLen = (unsigned)(firstSpace - read->getId());
    }

    int charsInString = snprintf(buffer, bufferSpace, "%.*s\t%d\t%s\t%u\t%d\t%s\t%s\t%u\t%lld\t%.*s\t%.*s\n",
        readLen, read->getId(),
        flags,
        pieceName,
        positionInPiece,
        mapQuality,
        cigar,
        matePieceName,
        matePositionInPiece,
        templateLength,
        fullLength, data,
        fullLength, quality);

    if (charsInString > bufferSpace) {
        //
        // Out of buffer space.
        //
        return false;
    }

    if (NULL != spaceUsed) {
        *spaceUsed = charsInString;
    }
    return true;
}


SimpleSAMWriter::SimpleSAMWriter()
{
    file = NULL;
}


SimpleSAMWriter::~SimpleSAMWriter()
{
    if (file != NULL) {
        close();
    }
}


bool SimpleSAMWriter::open(const char* fileName, const Genome *genome)
{
    buffer = (char *) BigAlloc(BUFFER_SIZE);
    if (buffer == NULL) {
        fprintf(stderr, "allocating write buffer failed\n");
        return false;
    }
    file = fopen(fileName, "w");
    if (file == NULL) {
        fprintf(stderr, "fopen failed\n");
        return false;
    }
    setvbuf(file, buffer, _IOFBF, BUFFER_SIZE);
    this->genome = genome;

    // Write out SAM header
    char *headerBuffer = new char[HEADER_BUFFER_SIZE];
    size_t headerSize;
    if (!generateHeader(genome,headerBuffer,HEADER_BUFFER_SIZE,&headerSize, false)) {
        fprintf(stderr,"SimpleSAMWriter: unable to generate SAM header\n");
        return false;
    }
    fprintf(file, "%s", headerBuffer);
    delete[] headerBuffer;

    return true;
}


bool SimpleSAMWriter::write(Read *read, AlignmentResult result, unsigned genomeLocation, bool isRC)
{
    if (file == NULL) {
        return false;
    }

    write(read, result, genomeLocation, isRC, false, false, NULL, NotFound, 0, false);
    return true;
}


bool SimpleSAMWriter::writePair(Read *read0, Read *read1, PairedAlignmentResult *result)
{
    if (file == NULL) {
        return false;
    }

    write(read0, result->status[0], result->location[0], result->isRC[0], true, true,
          read1, result->status[1], result->location[1], result->isRC[1]);
    write(read1, result->status[1], result->location[1], result->isRC[1], true, false,
          read0, result->status[0], result->location[0], result->isRC[0]);
    return true;
}


void SimpleSAMWriter::write(
        Read *read,
        AlignmentResult result,
        unsigned genomeLocation,
        bool isRC,
        bool hasMate,
        bool firstInPair,
        Read *mate,
        AlignmentResult mateResult,
        unsigned mateLocation,
        bool mateIsRC)
{
    const unsigned maxLineLength = 25000;
    char outputBuffer[maxLineLength];
    size_t outputBufferUsed;

    if (!generateSAMText(read,result,genomeLocation,isRC,hasMate,firstInPair,mate,mateResult,
            mateLocation,mateIsRC, genome, &lv, outputBuffer, maxLineLength,&outputBufferUsed)) {
        fprintf(stderr,"SimpleSAMWriter: tried to generate too long of a SAM line (> %d)\n",maxLineLength);
        exit(1);
    }
    
    if (1 != fwrite(outputBuffer,outputBufferUsed,1,file)) {
        fprintf(stderr,"Unable to write to SAM file.\n");
        exit(1);
    }
}


bool SimpleSAMWriter::close()
{
    if (file == NULL) {
        return false;
    }
    fclose(file);
    BigDealloc(buffer);
    file = NULL;
    buffer = NULL;
    return true;
}

ThreadSAMWriter::ThreadSAMWriter(size_t i_bufferSize)
    : remainingBufferSpace(i_bufferSize), bufferBeingCreated(0), bufferSize(i_bufferSize)
{
    buffer[0] = NULL;
    buffer[1] = NULL;
}

    bool
ThreadSAMWriter::initialize(AsyncFile* file, const Genome *i_genome, volatile _int64 *i_nextWriteOffset)
{
    genome = i_genome;
    nextWriteOffset = i_nextWriteOffset;
    buffer[0] = (char *)BigAlloc(bufferSize);
    buffer[1] = (char *)BigAlloc(bufferSize);
    writer[0] = file->getWriter();
    writer[1] = file->getWriter();

    if (NULL == buffer[0] || NULL == buffer[1] || NULL == writer[0] || NULL == writer[1]) {
        fprintf(stderr,"ThreadSAMWriter: failed to initialize\n");
        return false;
    }
    return true;
}
    
ThreadSAMWriter::~ThreadSAMWriter()
{
    BigDealloc(buffer[0]);
    BigDealloc(buffer[1]);
    delete writer[0];
    delete writer[1];
}

    bool
ThreadSAMWriter::close()
{
    if (remainingBufferSpace != bufferSize) {
        if (!startIo()) {
            fprintf(stderr,"WindowsSAMWriter::close(): startIo failed\n");
            return false;
        }

        if (!waitForIoCompletion()) {
            fprintf(stderr,"WindowsSAMWriter::close(): waitForIoCompletion failed\n");
            return false;
        }
    }
    bool ok = writer[0]->close();
    ok &= writer[1]->close();
    if (! ok) {
        fprintf(stderr, "WindowsSAMWriter::close closing writer failed\n");
    }

    return true;
}
    
    bool
ThreadSAMWriter::write(Read *read, AlignmentResult result, unsigned genomeLocation, bool isRC)
{
    size_t sizeUsed;
    if (!generateSAMText(read,result,genomeLocation,isRC,false,true,NULL,UnknownAlignment,0,false,genome,&lv,
            buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace,remainingBufferSpace, &sizeUsed)) {

        if (!startIo()) {
            return false;
        }

        if (!generateSAMText(read,result,genomeLocation,isRC,false,true,NULL,UnknownAlignment,0,false,genome,&lv,
                buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace,remainingBufferSpace, &sizeUsed)) {

            fprintf(stderr,"WindowsSAMWriter: create SAM string into fresh buffer failed\n");
            return false;
        }
    }
    size_t bufferOffset = bufferSize - remainingBufferSpace;
    remainingBufferSpace -= sizeUsed;
    afterWrite(result != NotFound ? genomeLocation : UINT32_MAX, bufferOffset, sizeUsed);
    return true;
}

    bool
ThreadSAMWriter::writePair(Read *read0, Read *read1, PairedAlignmentResult *result)
{
    //
    // We need to write both halves of the pair into the same buffer, so that a write from
    // some other thread doesn't separate them.  So, try the writes and if either doesn't
    // work start IO and try again.
    //
    size_t sizeUsed[2];
    bool writesFit = generateSAMText(read0,result->status[0],result->location[0],result->isRC[0],true,true,
                                     read1,result->status[1],result->location[1],result->isRC[1],
                        genome, &lv, buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace,remainingBufferSpace,&sizeUsed[0]);
    if (writesFit) {
        writesFit = generateSAMText(read1,result->status[1],result->location[1],result->isRC[1],true,false,
                                    read0,result->status[0],result->location[0],result->isRC[0],
                        genome, &lv, buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace + sizeUsed[0],remainingBufferSpace-sizeUsed[0],&sizeUsed[1]);
    }

    if (!writesFit) {
        if (!startIo()) {
            return false;
        }

        if (!generateSAMText(read0,result->status[0],result->location[0],result->isRC[0],true,true,
                             read1,result->status[1],result->location[1],result->isRC[1],
                genome, &lv, buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace,remainingBufferSpace,&sizeUsed[0]) ||
            !generateSAMText(read1,result->status[1],result->location[1],result->isRC[1],true,false,
                             read0,result->status[0],result->location[0],result->isRC[0],
                genome, &lv, buffer[bufferBeingCreated] + bufferSize - remainingBufferSpace + sizeUsed[0],remainingBufferSpace-sizeUsed[0],&sizeUsed[1])) {

                fprintf(stderr,"WindowsSAMWriter: create SAM string into fresh buffer failed\n");
                return false;
            }
    }

    size_t bufferOffset[2] = {bufferSize - sizeUsed[0] - sizeUsed[1], bufferSize - sizeUsed[1]};
    remainingBufferSpace -= (sizeUsed[0] + sizeUsed[1]);
    afterWrite(result->status[0] != NotFound ? result->location[0] : UINT32_MAX, bufferOffset[0], sizeUsed[0]);
    afterWrite(result->status[1] != NotFound ? result->location[1] : UINT32_MAX, bufferOffset[1], sizeUsed[1]);
    return true;
}

    ParallelSAMWriter *
ParallelSAMWriter::create(
    const char *fileName,
    const Genome *genome,
    unsigned nThreads,
    bool sort,
    size_t sortBufferMemory)
{
    ParallelSAMWriter *parallelWriter = sort
        ? new SortedParallelSAMWriter(sortBufferMemory)
        : new ParallelSAMWriter();
    if (!parallelWriter->initialize(fileName, genome, nThreads, sort)) {
        fprintf(stderr, "unable to initialize parallel SAM writer\n");
        delete parallelWriter;
        return NULL;
    }
    return parallelWriter;
}

    bool

ParallelSAMWriter::initialize(const char *fileName, const Genome *genome, unsigned i_nThreads, bool sorted)
{
    file = AsyncFile::open(fileName, true);
    if (NULL == file) {
        fprintf(stderr,"Unable to create SAM file '%s'\n",fileName);
        return false;
    }

    const size_t headerBufferSize = 20000;
    char headerBuffer[headerBufferSize];
    size_t headerActualSize;

    if (!SAMWriter::generateHeader(genome,headerBuffer,headerBufferSize,&headerActualSize, sorted)) {
        fprintf(stderr,"WindowsParallelSAMWriter: unable to generate SAM header.\n");
        return false;
    }

    AsyncFile::Writer* hwriter = file->getWriter();

    if (NULL == hwriter) {
        fprintf(stderr,"ParallelSAMWriter: unable to create writer\n");
        return false;
    }

    size_t bytesWritten;
    if (! hwriter->beginWrite(headerBuffer, headerActualSize, 0, &bytesWritten)) {
        fprintf(stderr,"ParallelSAMWriter: unable to write header to file\n");
        return false;
    }

    if (! hwriter->waitForCompletion()) {
        fprintf(stderr,"ParallelSAMWriter: failed to complete\n");
        return false;
    }
    if (! hwriter->close()) {
        fprintf(stderr, "ParallelSAMWriter: failed to close\n");
    }
    delete hwriter;

    nextWriteOffset = headerActualSize;

    nThreads = i_nThreads;
    writer = new ThreadSAMWriter *[nThreads];

    if (!createThreadWriters(genome)) {
        fprintf(stderr,"Unable to create SAM writer.\n");
        return false;
    }

    return true;
}

    bool
ParallelSAMWriter::createThreadWriters(const Genome* genome)
{
    bool worked = true;
    for (unsigned i = 0; i < nThreads; i++) {
        writer[i] = new ThreadSAMWriter(UnsortedBufferSize);
        worked &= writer[i]->initialize(file, genome, &nextWriteOffset);
    }
    return worked;
}

ParallelSAMWriter::~ParallelSAMWriter()
{
    delete [] writer;
}

    SAMWriter*
ParallelSAMWriter::getWriterForThread(
    unsigned whichThread)
{
    _ASSERT(whichThread < nThreads);
    return writer[whichThread];
}

    bool
ParallelSAMWriter::close()
{
    for (unsigned i = 0; i < nThreads; i++) {
        // writers were already closed by each thread so they could flush in parallel
        delete writer[i];
        writer[i] = NULL;
    }
    if (! file->close()) {
        fprintf(stderr, "ParallelSAMWriter::close file close failed\n");
    }
    delete file;
    file = NULL;
    return true;
}

    bool
ThreadSAMWriter::startIo()
{
    //
    // It didn't fit in the buffer.  Start writing it.
    //
    _int64 writeOffset = InterlockedAdd64AndReturnNewValue(nextWriteOffset,bufferSize - remainingBufferSpace) - (bufferSize - remainingBufferSpace);
    if (!beforeFlush(writeOffset, bufferSize - remainingBufferSpace)) {
        fprintf(stderr, "ThreadSAMWriter: beforeFlush failed\n");
        return false;
    }
    if (!writer[bufferBeingCreated]->beginWrite(buffer[bufferBeingCreated], bufferSize - remainingBufferSpace, writeOffset, NULL)) {
        fprintf(stderr,"ThreadSAMWriter: WriteFile failed\n");
        return false;
    }

    //
    // If necessary, wait for the other buffer to finish writing.
    //
    if (!waitForIoCompletion()) {
        fprintf(stderr,"ThreadSAMWriter: waitForIoCompletion failed\n");
        return false;
    }
    bufferBeingCreated = 1 - bufferBeingCreated;
    remainingBufferSpace = bufferSize;

    return true;
}

    bool
ThreadSAMWriter::waitForIoCompletion()
{
    return writer[1 - bufferBeingCreated]->waitForCompletion();
}

SortBlock::SortBlock()
    : entries(0), fileOffset(0), fileBytes(0), index(0), reader()
{
}

SortBlock::SortBlock(size_t capacity)
    : entries(capacity), fileOffset(0), fileBytes(0), index(0), reader()
{
}

SortBlock::SortBlock(
    SortBlock& other)
{
    entries = other.entries;
    fileOffset = other.fileOffset;
    fileBytes = other.fileBytes;
    index = other.index;
    reader = other.reader;
}

    void
SortBlock::operator=(
    SortBlock& other)
{
    entries = other.entries;
    fileOffset = other.fileOffset;
    fileBytes = other.fileBytes;
    index = other.index;
    index = other.index;
    reader = other.reader;
}

SortedThreadSAMWriter::SortedThreadSAMWriter(size_t i_bufferSize)
    : ThreadSAMWriter(i_bufferSize),
        parent(NULL),
        largest(1000),
        locations(1000)
{
}

SortedThreadSAMWriter::~SortedThreadSAMWriter()
{
}

    bool
SortedThreadSAMWriter::initialize(
    SortedParallelSAMWriter* i_parent,
    const Genome* i_genome)
{
    parent = i_parent;
    return ThreadSAMWriter::initialize(parent->file, i_genome, &parent->nextWriteOffset);
}

    void
SortedThreadSAMWriter::afterWrite(
    unsigned location,
    size_t bufferOffset,
    unsigned length)
{
    SortEntry entry(bufferOffset, length, location);
    locations.entries.push_back(entry);
}

    bool
SortedThreadSAMWriter::beforeFlush(_int64 fileOffset, _int64 length)
{
    // sort buffered reads by location for later merge sort
    std::sort(locations.entries.begin(), locations.entries.end(), SortEntry::comparator);
    
    // wait for IO of other buffer to finish since we're going to sort into there
    if (! waitForIoCompletion()) {
        fprintf(stderr, "SortedThreadSAMWriter waitForIoCompletion failed\n");
        return false;
    }
    
    // copy into other buffer in sorted order & switch buffers
    unsigned target = 0;
    for (VariableSizeVector<SortEntry>::iterator i = locations.entries.begin(); i != locations.entries.end(); i++) {
        memcpy(buffer[1 - bufferBeingCreated] + target, buffer[bufferBeingCreated] + i->offset, i->length);
        i->offset = fileOffset + target;
        target += i->length;
    }
    bufferBeingCreated = 1 - bufferBeingCreated;
    
    // remember offsets for full-file sort, get a new vector of appropriate size
    locations.fileOffset = fileOffset;
    locations.fileBytes = length;
    parent->addLocations(locations);
    if (locations.entries.size() > largest) {
        largest = (locations.entries.size() * 6) / 5; // grow by 20% if it exceeds prior max
    }
    locations.entries.reserve(largest);

    return true;
}

    bool
SortedParallelSAMWriter::initialize(
    const char *fileName,
    const Genome *genome,
    unsigned i_nThreads,
    bool sorted)
{
    InitializeExclusiveLock(&lock);
    sortedFile = fileName;
    tempFile = (char*) malloc(strlen(fileName) + 5);
    strcpy(tempFile, fileName);
    strcat(tempFile, ".tmp");
    bool ok = ParallelSAMWriter::initialize(tempFile, genome, i_nThreads, sorted);
    headerSize = nextWriteOffset;
    return ok;
}
    
    bool
SortedParallelSAMWriter::createThreadWriters(const Genome* genome)
{
    size_t bufferSize = totalMemory / nThreads / 2;
    bool worked = true;
    for (unsigned i = 0; i < nThreads; i++) {
        SortedThreadSAMWriter* w = new SortedThreadSAMWriter(bufferSize);
        writer[i] = w;
        worked &= w->initialize(this, genome);
    }
    return worked;
}

    bool
SortedParallelSAMWriter::close()
{
    if (! ParallelSAMWriter::close()) {
        return false;
    }
    DestroyExclusiveLock(&lock);

    // merge sort from temp file into sorted file
    printf("sorting...");
    _int64 start = timeInMillis();

    // first replace offset in entries with block index, and sort them all in one large array
    size_t total = 0;
    for (VariableSizeVector<SortBlock>::iterator i = locations.begin(); i != locations.end(); i++) {
        total += i->entries.size();
    }
    SortEntry* entries = (SortEntry*) BigAlloc(total * sizeof(SortEntry));
    size_t offset = 0;
    for (VariableSizeVector<SortBlock>::iterator i = locations.begin(); i != locations.end(); i++) {
        unsigned n = i->entries.size();
        unsigned blockIndex = i - locations.begin();
        memcpy(entries + offset, i->entries.begin(), n * (size_t) sizeof(SortEntry));
        for (unsigned j = 0; j < n; j++) {
            entries[offset + j].offset = blockIndex;
        }
        offset += i->entries.size();
    }
    std::sort(entries, entries + total, SortEntry::comparator);
    printf(" %ld s\nwriting sorted reads...", (timeInMillis() - start) / 1000);
    start = timeInMillis();

    // setup - open all files, read first block, begin read for second
    AsyncFile* temp = AsyncFile::open(tempFile, false);
    const size_t ReadBufferSize = UnsortedBufferSize;
    char* buffers = (char*) BigAlloc(ReadBufferSize * 2 * locations.size());
    unsigned j = 0;
    if (timeInMillis() - start > 1000) {
        printf(" (allocated %lld Mb in %lld s)",
            ReadBufferSize * 2 * locations.size() / (2 << 20), (timeInMillis() - start) / 1000);
    }
    for (VariableSizeVector<SortBlock>::iterator i = locations.begin(); i != locations.end(); i++, j++) {
        i->reader.open(temp, i->fileOffset, i->fileBytes, ReadBufferSize, true,
            buffers + j * 2 * ReadBufferSize, buffers + (j * 2 + 1) * ReadBufferSize);
    }
    for (VariableSizeVector<SortBlock>::iterator i = locations.begin(); i != locations.end(); i++) {
        i->reader.endOpen();
{'fraction_non_alphanumeric': 0.11328732097795752, 'fraction_numerical': 0.008890522903123381, 'mean_word_length': 2.572110792741165, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 239, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 8, 'word_list_counts': {'cursed_substrings.json': 22, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6134040', 'n_tokens_mistral': 18181, 'n_tokens_neox': 17186, 'n_words': 4516}