reference_database_file.cuh

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <bit>
#include <cctype>
#include <cerrno>
#include <condition_variable>
#include <cstdio>
#include <cstdlib>
#include <exception>
#include <filesystem>
#include <fstream>
#include <limits>
#include <mutex>
#include <optional>
#include <span>
#include <string>
#include <system_error>
#include <thread>
#include <vector>
 
#include <fcntl.h>
#include <libdeflate.h>
#include <unistd.h>
 
#include <cub/device/device_reduce.cuh>
#include <cuda/iterator>
#include <cuda/std/functional>
 
#include <cuddl/detail/database_staging.cuh>
#include <cuddl/detail/sequence_encode.cuh>
#include <cuddl/fastx.hpp>
#include <cuddl/reference_database.cuh>
#include <cuddl/sketch.cuh>
 
namespace cuddl::detail {
 
// The file format uses little-endian integers, never native struct representations.
template <typename T>
inline T database_file_little_endian(T value) {
    if constexpr (std::endian::native == std::endian::little) {
        return value;
    } else {
        T result = 0;
        for (size_t i = 0; i < sizeof(T); ++i) {
            result = static_cast<T>((result << 8U) | (value & 0xffU));
            value >>= 8U;
        }
        return result;
    }
}
 
struct database_file_reader {
    std::ifstream input;
    uint64_t remaining{};
    // libdeflate's CRC-32 equals zlib's and uses carry-less multiplication where available.
    uint32_t checksum = 0;
 
    Result<void> bytes(void* destination, size_t size) {
        if (size > remaining ||
            size > static_cast<size_t>(std::numeric_limits<std::streamsize>::max())) {
            return Err(Error::invalid_argument("truncated reference database file"));
        }
        if (size != 0) {
            input.read(static_cast<char*>(destination), static_cast<std::streamsize>(size));
            if (!input) return Err(Error::resource("cannot read reference database file"));
            checksum = libdeflate_crc32(checksum, destination, size);
        }
        remaining -= size;
        return Ok();
    }
 
    template <typename T>
    Result<void> value(T& item) {
        CUDDL_TRY(bytes(&item, sizeof(item)));
        item = database_file_little_endian(item);
        return Ok();
    }
 
    template <typename T>
    Result<void> words(std::vector<T>& items) {
        CUDDL_TRY(bytes(items.data(), items.size() * sizeof(T)));
        if constexpr (std::endian::native != std::endian::little) {
            for (auto& item : items) {
                item = database_file_little_endian(item);
            }
        }
        return Ok();
    }
};
 
struct database_file_writer {
    std::string temporary;
    FILE* output{};
    uint32_t checksum = 0;
 
    ~database_file_writer() {
        if (output) std::fclose(output);
        if (!temporary.empty()) {
            ::unlink(temporary.c_str());
        }
    }
 
    Result<void> bytes(void const* source, size_t size) {
        if (size != 0) {
            if (std::fwrite(source, 1, size, output) != size) {
                return Err(Error::resource("cannot write reference database file"));
            }
            checksum = libdeflate_crc32(checksum, source, size);
        }
        return Ok();
    }
 
    template <typename T>
    Result<void> value(T item) {
        item = database_file_little_endian(item);
        return bytes(&item, sizeof(item));
    }
 
    template <typename T>
    Result<void> words(std::vector<T> const& items) {
        if constexpr (std::endian::native == std::endian::little) {
            return bytes(items.data(), items.size() * sizeof(T));
        } else {
            for (auto item : items) {
                CUDDL_TRY(value(item));
            }
            return Ok();
        }
    }
};
 
template <typename IO>
Result<void>
database_file_metadata(IO& io, reference_database_metadata& metadata, uint32_t& blacklist_count) {
    auto& c = metadata.compatibility;
    CUDDL_TRY(io.value(c.kmer_length));
    CUDDL_TRY(io.value(c.exponent_bits));
    CUDDL_TRY(io.value(c.bucket_count));
    CUDDL_TRY(io.value(c.key_mask));
    CUDDL_TRY(io.value(metadata.reference_count));
    CUDDL_TRY(io.value(blacklist_count));
    CUDDL_TRY(io.value(c.blacklist_identity));
    return Ok();
}
 
}  // namespace cuddl::detail
 
