a48.hpp

Go to the documentation of this file.

#pragma once
 
#include <atomic>
#include <cstdint>
#include <exception>
#include <optional>
#include <string>
#include <string_view>
#include <thread>
#include <vector>
 
#include <cuddl/error.hpp>
 
namespace cuddl::a48 {
 
struct database_metadata {
    uint32_t kmer_length{};
    bool has_kmer_length{};
    uint64_t seed{};
    bool has_seed{};
    uint8_t exponent_bits{6U};
    bool has_exponent{};
    std::string blacklist;
    size_t record_count{};
};
 
struct record_metadata {
    int64_t id{-1};
    int64_t tax_id{-1};
    std::string name;
    std::string filename;
    int64_t bases{-1};
    int32_t contigs{-1};
    float gc{-1.0f};
    std::string origin;
    std::string lineage;
    uint32_t len{};
    int32_t offset{-1};
};
 
struct record {
    uint32_t ordinal{};
    record_metadata metadata;
    std::vector<uint16_t> scores;
};
 
struct database {
    database_metadata metadata;
    std::vector<record> records;
};
 
[[nodiscard]] inline std::optional<uint64_t> decode_a48_token(std::string_view token) {
    if (token.empty()) {
        return uint64_t{0};
    }
    uint64_t value = 0;
    for (char const character : token) {
        auto const code = static_cast<uint8_t>(character);
        if (code < 48U || code > 111U) {
            return std::nullopt;
        }
        value = (value << 6U) | static_cast<uint64_t>(code - 48U);
    }
    return value;
}
 
[[nodiscard]] inline std::string encode_a48_token(uint64_t value) {
    if (value == 0U) {
        return std::string(1, '0');
    }
    std::string reversed;
    while (value != 0U) {
        reversed.push_back(static_cast<char>(48U + (value & 0x3fU)));
        value >>= 6U;
    }
    return {reversed.rbegin(), reversed.rend()};
}
 
[[nodiscard]] inline uint16_t
promote_score(uint16_t loaded, int32_t offset, uint8_t exponent_bits) {
    if (loaded == 0U || offset < 0) {
        return loaded;
    }
    auto const mantissa_bits = 16U - exponent_bits;
    auto const mask = (1U << mantissa_bits) - 1U;
    auto const rel_nlz = static_cast<uint32_t>(loaded >> mantissa_bits);
    auto const abs_nlz = rel_nlz + static_cast<uint32_t>(offset);
    return static_cast<uint16_t>((abs_nlz << mantissa_bits) | (loaded & mask));
}
 
[[nodiscard]] inline std::vector<std::string_view>
split_fields(std::string_view text, char separator, bool keep_trailing_empty = false) {
    std::vector<std::string_view> fields;
    size_t begin = 0;
    while (true) {
        auto const end = text.find(separator, begin);
        if (end == std::string_view::npos) {
            fields.push_back(text.substr(begin));
            break;
        }
        fields.push_back(text.substr(begin, end - begin));
        begin = end + 1U;
    }
    if (!keep_trailing_empty && !fields.empty() && fields.back().empty()) {
        fields.pop_back();
    }
    return fields;
}
 
namespace detail {
 
template <typename T, typename Parser>
[[nodiscard]] Result<T> parse_field(Parser&& parser, std::string_view value, char const* field) {
    try {
        return Result<T>::ok(static_cast<T>(parser(std::string(value))));
    } catch (std::exception const&) {
        return Err(Error::invalid_argument(std::string("malformed A48 ") + field + " value"));
    }
}
 
}  // namespace detail
 
[[nodiscard]] inline bool is_record_header(std::string_view key) {
    return key == "id" || key == "tid" || key == "name" || key == "file" || key == "bases" ||
           key == "contigs" || key == "gc" || key == "origin" || key == "lineage" || key == "len" ||
           key == "offset";
}
 
[[nodiscard]] inline Result<record_metadata> parse_record_headers(std::string_view block) {
    record_metadata meta;
    for (auto const line : split_fields(block, '\n')) {
        if (line.empty() || line.front() != '#') {
            continue;
        }
        auto const fields = split_fields(line, '\t');
        if (fields.empty()) {
            continue;
        }
        auto key = fields.front();
        key.remove_prefix(1);  // drop '#'
        if (fields.size() < 2U) {
            continue;
        }
        auto const value = fields[1];
        if (key == "id") {
            auto const parsed = detail::parse_field<int64_t>(
                [](std::string const& s) { return std::stoll(s); }, value, "id"
            );
            if (!parsed) {
                return Err(parsed.error());
            }
            meta.id = *parsed;
        } else if (key == "tid") {
            auto const parsed = detail::parse_field<int64_t>(
                [](std::string const& s) { return std::stoll(s); }, value, "tid"
            );
            if (!parsed) {
                return Err(parsed.error());
            }
            meta.tax_id = *parsed;
        } else if (key == "name") {
            meta.name = std::string(value);
        } else if (key == "file") {
            meta.filename = std::string(value);
        } else if (key == "bases") {
            auto const parsed = detail::parse_field<int64_t>(
                [](std::string const& s) { return std::stoll(s); }, value, "bases"
            );
            if (!parsed) {
                return Err(parsed.error());
            }
            meta.bases = *parsed;
        } else if (key == "contigs") {
            auto const parsed = detail::parse_field<int32_t>(
                [](std::string const& s) { return std::stol(s); }, value, "contigs"
            );
            if (!parsed) {
                return Err(parsed.error());
            }
            meta.contigs = *parsed;
        } else if (key == "gc") {
            auto const parsed = detail::parse_field<float>(
                [](std::string const& s) { return std::stof(s); }, value, "gc"
            );
            if (!parsed) {
                return Err(parsed.error());
            }
            meta.gc = *parsed;
        } else if (key == "origin") {
            meta.origin = std::string(value);
        } else if (key == "lineage") {
            meta.lineage = std::string(value);
        } else if (key == "len") {
            auto const parsed = detail::parse_field<uint32_t>(
                [](std::string const& s) { return std::stoul(s); }, value, "len"
            );
            if (!parsed) {
                return Err(parsed.error());
            }
            meta.len = *parsed;
        } else if (key == "offset") {
            auto const parsed = detail::parse_field<int32_t>(
                [](std::string const& s) { return std::stol(s); }, value, "offset"
            );
            if (!parsed) {
                return Err(parsed.error());
            }
            meta.offset = *parsed;
        }
    }
    return Result<record_metadata>::ok(std::move(meta));
}
 
[[nodiscard]] inline Result<std::vector<uint16_t>>
decode_a48_row(std::string_view row, record_metadata const& metadata, uint8_t exponent_bits) {
    auto const fields = split_fields(row, '\t', true);
    std::vector<uint16_t> scores;
    scores.reserve(fields.size());
    for (auto const field : fields) {
        auto const loaded = decode_a48_token(field);
        if (!loaded) {
            return Err(
                Error::invalid_argument(
                    "A48 data row contains a character outside the A48 alphabet"
                )
            );
        }
        scores.push_back(
            promote_score(static_cast<uint16_t>(*loaded), metadata.offset, exponent_bits)
        );
    }
    return Result<std::vector<uint16_t>>::ok(std::move(scores));
}
 
[[nodiscard]] inline Result<database> decode_a48_tsv(std::string_view input) {
    database result;
    uint8_t exponent_bits = 6U;
    bool has_exponent = false;
    size_t expected_bucket_count = 0;
    bool bucket_count_set = false;
 
    std::string header_block;
    size_t ordinal = 0;
    size_t begin = 0;
    while (begin <= input.size()) {
        auto const end = input.find('\n', begin);
        auto line =
            input.substr(begin, end == std::string_view::npos ? input.size() - begin : end - begin);
        if (!line.empty() && line.back() == '\r') {
            line.remove_suffix(1);
        }
 
        if (line.empty()) {
            if (end == std::string_view::npos) {
                break;
            }
            begin = end + 1U;
            continue;
        }
 
        if (line.front() == '#') {
            auto const fields = split_fields(line, '\t');
            if (fields.empty()) {
                break;
            }
            auto key = fields.front();
            key.remove_prefix(1);  // drop '#'
            auto const value = fields.size() > 1U ? fields[1] : std::string_view{};
 
            if (key == "k") {
                auto const parsed = detail::parse_field<uint32_t>(
                    [](std::string const& s) { return std::stoul(s); }, value, "k"
                );
                if (!parsed) {
                    return Err(parsed.error());
                }
                result.metadata.kmer_length = *parsed;
                result.metadata.has_kmer_length = true;
            } else if (key == "seed") {
                auto const parsed = detail::parse_field<uint64_t>(
                    [](std::string const& s) { return std::stoull(s); }, value, "seed"
                );
                if (!parsed) {
                    return Err(parsed.error());
                }
                result.metadata.seed = *parsed;
                result.metadata.has_seed = true;
            } else if (key == "exponent") {
                auto const parsed = detail::parse_field<uint32_t>(
                    [](std::string const& s) { return std::stoul(s); }, value, "exponent"
                );
                if (!parsed) {
                    return Err(parsed.error());
                }
                if (*parsed < 1U || *parsed > 15U) {
                    return Err(Error::invalid_argument("A48 exponent must be in [1, 15]"));
                }
                exponent_bits = *parsed;
                result.metadata.exponent_bits = *parsed;
                result.metadata.has_exponent = true;
                has_exponent = true;
            } else if (key == "blacklist") {
                result.metadata.blacklist = std::string(value);
            }
 
            if (is_record_header(key)) {
                if (!header_block.empty()) {
                    header_block.push_back('\n');
                }
                header_block.append(line);
            }
 
            if (end == std::string_view::npos) {
                break;
            }
            begin = end + 1U;
            continue;
        }
 
        // A data row terminates the currently accumulated record header block.
        if (header_block.empty()) {
            return Err(Error::invalid_argument("A48 data row appears before any record header"));
        }
        auto meta = CUDDL_TRY(parse_record_headers(header_block));
        header_block.clear();
 
        if (meta.offset >= 0 && !has_exponent) {
            return Err(
                Error::invalid_argument("relative-encoded A48 record requires a #exponent header")
            );
        }
        auto scores = CUDDL_TRY(decode_a48_row(line, meta, exponent_bits));
        if (scores.empty()) {
            return Err(Error::invalid_argument("A48 record has no bucket values"));
        }
        if (!bucket_count_set) {
            expected_bucket_count = scores.size();
            bucket_count_set = true;
        } else if (scores.size() != expected_bucket_count) {
            return Err(Error::invalid_argument("A48 data row width changed between records"));
        }
        if (meta.len != 0U && meta.len != scores.size()) {
            return Err(Error::invalid_argument("#len disagrees with the A48 data row width"));
        }
 
        record decoded;
        decoded.ordinal = static_cast<uint32_t>(ordinal);
        decoded.metadata = meta;
        decoded.scores = std::move(scores);
        result.records.push_back(std::move(decoded));
        ++ordinal;
 
        if (end == std::string_view::npos) {
            break;
        }
        begin = end + 1U;
    }
 
    if (!header_block.empty()) {
        return Err(Error::invalid_argument("A48 file ends inside a record header block"));
    }
    if (result.records.empty()) {
        return Err(Error::invalid_argument("A48 file contains no records"));
    }
 
    result.metadata.record_count = result.records.size();
    return Result<database>::ok(std::move(result));
}
 
[[nodiscard]] inline Result<database> decode_a48_tsv_parallel(
    std::string_view input,
    unsigned requested_threads = std::thread::hardware_concurrency()
) {
    struct record_job {
        std::string header;
        std::string_view data;
        uint32_t ordinal{};
    };
 
    database result;
    std::vector<record_job> jobs;
    uint8_t exponent_bits = 6U;
    bool has_exponent = false;
    size_t expected_bucket_count = 0;
    bool bucket_count_set = false;
 
    std::string header_block;
    size_t ordinal = 0;
    size_t begin = 0;
    while (begin <= input.size()) {
        auto const end = input.find('\n', begin);
        auto line =
            input.substr(begin, end == std::string_view::npos ? input.size() - begin : end - begin);
        if (!line.empty() && line.back() == '\r') {
            line.remove_suffix(1);
        }
        if (line.empty()) {
            if (end == std::string_view::npos) {
                break;
            }
            begin = end + 1U;
            continue;
        }
 
        if (line.front() == '#') {
            auto const fields = split_fields(line, '\t');
            if (fields.empty()) {
                break;
            }
            auto key = fields.front();
            key.remove_prefix(1);
            auto const value = fields.size() > 1U ? fields[1] : std::string_view{};
            if (key == "k") {
                auto const parsed = detail::parse_field<uint32_t>(
                    [](std::string const& s) { return std::stoul(s); }, value, "k"
                );
                if (!parsed) {
                    return Err(parsed.error());
                }
                result.metadata.kmer_length = *parsed;
                result.metadata.has_kmer_length = true;
            } else if (key == "seed") {
                auto const parsed = detail::parse_field<uint64_t>(
                    [](std::string const& s) { return std::stoull(s); }, value, "seed"
                );
                if (!parsed) {
                    return Err(parsed.error());
                }
                result.metadata.seed = *parsed;
                result.metadata.has_seed = true;
            } else if (key == "exponent") {
                auto const parsed = detail::parse_field<uint32_t>(
                    [](std::string const& s) { return std::stoul(s); }, value, "exponent"
                );
                if (!parsed) {
                    return Err(parsed.error());
                }
                if (*parsed < 1U || *parsed > 15U) {
                    return Err(Error::invalid_argument("A48 exponent must be in [1, 15]"));
                }
                exponent_bits = *parsed;
                result.metadata.exponent_bits = *parsed;
                result.metadata.has_exponent = true;
                has_exponent = true;
            } else if (key == "blacklist") {
                result.metadata.blacklist = std::string(value);
            }
 
            if (is_record_header(key)) {
                if (!header_block.empty()) {
                    header_block.push_back('\n');
                }
                header_block.append(line);
            }
            if (end == std::string_view::npos) {
                break;
            }
            begin = end + 1U;
            continue;
        }
 
        if (header_block.empty()) {
            return Err(Error::invalid_argument("A48 data row appears before any record header"));
        }
        auto const record_ordinal = static_cast<uint32_t>(ordinal);
        if (record_ordinal == 0U) {
            auto meta = CUDDL_TRY(parse_record_headers(header_block));
            header_block.clear();
            if (meta.offset >= 0 && !has_exponent) {
                return Err(
                    Error::invalid_argument(
                        "relative-encoded A48 record requires a #exponent header"
                    )
                );
            }
            auto scores = CUDDL_TRY(decode_a48_row(line, meta, exponent_bits));
            if (scores.empty()) {
                return Err(Error::invalid_argument("A48 record has no bucket values"));
            }
            expected_bucket_count = scores.size();
            bucket_count_set = true;
            if (meta.len != 0U && meta.len != scores.size()) {
                return Err(Error::invalid_argument("#len disagrees with the A48 data row width"));
            }
            record decoded;
            decoded.ordinal = record_ordinal;
            decoded.metadata = std::move(meta);
            decoded.scores = std::move(scores);
            result.records.push_back(std::move(decoded));
        } else {
            jobs.push_back(
                record_job{std::move(header_block), line, static_cast<uint32_t>(ordinal)}
            );
            header_block.clear();
        }
        ++ordinal;
 
        if (end == std::string_view::npos) {
            break;
        }
        begin = end + 1U;
    }
 
    if (!header_block.empty()) {
        return Err(Error::invalid_argument("A48 file ends inside a record header block"));
    }
    if (result.records.empty()) {
        return Err(Error::invalid_argument("A48 file contains no records"));
    }
 
    result.records.resize(result.records.size() + jobs.size());
    auto const requested = requested_threads;
    auto const worker_count = static_cast<size_t>(std::max(1U, requested));
    std::vector<std::thread> workers;
    std::vector<std::exception_ptr> errors(worker_count);
    workers.reserve(worker_count);
 
    for (size_t worker = 0; worker < worker_count; ++worker) {
        auto const chunk_begin = jobs.size() * worker / worker_count;
        auto const chunk_end = jobs.size() * (worker + 1U) / worker_count;
        if (chunk_begin == chunk_end) {
            continue;
        }
        workers.emplace_back([&, chunk_begin, chunk_end, worker]() {
            try {
                for (size_t i = chunk_begin; i < chunk_end; ++i) {
                    auto const& job = jobs[i];
                    auto meta = parse_record_headers(job.header);
                    if (!meta) {
                        throw std::runtime_error(meta.error().message());
                    }
                    if (meta->offset >= 0 && !has_exponent) {
                        throw std::runtime_error(
                            "relative-encoded A48 record requires a #exponent header"
                        );
                    }
                    auto scores = decode_a48_row(job.data, *meta, exponent_bits);
                    if (!scores) {
                        throw std::runtime_error(scores.error().message());
                    }
                    if (scores->empty()) {
                        throw std::runtime_error("A48 record has no bucket values");
                    }
                    if (!bucket_count_set || scores->size() != expected_bucket_count) {
                        throw std::runtime_error("A48 data row width changed between records");
                    }
                    if (meta->len != 0U && meta->len != scores->size()) {
                        throw std::runtime_error("#len disagrees with the A48 data row width");
                    }
                    record decoded;
                    decoded.ordinal = job.ordinal;
                    decoded.metadata = std::move(*meta);
                    decoded.scores = std::move(*scores);
                    result.records[job.ordinal] = std::move(decoded);
                }
            } catch (...) {
                errors[worker] = std::current_exception();
            }
        });
    }
    for (auto& worker : workers) {
        worker.join();
    }
    for (auto const& error : errors) {
        if (error) {
            try {
                std::rethrow_exception(error);
            } catch (std::exception const& exception) {
                return Err(Error::invalid_argument(exception.what()));
            }
        }
    }
 
    result.metadata.record_count = result.records.size();
    return Result<database>::ok(std::move(result));
}
 
struct row_summary {
    uint32_t lower{};
    uint32_t equal{};
    uint32_t higher{};
    uint32_t both_empty{};
 
