Jungfraujoch/reader/HDF5ImageSource.cpp

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "HDF5ImageSource.h"
#include "../common/JFJochException.h"

#include <algorithm>

#ifdef _WIN32
#include <windows.h>
#else
#include <fcntl.h>
#include <unistd.h>
#endif

// Positional reads: pread() on POSIX, ReadFile() with an OVERLAPPED offset on Windows. Both take the
// offset as an argument instead of moving a shared file position, which is what lets every worker
// thread read through one handle at the same time.
HDF5ImageSource::RawFile::RawFile(const std::string &path) {
#ifdef _WIN32
    HANDLE h = CreateFileA(path.c_str(), GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
                           OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
    handle_ = (h == INVALID_HANDLE_VALUE) ? -1 : reinterpret_cast<intptr_t>(h);
#else
    handle_ = ::open(path.c_str(), O_RDONLY);
#endif
}

HDF5ImageSource::RawFile::~RawFile() {
    if (handle_ == -1)
        return;
#ifdef _WIN32
    CloseHandle(reinterpret_cast<HANDLE>(handle_));
#else
    ::close(static_cast<int>(handle_));
#endif
}

void HDF5ImageSource::RawFile::ReadAt(void *dst, size_t size, uint64_t address) const {
    auto *out = static_cast<uint8_t *>(dst);
    size_t done = 0;
    while (done < size) {
#ifdef _WIN32
        OVERLAPPED ov{};
        ov.Offset = static_cast<DWORD>((address + done) & 0xFFFFFFFFULL);
        ov.OffsetHigh = static_cast<DWORD>((address + done) >> 32);
        DWORD got = 0;
        const bool ok = ReadFile(reinterpret_cast<HANDLE>(handle_), out + done,
                                 static_cast<DWORD>(size - done), &got, &ov);
        const long long n = ok ? static_cast<long long>(got) : -1;
#else
        const long long n = ::pread(static_cast<int>(handle_), out + done, size - done, address + done);
#endif
        if (n <= 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, "Error reading image chunk from file");
        done += static_cast<size_t>(n);
    }
}

void HDF5ImageSource::Configure(HDF5ImageLocator::Layout layout) {
    dataset_cache_.clear();
    locator_.Configure(std::move(layout));
}

void HDF5ImageSource::Clear() {
    dataset_cache_.clear();
    locator_.Clear();
}

HDF5ImageLocator::Location HDF5ImageSource::Resolve(int64_t global) const {
    return locator_.Resolve(global);
}

StoredPixelFormat HDF5ImageSource::GetStoredPixelFormat() const {
    auto loc = locator_.Resolve(0);
    HDF5DataSet dataset(*loc.file, loc.dataset);
    HDF5DataType datatype(dataset);
    return {static_cast<int64_t>(datatype.GetElemSize()) * 8, datatype.IsSigned()};
}

std::vector<HDF5DataSourceMessage> HDF5ImageSource::GetSourceMapping(uint64_t first_image,
                                                                     std::optional<uint64_t> image_count,
                                                                     uint64_t total_images,
                                                                     uint64_t stride) const {
    return locator_.GetSourceMapping(first_image, image_count, total_images, stride);
}

const HDF5ImageSource::OpenDataset &
HDF5ImageSource::GetDataset(const HDF5ImageLocator::Location &loc) const {
    auto key = std::make_pair(loc.file.get(), loc.dataset);
    if (auto it = dataset_cache_.find(key); it != dataset_cache_.end())
        return it->second;

    OpenDataset entry;
    entry.file = loc.file;
    entry.dataset = std::make_unique<HDF5DataSet>(*loc.file, loc.dataset);

    HDF5DataSpace dataspace(*entry.dataset);
    HDF5DataType datatype(*entry.dataset);
    HDF5Dcpl dcpl(*entry.dataset);

    const auto rank = dataspace.GetNumOfDimensions();
    if (rank != 3 && rank != 4)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              loc.dataset + " dataset must be 3D or 4D");
    entry.multichannel = (rank == 4);

    if (datatype.IsFloat())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Float datasets not supported at this time");

    auto dim = dataspace.GetDimensions();
    entry.height = dim[rank - 2];
    entry.width = dim[rank - 1];
    entry.mode = CalcImageMode(datatype.GetElemSize(), datatype.IsFloat(), datatype.IsSigned());

    // One chunk per image: [1, h, w], or [1, 1, h, w] for a multichannel dataset
    auto chunk_size = dcpl.GetChunking();
    entry.direct_chunk = (chunk_size.size() == rank)
                         && std::all_of(chunk_size.begin(), chunk_size.end() - 2, [](hsize_t c) { return c == 1; })
                         && (chunk_size[rank - 2] == entry.height) && (chunk_size[rank - 1] == entry.width);
    if (entry.direct_chunk)
        entry.algorithm = dcpl.GetCompression();

    if (entry.direct_chunk) {
        // A positional read needs the file the PIXELS are in, and that is not always the file that
        // was opened: a dataset reached through an external link lives elsewhere, and a chunk
        // address is an offset into ITS file - as is the user block the address counts from. Ask
        // HDF5 which file the dataset ended up in rather than assuming it is the one we asked.
        std::string path = loc.path;
        const hid_t owner = H5Iget_file_id(entry.dataset->GetID());
        if (owner >= 0) {
            const ssize_t name_size = H5Fget_name(owner, nullptr, 0);
            if (name_size > 0) {
                std::string name(static_cast<size_t>(name_size), '\0');
                if (H5Fget_name(owner, name.data(), name.size() + 1) >= 0)
                    path = name;
            }
            const hid_t fcpl = H5Fget_create_plist(owner);
            if (fcpl >= 0) {
                hsize_t user_block = 0;
                if (H5Pget_userblock(fcpl, &user_block) >= 0)
                    entry.user_block = user_block;
                H5Pclose(fcpl);
            }
            H5Fclose(owner);
        }
        if (!path.empty()) {
            entry.raw = std::make_shared<RawFile>(path);
            if (!entry.raw->IsOpen())
                entry.raw.reset();
        }
    }

    return dataset_cache_.emplace(std::move(key), std::move(entry)).first->second;
}

std::optional<HDF5ImageSource::DirectChunk>
HDF5ImageSource::PrepareDirectRead(const HDF5ImageLocator::Location &loc) const {
    const auto &ds = GetDataset(loc);
    if (!ds.raw)
        return {};

    const hsize_t coord_3d[3] = {loc.local_index, 0, 0};
    const hsize_t coord_4d[4] = {loc.local_index, loc.channel, 0, 0};
    const hsize_t *coord = ds.multichannel ? coord_4d : coord_3d;
    unsigned filter_mask = 0;
    haddr_t address = HADDR_UNDEF;
    hsize_t size = 0;
    if (H5Dget_chunk_info_by_coord(ds.dataset->GetID(), coord, &filter_mask, &address, &size) < 0)
        return {};
    // A chunk nobody ever wrote has no address and no bytes; only HDF5 knows it reads as the fill
    // value, so hand those back to it.
    if (address == HADDR_UNDEF || size == 0)
        return {};

    return DirectChunk{ds.raw, ds.user_block + address, static_cast<uint32_t>(size),
                       ds.width, ds.height, ds.mode, ds.algorithm};
}

CompressedImage HDF5ImageSource::ReadDirect(RawByteBuffer &buffer, const DirectChunk &chunk) {
    buffer.resize(chunk.size);
    chunk.file->ReadAt(buffer.data(), chunk.size, chunk.address);
    return {buffer.data(), buffer.size(), chunk.width, chunk.height, chunk.mode, chunk.algorithm};
}