namespace cuddl {
 
using transfer_mode = detail::transfer_mode;
 
using decompression_backend = detail::decompression_backend;
 
using sequence_record = detail::sequence_record;
 
using sequence_genome = detail::sequence_genome;
 
[[nodiscard]] inline unsigned default_parser_workers() {
    auto const hardware = std::thread::hardware_concurrency();
    return hardware == 0U ? 8U : hardware;
}
 
struct path_build_options {
    reference_build_statistics* statistics = nullptr;
    unsigned parser_workers = default_parser_workers();
    std::optional<size_t> staging_bytes{};
    transfer_mode transfer = transfer_mode::automatic;
    decompression_backend decompression = decompression_backend::automatic;
    bool direct_reads = false;
    std::optional<std::reference_wrapper<device_blacklist const>> blacklist = std::nullopt;
};
 
struct sequence_build_options {
    reference_build_statistics* statistics = nullptr;
    std::optional<size_t> staging_bytes{};
    std::optional<std::reference_wrapper<device_blacklist const>> blacklist = std::nullopt;
};
 
[[nodiscard]] inline Result<std::vector<std::filesystem::path>> genome_paths(
    std::filesystem::path const& directory
) try {
    std::vector<std::filesystem::path> paths;
    for (auto const& entry : std::filesystem::recursive_directory_iterator(
             directory, std::filesystem::directory_options::skip_permission_denied
         )) {
        if (!entry.is_symlink() && !entry.is_regular_file()) continue;
        auto name = entry.path().filename().string();
        std::ranges::transform(name, name.begin(), [](unsigned char c) {
            return static_cast<char>(std::tolower(c));
        });
        auto extension = std::filesystem::path(name).extension().string();
        if (extension == ".gz" || extension == ".bgz" || extension == ".bgzf") {
            extension = std::filesystem::path(name).stem().extension().string();
        }
        if (extension == ".fa" || extension == ".fna" || extension == ".fasta" ||
            extension == ".ffn" || extension == ".frn" || extension == ".fq" ||
            extension == ".fastq") {
            paths.push_back(entry.path());
        }
    }
    std::ranges::sort(paths);
    return paths;
} catch (std::filesystem::filesystem_error const& error) {
    return Err(Error::resource(error.what()));
}
 
[[nodiscard]] inline Result<std::vector<std::filesystem::path>> list_genomes(
    std::filesystem::path const& input
) {
    if (std::filesystem::is_directory(input)) return genome_paths(input);
    std::ifstream list(input);
    if (!list) return Err(Error::resource("cannot read genome list " + input.string()));
    std::vector<std::filesystem::path> paths;
    for (std::string line; std::getline(list, line);) {
        if (!line.empty()) paths.emplace_back(line);
    }
    return paths;
}
 
namespace detail { struct reference_file_assembly; }  // namespace detail
 
class reference_database_file {
    static constexpr uint16_t database_file_version = 1U;
 
   public:
    [[nodiscard]] kmer_blacklist const& blacklist() const noexcept {
        return blacklist_;
    }
 
    [[nodiscard]] reference_database_metadata metadata() const noexcept {
        return metadata_;
    }
    [[nodiscard]] std::span<std::string const> names() const noexcept {
        return names_;
    }
    [[nodiscard]] std::span<uint16_t const> rows() const noexcept {
        return rows_;
    }
 