    friend bool operator==(row_summary const&, row_summary const&) = default;
};
 
struct oracle_result {
    std::vector<uint32_t> match_counts;
    std::vector<row_summary> summaries;
};
 
[[nodiscard]] inline oracle_result exhaustive_oracle(
    std::vector<uint16_t> const& rows,
    std::vector<uint16_t> const& query,
    size_t bucket_count,
    uint32_t minimum_matches
) {
    auto const reference_count = rows.size() / bucket_count;
    oracle_result result;
    result.match_counts.assign(reference_count, 0U);
    result.summaries.assign(reference_count, row_summary{});
    for (size_t reference = 0; reference < reference_count; ++reference) {
        auto summary = row_summary{};
        auto const* reference_row = rows.data() + reference * bucket_count;
        auto const* query_row = query.data();
        for (size_t bucket = 0; bucket < bucket_count; ++bucket) {
            auto const query_score = query_row[bucket];
            auto const reference_score = reference_row[bucket];
            if (query_score == reference_score) {
                if (query_score == 0U) {
                    ++summary.both_empty;
                } else {
                    ++summary.equal;
                    ++result.match_counts[reference];
                }
            } else if (query_score < reference_score) {
                ++summary.lower;
            } else {
                ++summary.higher;
            }
        }
        if (result.match_counts[reference] >= minimum_matches) {
            result.summaries[reference] = summary;
        }
    }
    return result;
}
 
[[nodiscard]] inline oracle_result exhaustive_oracle(
    std::vector<std::vector<uint16_t>> const& rows,
    std::vector<uint16_t> const& query,
    uint32_t minimum_matches
) {
    std::vector<uint16_t> flat;
    auto const bucket_count = rows.empty() ? size_t{0} : rows.front().size();
    flat.reserve(rows.size() * bucket_count);
    for (auto const& row : rows) {
        flat.insert(flat.end(), row.begin(), row.end());
    }
    return exhaustive_oracle(flat, query, bucket_count, minimum_matches);
}
 
}  // namespace cuddl::a48