Fastx.hpp

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <concepts>
#include <cstdint>
#include <filesystem>
#include <format>
#include <fstream>
#include <istream>
#include <memory>
#include <span>
#include <stdexcept>
#include <string>
#include <string_view>
#include <utility>
#include <vector>
 
#include <cusbf/error.hpp>
#include <cusbf/gzstreambuf.hpp>
 
namespace cusbf {
 
struct RecordRange {
    uint64_t sequenceOffset{};
    uint64_t sequenceBytes{};
};
 
struct RecordBatchView {
    std::string_view sequence{};
    std::span<const RecordRange> records{};
};
 
class DenseRecordBatchBuilder {
   public:
    explicit DenseRecordBatchBuilder(uint64_t reserve_bytes = 0) {
        if (reserve_bytes != 0) {
            sequence_.reserve(static_cast<size_t>(reserve_bytes));
        }
    }
 
    [[nodiscard]] std::string_view sequence_view() const noexcept {
        return external_sequence_.empty() ? std::string_view{sequence_} : external_sequence_;
    }
 
    [[nodiscard]] RecordBatchView view() const noexcept {
        return RecordBatchView{
            sequence_view(),
            std::span<const RecordRange>{ranges_.data(), ranges_.size()},
        };
    }
 
    void appendRecord(std::string_view record_sequence) {
        assert(external_sequence_.empty() && "cannot appendRecord when external_sequence_ is set");
        ranges_.push_back(
            RecordRange{
                static_cast<uint64_t>(sequence_.size()),
                static_cast<uint64_t>(record_sequence.size()),
            }
        );
        sequence_.append(record_sequence);
    }
 
    void push_range(RecordRange range) {
        ranges_.push_back(range);
    }
 
    [[nodiscard]] std::string& sequence_buffer() noexcept {
        return sequence_;
    }
 
    [[nodiscard]] std::string_view& external_sequence_slot() noexcept {
        return external_sequence_;
    }
 
    [[nodiscard]] bool empty() const noexcept {
        return ranges_.empty();
    }
 
    [[nodiscard]] uint64_t recordCount() const noexcept {
        return static_cast<uint64_t>(ranges_.size());
    }
 
    [[nodiscard]] uint64_t raw_sequence_bytes() const noexcept {
        if (!external_sequence_.empty()) {
            uint64_t total = 0;
            for (const RecordRange& range : ranges_) {
                total += range.sequenceBytes;
            }
            return total;
        }
        return static_cast<uint64_t>(sequence_.size());
    }
 
    [[nodiscard]] const std::vector<RecordRange>& ranges() const noexcept {
        return ranges_;
    }
 
    void clear() {
        sequence_.clear();
        ranges_.clear();
        external_sequence_ = {};
    }
 
    void clear_and_shrink() {
        clear();
        std::string{}.swap(sequence_);
        std::vector<RecordRange>{}.swap(ranges_);
    }
 
   private:
    std::string sequence_;
    std::string_view external_sequence_{};
    std::vector<RecordRange> ranges_;
};
 
struct RecordQueryView {
    uint64_t record_index{};
    std::string_view sequence{};
    uint64_t queriedBases{};
    uint64_t queriedKmers{};
    uint64_t positive_kmers{};
    std::span<const uint8_t> hits{};
};
 
struct FastxInsertReport {
    uint64_t recordsIndexed{};
    uint64_t indexedBases{};
    uint64_t insertedKmers{};
};
 
struct FastxQueryReport {
    uint64_t recordsQueried{};
    uint64_t queriedBases{};
    uint64_t queriedKmers{};
    uint64_t positive_kmers{};
};
 
struct FastxRecordView {
    uint64_t record_index{};
    std::string_view header{};
    std::string_view sequence{};
    uint64_t queriedBases{};
    uint64_t queriedKmers{};
    uint64_t positive_kmers{};
    std::span<const uint8_t> hits{};
};
 
struct FastxDetailedQueryRecord {
    uint64_t record_index{};
    std::string header;
    std::string sequence;
    uint64_t queriedBases{};
    uint64_t queriedKmers{};
    uint64_t positive_kmers{};
    std::vector<uint8_t> hits;
};
 
struct FastxDetailedQueryReport {
    FastxQueryReport summary{};
    std::vector<FastxDetailedQueryRecord> records;
};
 
template <typename Functor>
concept RecordQueryConsumer = requires(Functor f, const RecordQueryView& view) {
    { f(view) } -> std::same_as<void>;
};
 
template <typename Functor>
concept FastxRecordConsumer = requires(Functor f, const FastxRecordView& record) {
    { f(record) } -> std::same_as<void>;
};
 
namespace detail {
 
enum class FastxFormat : uint8_t {
    unknown,
    fasta,
    fastq,
};
 
struct FastxRecord {
    std::string header;
    std::string sequence;
};
 
inline void trimTrailingCarriageReturn(std::string& line) {
    if (!line.empty() && line.back() == '\r') {
        line.pop_back();
    }
}
 
[[nodiscard]] inline uint32_t fastx_column_at(std::string_view line, size_t byte_index) {
    if (line.empty()) {
        return 1;
    }
    return static_cast<uint32_t>(std::min(byte_index, line.size() - 1) + 1);
}
 
[[nodiscard]] inline uint32_t fastx_quality_excess_column(
    uint64_t quality_length,
    uint64_t expected_length,
    std::string_view line
) {
    const uint64_t before = quality_length - line.size();
    const size_t byte_index = expected_length > before ? expected_length - before : line.size();
    return fastx_column_at(line, byte_index);
}
 
[[nodiscard]] inline uint32_t fastx_quality_short_column(std::string_view line) {
    return static_cast<uint32_t>(line.empty() ? 1 : line.size() + 1);
}
 
class FastxReader {
   public:
    explicit FastxReader(std::istream& input, std::string_view source_name = "<stream>")
        : input_(input), source_name_(source_name) {}
 