    template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
    [[nodiscard]] static Result<reference_database_file> build(
        std::span<std::filesystem::path const> paths,
        cuda::stream_ref stream,
        path_build_options options = {}
    ) try {
        // NVCC 13.3 crashes on CUDDL_TRY directly inside a try block. Keep its GNU statement
        // expressions in a separate lambda scope, outside the exception-catching function.
        return [&]() -> Result<reference_database_file> {
            using database_type = reference_database<K, BucketCount, Layout>;
            if (paths.size() > std::numeric_limits<uint32_t>::max() ||
                paths.size() >
                    std::numeric_limits<size_t>::max() / (BucketCount * sizeof(uint32_t))) {
                return Err(Error::resource("reference collection exceeds database capacity"));
            }
            reference_database_file result;
            result.metadata_ = {
                score_compatibility::current<K, BucketCount, Layout>(),
                static_cast<uint32_t>(paths.size())
            };
            CUDDL_TRY((result.set_blacklist<K, BucketCount>(options.blacklist, stream)));
            result.rows_.resize(paths.size() * BucketCount);
            result.names_.reserve(paths.size());
            // No inputs means no staging budget to resolve and nothing to encode.
            if (paths.empty()) return result;
            CUDDL_TRY((detail::stage_paths<K, BucketCount, Layout>(
                paths,
                stream,
                options.staging_bytes,
                options.parser_workers,
                options.transfer,
                options.decompression,
                options.direct_reads,
                options.statistics,
                result.download_rows<BucketCount>(stream),
                options.blacklist ? options.blacklist->get().view() : cuda::std::nullopt
            )));
            for (auto const& path : paths) result.names_.push_back(path.string());
            // Instantiate the same constraints as the destination GPU database.
            static_assert(sizeof(database_type) > 0);
            return result;
        }();
    } catch (cuda::cuda_error const& error) {
        return Err(Error::cuda(static_cast<cudaError_t>(error.status())));
    } catch (std::system_error const& error) {
        return Err(Error::resource(error.what()));
    } catch (std::bad_alloc const& error) {
        return Err(Error::resource(error.what()));
    }
 
    template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
    [[nodiscard]] static Result<reference_database_file> build_from_sequences(
        std::span<sequence_genome const> genomes,
        cuda::stream_ref stream,
        sequence_build_options options = {}
    ) try {
        return [&]() -> Result<reference_database_file> {
            using database_type = reference_database<K, BucketCount, Layout>;
            if (genomes.size() > std::numeric_limits<uint32_t>::max() ||
                genomes.size() >
                    std::numeric_limits<size_t>::max() / (BucketCount * sizeof(uint32_t))) {
                return Err(Error::resource("reference collection exceeds database capacity"));
            }
            reference_database_file result;
            result.metadata_ = {
                score_compatibility::current<K, BucketCount, Layout>(),
                static_cast<uint32_t>(genomes.size())
            };
            CUDDL_TRY((result.set_blacklist<K, BucketCount>(options.blacklist, stream)));
            result.rows_.resize(genomes.size() * BucketCount);
            result.names_.reserve(genomes.size());
            if (genomes.empty()) return result;
            CUDDL_TRY((detail::stage_sequences<K, BucketCount, Layout>(
                genomes,
                stream,
                options.staging_bytes,
                options.statistics,
                result.download_rows<BucketCount>(stream),
                options.blacklist ? options.blacklist->get().view() : cuda::std::nullopt
            )));
            for (auto const& genome : genomes) result.names_.emplace_back(genome.name);
            static_assert(sizeof(database_type) > 0);
            return result;
        }();
    } catch (cuda::cuda_error const& error) {
        return Err(Error::cuda(static_cast<cudaError_t>(error.status())));
    } catch (std::system_error const& error) {
        return Err(Error::resource(error.what()));
    } catch (std::bad_alloc const& error) {
        return Err(Error::resource(error.what()));
    }
 
    template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
    [[nodiscard]] static Result<reference_database_file> from_store(
        device_span<uint32_t const> store,
        std::span<std::string const> names,
        cuda::stream_ref stream,
        kmer_blacklist const& blacklist
    ) try {
        return [&]() -> Result<reference_database_file> {
            if (store.size() % BucketCount != 0) {
                return Err(Error::invalid_argument("a store must contain whole sketches"));
            }
            size_t const count = store.size() / BucketCount;
            if (!names.empty() && names.size() != count) {
                return Err(Error::invalid_argument("labels must be empty or name every reference"));
            }
            if (count != 0 && store.data() == nullptr) {
                return Err(Error::invalid_argument("a nonempty store must not be null"));
            }
            reference_database_file result;
            result.metadata_ = {
                score_compatibility::current<K, BucketCount, Layout>(), static_cast<uint32_t>(count)
            };
            if (!blacklist.keys().empty() && blacklist.kmer_length() != K) {
                return Err(Error::invalid_argument("store blacklist k mismatch"));
            }
            result.blacklist_ = blacklist;
            result.metadata_.compatibility.blacklist_identity = blacklist.identity();
            result.metadata_.compatibility.blacklist_version = blacklist.version();
            result.rows_.resize(count * BucketCount);
            if (count == 0) return result;
            CUDDL_TRY(detail::download_store<BucketCount>(store, result.rows_, stream));
            CUDDL_CUDA_TRY(stream.sync());
            result.names_.assign(names.begin(), names.end());
            return result;
        }();
    } catch (cuda::cuda_error const& error) {
        return Err(Error::cuda(static_cast<cudaError_t>(error.status())));
    } catch (std::bad_alloc const& error) {
        return Err(Error::resource(error.what()));
    }
 
    template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
    [[nodiscard]] Result<reference_database<K, BucketCount, Layout>> upload(
        cuda::stream_ref stream
    ) const {
        CUDDL_TRY((detail::validate_indexed_score_compatibility<K, BucketCount, Layout>(
            metadata_.compatibility
        )));
        auto rows =
            CUDDL_CUDA_TRY(cuda::make_device_buffer<uint16_t>(stream, stream.device(), rows_));
        auto database = CUDDL_TRY((reference_database<K, BucketCount, Layout>::build_async(
            {rows.data(), rows.size()}, metadata_.compatibility, stream
        )));
        database.names_ = names_;
        CUDDL_CUDA_TRY(stream.sync());
        return std::move(database);
    }
 
    [[nodiscard]] Result<void> save(std::filesystem::path const& path) const try {
        return [&]() -> Result<void> {
            detail::database_file_writer writer;
            writer.temporary = path.string() + ".tmp.XXXXXX";
            auto fd = ::mkstemp(writer.temporary.data());
            if (fd < 0) {
                writer.temporary.clear();
                return Err(Error::resource("cannot create database file beside: " + path.string()));
            }
            writer.output = ::fdopen(fd, "wb");
            if (!writer.output) {
                ::close(fd);
                return Err(Error::resource("cannot open temporary database file"));
            }
            CUDDL_TRY(writer.bytes("CUDDLDB\0", 8));
            CUDDL_TRY(writer.value(database_file_version));
            auto metadata = metadata_;
            if (blacklist_.identity() != metadata.compatibility.blacklist_identity ||
                blacklist_.version() != metadata.compatibility.blacklist_version) {
                return Err(
                    Error::invalid_argument("database is missing its construction blacklist")
                );
            }
            auto blacklist_count = static_cast<uint32_t>(blacklist_.keys().size());
            CUDDL_TRY(detail::database_file_metadata(writer, metadata, blacklist_count));
            CUDDL_TRY(writer.words(blacklist_.keys()));
            auto const prefix = ::ftello(writer.output);
            auto const limit = static_cast<uint64_t>(std::numeric_limits<off_t>::max());
            if (prefix < 0 || static_cast<uint64_t>(prefix) > limit - sizeof(uint32_t)) {
                return Err(Error::resource("cannot size temporary database file"));
            }
            uint64_t bytes = static_cast<uint64_t>(prefix) + sizeof(uint32_t);
            if (rows_.size() > (limit - bytes) / sizeof(uint16_t)) {
                return Err(Error::resource("database exceeds file size capacity"));
            }
            bytes += rows_.size() * sizeof(uint16_t);
            for (auto const& name : names_) {
                if (name.size() > std::numeric_limits<uint32_t>::max()) {
                    return Err(Error::resource("reference label exceeds binary format capacity"));
                }
                if (bytes > limit - sizeof(uint32_t) ||
                    name.size() > limit - bytes - sizeof(uint32_t)) {
                    return Err(Error::resource("database exceeds file size capacity"));
                }
                bytes += sizeof(uint32_t) + name.size();
            }
            // Allocate before names spill from stdio, avoiding delayed extent allocation when
            // the completed temporary file atomically replaces an existing database.
            auto const allocation = ::posix_fallocate(fd, 0, static_cast<off_t>(bytes));
            if (allocation != 0 && allocation != EOPNOTSUPP && allocation != ENOSYS) {
                return Err(
                    Error::resource(
                        "cannot allocate database file: " +
                        std::generic_category().message(allocation)
                    )
                );
            }
            for (auto const& name : names_) {
                CUDDL_TRY(writer.value(static_cast<uint32_t>(name.size())));
                CUDDL_TRY(writer.bytes(name.data(), name.size()));
            }
            CUDDL_TRY(writer.words(rows_));
            CUDDL_TRY(writer.value(static_cast<uint32_t>(writer.checksum)));
            auto closed = std::fclose(std::exchange(writer.output, nullptr));
            if (closed != 0) return Err(Error::resource("cannot close database file"));
            std::filesystem::rename(writer.temporary, path);
            writer.temporary.clear();
            return Ok();
        }();
    } catch (std::filesystem::filesystem_error const& error) {
        return Err(Error::resource(error.what()));
    } catch (std::bad_alloc const& error) {
        return Err(Error::resource(error.what()));
    }
 