    [[nodiscard]] Result<bool> nextRecord(FastxRecord& record) {
        record.header.clear();
        record.sequence.clear();
 
        const auto header = readHeaderLine();
        if (!header) {
            return Err(header.error());
        }
        if (header->empty()) {
            return false;
        }
 
        const char headerTag = header->front();
 
        if (format_ == FastxFormat::unknown) {
            if (headerTag == '>') {
                format_ = FastxFormat::fasta;
            } else if (headerTag == '@') {
                format_ = FastxFormat::fastq;
            } else {
                return Err(
                    parseError("expected FASTA or FASTQ header", fastx_column_at(*header, 0))
                );
            }
        }
 
        const char expectedHeader = format_ == FastxFormat::fasta ? '>' : '@';
        if (headerTag != expectedHeader) {
            return Err(parseError(
                "mixed FASTA and FASTQ records are not supported", fastx_column_at(*header, 0)
            ));
        }
 
        record.header.assign(header->substr(1));
        if (format_ == FastxFormat::fasta) {
            CUSBF_TRY(readFastaSequence(record.sequence));
        } else {
            CUSBF_TRY(readFastqSequence(record.sequence));
        }
        return true;
    }
 
   private:
    std::istream& input_;
    std::string source_name_;
    std::string pendingHeader_;
    std::string lineBuffer_;
    FastxFormat format_{FastxFormat::unknown};
    uint64_t lineNumber_{};
 
    [[nodiscard]] Error parseError(std::string_view message, uint32_t column) const {
        return Error::fastx_parse(
            SourceLocation::fastx(source_name_, static_cast<uint32_t>(lineNumber_), column), message
        );
    }
 
    [[nodiscard]] Result<std::string_view> readHeaderLine() {
        if (!pendingHeader_.empty()) {
            lineBuffer_ = std::move(pendingHeader_);
            pendingHeader_.clear();
            return std::string_view{lineBuffer_};
        }
 
        while (std::getline(input_, lineBuffer_)) {
            ++lineNumber_;
            trimTrailingCarriageReturn(lineBuffer_);
            if (!lineBuffer_.empty()) {
                return std::string_view{lineBuffer_};
            }
        }
 
        if (input_.bad()) {
            return Err(
                Error::io(std::format("Failed to read FASTA/FASTQ input from {}", source_name_))
            );
        }
        return std::string_view{};
    }
 
    [[nodiscard]] Result<void> readFastaSequence(std::string& sequence) {
        sequence.reserve(4096);
        while (std::getline(input_, lineBuffer_)) {
            ++lineNumber_;
            trimTrailingCarriageReturn(lineBuffer_);
            if (!lineBuffer_.empty() && lineBuffer_.front() == '>') {
                pendingHeader_ = std::move(lineBuffer_);
                return {};
            }
            sequence += lineBuffer_;
        }
 
        if (input_.bad()) {
            return Err(
                Error::io(std::format("Failed to read FASTA/FASTQ input from {}", source_name_))
            );
        }
        return {};
    }
 
    [[nodiscard]] Result<void> readFastqSequence(std::string& sequence) {
        sequence.reserve(4096);
        while (std::getline(input_, lineBuffer_)) {
            ++lineNumber_;
            trimTrailingCarriageReturn(lineBuffer_);
            if (!lineBuffer_.empty() && lineBuffer_.front() == '+') {
                CUSBF_TRY(readFastqQualities(sequence.size()));
                return {};
            }
            sequence += lineBuffer_;
        }
 
        if (input_.bad()) {
            return Err(
                Error::io(std::format("Failed to read FASTA/FASTQ input from {}", source_name_))
            );
        }
        return Err(parseError(
            "unterminated FASTQ record: missing '+' separator",
            fastx_column_at(lineBuffer_, lineBuffer_.size() > 0 ? lineBuffer_.size() - 1 : 0)
        ));
    }
 
    [[nodiscard]] Result<void> readFastqQualities(uint64_t expectedLength) {
        uint64_t qualityLength = 0;
        uint32_t short_column = 1;
        while (qualityLength < expectedLength && std::getline(input_, lineBuffer_)) {
            ++lineNumber_;
            trimTrailingCarriageReturn(lineBuffer_);
            short_column = fastx_quality_short_column(lineBuffer_);
            qualityLength += lineBuffer_.size();
            if (qualityLength > expectedLength) {
                return Err(parseError(
                    "FASTQ quality length exceeds sequence length",
                    fastx_quality_excess_column(qualityLength, expectedLength, lineBuffer_)
                ));
            }
        }
 
        if (qualityLength == expectedLength) {
            return {};
        }
        if (input_.bad()) {
            return Err(
                Error::io(std::format("Failed to read FASTA/FASTQ input from {}", source_name_))
            );
        }
        return Err(parseError("FASTQ quality length does not match sequence length", short_column));
    }
};
 
[[nodiscard]] inline Result<std::unique_ptr<std::istream>> openFastxFile(
    const std::filesystem::path& path
) {
    if (isGzipFile(path)) {
        return GzIstream::open(path);
    }
    auto input = std::make_unique<std::ifstream>(path);
    if (!input->is_open()) {
        return Err(Error::io(std::format("Failed to open FASTA/FASTQ file: {}", path.string())));
    }
    return input;
}
 
}  // namespace detail
 
}  // namespace cusbf