    [[nodiscard]] static Result<reference_database_file> load(
        std::filesystem::path const& path
    ) try {
        return [&]() -> Result<reference_database_file> {
            detail::database_file_reader reader{
                std::ifstream(path, std::ios::binary | std::ios::ate)
            };
            auto size = reader.input.tellg();
            if (size < 0) {
                return Err(Error::resource("cannot open database file: " + path.string()));
            }
            reader.remaining = static_cast<uint64_t>(size);
            reader.input.seekg(0);
            char magic[8];
            CUDDL_TRY(reader.bytes(magic, sizeof(magic)));
            if (std::string_view(magic, sizeof(magic)) != std::string_view("CUDDLDB\0", 8)) {
                return Err(Error::invalid_argument("not a cuDDL reference database file"));
            }
            uint16_t version{};
            CUDDL_TRY(reader.value(version));
            if (version != database_file_version) {
                return Err(Error::invalid_argument("unsupported database file version"));
            }
            reference_database_file result;
            uint32_t blacklist_count = 0;
            CUDDL_TRY(detail::database_file_metadata(reader, result.metadata_, blacklist_count));
            auto& c = result.metadata_.compatibility;
            if (c.kmer_length < 1 || c.kmer_length > 31 || c.bucket_count < 2048 ||
                c.bucket_count > 8192 || !std::has_single_bit(c.bucket_count) ||
                (c.key_mask != 0xffffU && c.key_mask != 0x7fffU) || c.exponent_bits == 0 ||
                c.exponent_bits >= 16) {
                return Err(Error::invalid_argument("unsupported database construction metadata"));
            }
            c.indexed_bucket_count = c.bucket_count;
            c.mantissa_bits = 16U - c.exponent_bits;
            c.score_encoder_identity = 1U;
            c.hash_identity = 1U;
            c.hash_seed = detail::seed;
            c.canonicalisation_policy = 1U;
            c.blacklist_version = blacklist_count == 0 ? 0U : 1U;
            if (blacklist_count > reader.remaining / sizeof(uint64_t)) {
                return Err(Error::invalid_argument("invalid blacklist extent"));
            }
            std::vector<uint64_t> keys(blacklist_count);
            CUDDL_TRY(reader.words(keys));
            if (!std::is_sorted(keys.begin(), keys.end()) ||
                std::adjacent_find(keys.begin(), keys.end()) != keys.end()) {
                return Err(Error::invalid_argument("blacklist payload is not sorted and unique"));
            }
            result.blacklist_ = kmer_blacklist(c.kmer_length, std::move(keys));
            if (result.blacklist_.identity() != c.blacklist_identity ||
                result.blacklist_.version() != c.blacklist_version) {
                return Err(Error::invalid_argument("blacklist identity mismatch"));
            }
            uint64_t const count = result.metadata_.reference_count;
            uint64_t const words = count * c.bucket_count;
            uint64_t const payload_bytes = words * sizeof(uint16_t);
            if (words > std::numeric_limits<size_t>::max() / sizeof(uint16_t) ||
                payload_bytes + count * sizeof(uint32_t) + sizeof(uint32_t) > reader.remaining) {
                return Err(Error::invalid_argument("invalid reference database extents"));
            }
            result.names_.reserve(count);
            for (uint64_t id = 0; id < count; ++id) {
                uint32_t length{};
                CUDDL_TRY(reader.value(length));
                auto const required =
                    payload_bytes + (count - id - 1) * sizeof(uint32_t) + sizeof(uint32_t);
                if (required > reader.remaining || length > reader.remaining - required) {
                    return Err(Error::invalid_argument("invalid reference label extent"));
                }
                std::string name(length, '\0');
                CUDDL_TRY(reader.bytes(name.data(), name.size()));
                result.names_.push_back(std::move(name));
            }
            if (reader.remaining != payload_bytes + sizeof(uint32_t)) {
                return Err(Error::invalid_argument("unexpected reference database payload size"));
            }
            result.rows_.resize(words);
            CUDDL_TRY(reader.words(result.rows_));
            auto const expected_checksum = static_cast<uint32_t>(reader.checksum);
            uint32_t checksum{};
            CUDDL_TRY(reader.value(checksum));
            if (checksum != expected_checksum) {
                return Err(Error::invalid_argument("database checksum mismatch"));
            }
            return result;
        }();
    } catch (std::invalid_argument const& error) {
        return Err(Error::invalid_argument(error.what()));
    } catch (std::bad_alloc const& error) {
        return Err(Error::resource(error.what()));
    }
 
   private:
    friend class reference_index_file;
    friend struct detail::reference_file_assembly;
    reference_database_file() = default;
 
    template <size_t BucketCount>
    [[nodiscard]] auto download_rows(cuda::stream_ref stream) {
        return [this, stream](device_span<uint32_t const> group, size_t base, size_t count) {
            return detail::download_store<BucketCount>(
                group, std::span{rows_}.subspan(base * BucketCount, count * BucketCount), stream
            );
        };
    }
 
    template <uint32_t K, size_t BucketCount>
    Result<void> set_blacklist(
        std::optional<std::reference_wrapper<device_blacklist const>> filter,
        cuda::stream_ref stream
    ) {
        if (!filter) return Ok();
        CUDDL_TRY(filter->get().validate(K, BucketCount, stream));
        blacklist_ = filter->get().source();
        metadata_.compatibility.blacklist_identity = blacklist_.identity();
        metadata_.compatibility.blacklist_version = blacklist_.version();
        return Ok();
    }
 
    kmer_blacklist blacklist_;
    reference_database_metadata metadata_{};
    std::vector<std::string> names_;
    std::vector<uint16_t> rows_;
};
 
template <uint32_t K, size_t BucketCount, typename Layout = default_register_layout>
[[nodiscard]] Result<cuda::device_buffer<uint32_t>> build_sketch_store(
    std::span<std::filesystem::path const> paths,
    cuda::stream_ref stream,
    path_build_options options = {}
) try {
    return [&]() -> Result<cuda::device_buffer<uint32_t>> {
        if (options.blacklist) CUDDL_TRY(options.blacklist->get().validate(K, BucketCount, stream));
 
        auto store = cuda::make_device_buffer<uint32_t>(
            stream, stream.device(), paths.size() * BucketCount, cuda::no_init
        );
 
        CUDDL_TRY((detail::stage_paths<K, BucketCount, Layout>(
            paths,
            stream,
            options.staging_bytes,
            options.parser_workers,
            options.transfer,
            options.decompression,
            options.direct_reads,
            options.statistics,
            [&](device_span<uint32_t const> group, size_t base, size_t count) -> Result<void> {
                CUDDL_CUDA_TRY(
                    cuda::copy_bytes(
                        stream,
                        group,
                        device_span<uint32_t>{
                            store.data() + base * BucketCount, count * BucketCount
                        }
                    )
                );
                return Ok();
            },
            options.blacklist ? options.blacklist->get().view() : cuda::std::nullopt
        )));
        return store;
    }();
} catch (cuda::cuda_error const& error) {
    return Err(Error::cuda(static_cast<cudaError_t>(error.status())));
} catch (std::system_error const& error) {
    return Err(Error::resource(error.what()));
} catch (std::bad_alloc const& error) {
    return Err(Error::resource(error.what()));
}
 
}  // namespace cuddl