Jungfraujoch/reader/MiniCBF.cpp

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

#include "MiniCBF.h"

#include <zlib.h>

#include <algorithm>
#include <cctype>
#include <cstring>
#include <fstream>
#include <limits>
#include <map>
#include <mutex>
#include <regex>
#include <stdexcept>

#if defined(__SSE2__) || defined(_M_X64) || defined(_M_AMD64)
#  include <emmintrin.h>
#  define MINICBF_STREAM_STORE 1
#endif

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

namespace minicbf {

namespace {

// One capture group, first match, or nothing. The headers are a few kB, so a regex per field is
// cheap and keeps each rule next to the thing it reads - but COMPILING one is not: measured at ~60 us
// here, and a header parse runs 26 of them. Parsing happens once per image read across a whole sweep,
// so the same two dozen patterns were being recompiled thousands of times. They are compiled once and
// kept; std::map never invalidates a reference, so the pointer outlives the lock.
std::optional<std::string> Match(const std::string &text, const char *pattern) {
    static std::mutex cache_mutex;
    static std::map<std::string, std::regex> cache;
    const std::regex *re;
    {
        const std::lock_guard<std::mutex> lock(cache_mutex);
        auto it = cache.find(pattern);
        if (it == cache.end())
            it = cache.emplace(pattern, std::regex(pattern)).first;
        re = &it->second;
    }
    std::smatch m;
    if (!std::regex_search(text, m, *re) || m.size() < 2)
        return {};
    return m[1].str();
}

// The captures are character classes, not number grammars: "[\d.eE+-]+" matches "." and "+-", and
// "(\d+)" matches a digit string too long for int64. std::stod and std::stoll answer both with a raw
// std:: exception, which would leave the format probe below - CanRead catches JFJochException only -
// and reach the caller as an unhandled throw from merely LOOKING at a file. A header field that does
// not parse is a malformed header, so say that.
double Num(const std::string &text, const char *pattern, double fallback = 0.0) {
    const auto s = Match(text, pattern);
    if (!s.has_value())
        return fallback;
    try {
        return std::stod(*s);
    } catch (const std::exception &) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Malformed number in CBF header: '" + *s + "'");
    }
}

int64_t Int(const std::string &text, const char *pattern, int64_t fallback = 0) {
    const auto s = Match(text, pattern);
    if (!s.has_value())
        return fallback;
    try {
        return std::stoll(*s);
    } catch (const std::exception &) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Malformed integer in CBF header: '" + *s + "'");
    }
}

// A goniometer angle, or nothing where the head has no such axis. Writers spell that -9999, and
// taking a sentinel for an angle would put the head somewhere it never was.
std::optional<double> Angle(const std::string &text, const char *pattern) {
    const auto s = Match(text, pattern);
    if (!s.has_value())
        return {};
    double v;
    try {
        v = std::stod(*s);
    } catch (const std::exception &) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Malformed angle in CBF header: '" + *s + "'");
    }
    if (v < -9998.0)
        return {};
    return v;
}

// One "loop_" of the imgCIF template block these headers carry, as rows keyed by tag. Tags and values
// are both read as whitespace-separated tokens rather than by line, because a template packs several
// tags onto one line ("_axis.vector[1] _axis.vector[2] _axis.vector[3]") and the rows that follow are
// laid out to match the tags, not the lines.
std::vector<std::map<std::string, std::string>> ParseLoop(const std::string &text, const std::string &tag) {
    // The loop_ that introduces the tag, not the tag's own position: everything before it is another
    // loop's data.
    const size_t tag_at = text.find("\n" + tag);
    if (tag_at == std::string::npos)
        return {};
    const size_t loop_at = text.rfind("loop_", tag_at);
    if (loop_at == std::string::npos)
        return {};

    // Bare tokens, and the quoted ones a CIF value may be - a quoted value holding spaces would
    // otherwise be counted as several columns and shift every row after it.
    std::vector<std::string> tokens;
    for (size_t i = loop_at + 5; i < text.size();) {
        while ((i < text.size()) && std::isspace(static_cast<unsigned char>(text[i])))
            i++;
        if (i >= text.size())
            break;
        size_t end;
        if ((text[i] == '\'') || (text[i] == '"')) {
            end = text.find(text[i], i + 1);
            if (end == std::string::npos)
                break;
            tokens.push_back(text.substr(i + 1, end - i - 1));
            end++;
        } else {
            end = i;
            while ((end < text.size()) && !std::isspace(static_cast<unsigned char>(text[end])))
                end++;
            tokens.push_back(text.substr(i, end - i));
        }
        // A second loop_, or a tag belonging to another category, ends this one.
        if ((tokens.back() == "loop_")
                || (tokens.back().starts_with("_") && !tokens.back().starts_with(tag.substr(0, tag.find('.') + 1)))) {
            tokens.pop_back();
            break;
        }
        i = end;
    }

    std::vector<std::string> names;
    size_t first_value = 0;
    while ((first_value < tokens.size()) && tokens[first_value].starts_with("_"))
        names.push_back(tokens[first_value++]);
    if (names.empty())
        return {};

    std::vector<std::map<std::string, std::string>> rows;
    for (size_t i = first_value; i + names.size() <= tokens.size(); i += names.size()) {
        std::map<std::string, std::string> row;
        for (size_t j = 0; j < names.size(); j++)
            row[names[j]] = tokens[i + j];
        rows.push_back(std::move(row));
    }
    return rows;
}

std::string Field(const std::map<std::string, std::string> &row, const std::string &name) {
    const auto it = row.find(name);
    return (it == row.end()) ? std::string() : it->second;
}

// The imgCIF axis table: which axis turns or translates in which laboratory direction, and which two
// axes the image's columns and rows run along. Absent from most headers, which say nothing about any
// of this and are left exactly as they were read before.
//
// The element vectors are stated in the frame of the axis they depend on. Between them and the
// detector's own rotation every header seen has translations only, so they describe the image in the
// unswung detector frame - which is where the image orientation belongs, with the arm applied on top.
// Following Hammersley, Bernstein & Westbrook (2006) Int. Tables Cryst. G, 444-458
void ParseAxisTable(const std::string &text, Header &h) {
    const auto axes = ParseLoop(text, "_axis.id");
    if (axes.empty())
        return;

    std::map<std::string, std::array<double, 3>> vector_of;
    for (const auto &row: axes) {
        const std::string id = Field(row, "_axis.id");
        const std::string type = Field(row, "_axis.type");
        const std::string equipment = Field(row, "_axis.equipment");
        const std::string depends_on = Field(row, "_axis.depends_on");
        std::array<double, 3> v{};
        try {
            for (int i = 0; i < 3; i++)
                v[i] = std::stod(Field(row, "_axis.vector[" + std::to_string(i + 1) + "]"));
        } catch (const std::exception &) {
            continue;   // "." for a vector: the table states no direction for this axis
        }
        vector_of[id] = v;

        // The base spindle and the detector arm are the rotations that hang off nothing: everything
        // further in is carried by them. Naming neither, so a beamline is free to call them anything.
        if ((type == "rotation") && (depends_on == ".")) {
            if (equipment == "goniometer")
                h.spindle_axis = v;
            else if (equipment == "detector")
                h.detector_axis = v;
        } else if ((type == "rotation") && (equipment == "goniometer")) {
            h.inner_spindle_axis = v;   // carried by the base axis: a fixed-chi or kappa inclination
        }
    }

    // Which axis the fast index runs along, and which the slow, through the two tables that say so.
    std::map<std::string, std::string> axis_of_set;
    for (const auto &row: ParseLoop(text, "_array_structure_list_axis.axis_set_id"))
        axis_of_set[Field(row, "_array_structure_list_axis.axis_set_id")]
                = Field(row, "_array_structure_list_axis.axis_id");

    for (const auto &row: ParseLoop(text, "_array_structure_list.array_id")) {
        const std::string set = Field(row, "_array_structure_list.axis_set_id");
        const auto id = axis_of_set.contains(set) ? axis_of_set[set] : set;
        const auto it = vector_of.find(id);
        if (it == vector_of.end())
            continue;
        std::array<double, 3> v = it->second;
        if (Field(row, "_array_structure_list.direction") == "decreasing")
            for (double &c: v)
                c = -c;
        const std::string index = Field(row, "_array_structure_list.index");
        if (index == "1")
            h.fast_direction = v;
        else if (index == "2")
            h.slow_direction = v;
    }
}

int16_t ReadI16(const uint8_t *p) { int16_t v; std::memcpy(&v, p, 2); return v; }
int32_t ReadI32(const uint8_t *p) { int32_t v; std::memcpy(&v, p, 4); return v; }
int64_t ReadI64(const uint8_t *p) { int64_t v; std::memcpy(&v, p, 8); return v; }

}  // namespace

bool ScansPhi(const Header &h) {
    if (h.phi_increment_deg != 0.0 || h.omega_increment_deg != 0.0 || h.chi_increment_deg != 0.0)
        return h.phi_increment_deg != 0.0;
    std::string name = h.axis_name;
    std::transform(name.begin(), name.end(), name.begin(),
                   [](unsigned char c) { return std::tolower(c); });
    return name.starts_with("phi");
}

std::optional<size_t> FindBinarySection(const uint8_t *data, size_t size) {
    if (size < sizeof(BINARY_SEPARATOR))
        return {};
    const auto *end = data + size;
    const auto *hit = std::search(data, end, std::begin(BINARY_SEPARATOR), std::end(BINARY_SEPARATOR));
    if (hit == end)
        return {};
    return static_cast<size_t>(hit - data) + sizeof(BINARY_SEPARATOR);
}

Header ParseHeader(const char *data, size_t size) {
    const std::string t(data, size);
    Header h;

    h.detector       = Match(t, R"(#\s*Detector:\s*([^\r\n]+))").value_or("PILATUS");
    h.pixel_x_m      = Num(t, R"(#\s*Pixel_size\s+([\d.eE+-]+)\s*m)");
    h.pixel_y_m      = Num(t, R"(#\s*Pixel_size\s+[\d.eE+-]+\s*m\s*x\s*([\d.eE+-]+)\s*m)");
    // The unit is optional: one ALBA set writes "thickness 0.001000" with no " m" after it.
    h.thickness_m    = Num(t, R"(sensor,\s*thickness\s+([\d.eE+-]+))");
    h.distance_m     = Num(t, R"(#\s*Detector_distance\s+([\d.eE+-]+))");
    h.beam_x_px      = Num(t, R"(#\s*Beam_xy\s*\(\s*([\d.eE+-]+))");
    h.beam_y_px      = Num(t, R"(#\s*Beam_xy\s*\([^,]+,\s*([\d.eE+-]+))");
    h.wavelength_A   = Num(t, R"(#\s*Wavelength\s+([\d.eE+-]+))");
    h.start_angle_deg     = Num(t, R"(#\s*Start_angle\s+([\d.eE+-]+))");
    h.angle_increment_deg = Num(t, R"(#\s*Angle_increment\s+([\d.eE+-]+))");
    h.two_theta_deg  = Num(t, R"(#\s*Detector_2theta\s+([\d.eE+-]+))");
    // The whitespace after each name is what keeps "Chi_increment" out of "Chi".
    h.chi_deg        = Angle(t, R"(#\s*Chi\s+([\d.eE+-]+))");
    h.omega_deg      = Angle(t, R"(#\s*Omega\s+([\d.eE+-]+))");
    h.chi_increment_deg   = Angle(t, R"(#\s*Chi_increment\s+([\d.eE+-]+))").value_or(0.0);
    h.phi_increment_deg   = Angle(t, R"(#\s*Phi_increment\s+([\d.eE+-]+))").value_or(0.0);
    h.omega_increment_deg = Angle(t, R"(#\s*Omega_increment\s+([\d.eE+-]+))").value_or(0.0);
    h.exposure_s     = Num(t, R"(#\s*Exposure_time\s+([\d.eE+-]+))");
    h.period_s       = Num(t, R"(#\s*Exposure_period\s+([\d.eE+-]+))");
    h.count_cutoff   = Int(t, R"(#\s*Count_cutoff\s+(\d+))");
    h.axis_name      = Match(t, R"(#\s*Oscillation_axis\s+(\S+))").value_or("omega");
    // "+SLOW" / "+FAST" on that same line: which of the image's two directions the spindle runs
    // along. Some writers state that instead of an axis name, and it is the only thing a header with
    // no axis table says about the spindle's direction at all.
    h.spindle_along_slow = Match(t, R"(#\s*Oscillation_axis[^\r\n]*\+(SLOW|slow))").has_value();
    ParseAxisTable(t, h);

    // "# Silicon sensor, ..." / "# CdTe sensor, ...". The rest of the code compares the material
    // against "CdTe" (BraggIntegrationEngine), so an unnormalised "Silicon" would silently give a
    // CdTe sensor silicon's attenuation length.
    if (const auto m = Match(t, R"(#\s*(\w+)\s+sensor,)")) {
        std::string s = *m;
        std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { return std::tolower(c); });
        h.material = (s == "cdte") ? "CdTe" : "Si";
    }

    h.nx    = Int(t, R"(X-Binary-Size-Fastest-Dimension:\s*(\d+))");
    h.ny    = Int(t, R"(X-Binary-Size-Second-Dimension:\s*(\d+))");
    h.nelem = Int(t, R"(X-Binary-Number-of-Elements:\s*(\d+))");

    const auto conv = Match(t, R"RE(conversions\s*=\s*"([^"]+)")RE").value_or("");
    h.byte_offset = conv.find("x-CBF_BYTE_OFFSET") != std::string::npos;

    if (h.period_s <= 0.0)
        h.period_s = h.exposure_s;

    return h;
}

namespace {

// Nothing on this side reads the decoded image back: it goes to the GPU by DMA, or a CPU engine
// walks it once and drops it. An ordinary store therefore pays twice over - it reads every cache
// line before overwriting it, and it evicts 72 MB of live cache to make room - so the pixels go
// straight past the cache. Measured on an 18 Mpx frame: 27 ms of ordinary stores against 11 ms of
// these. The bytes written are the same either way; the fence is what makes them visible to the
// copy that follows.
void StorePixel(int32_t *out, int32_t value) {
#ifdef MINICBF_STREAM_STORE
    _mm_stream_si32(out, value);
#else
    *out = value;
#endif
}

void StoreFence() {
#ifdef MINICBF_STREAM_STORE
    _mm_sfence();
#endif
}

}  // namespace

// The x-CBF_BYTE_OFFSET scheme: a running value, each pixel stored as a delta in the smallest
// container that holds it, escaping to the next size with that container's most negative value.
// Following Bernstein & Hammersley (2006) Int. Tables Cryst. G, 37-43
void DecodeByteOffset(const uint8_t *data, size_t size, int32_t *out, size_t n_pixels) {
    int64_t value = 0;
    size_t pos = 0;
    size_t written = 0;

    while (written < n_pixels) {
        if (pos >= size)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "miniCBF byte-offset stream ended after " + std::to_string(written)
                                  + " of " + std::to_string(n_pixels) + " pixels");

        int64_t delta = static_cast<int8_t>(data[pos]);
        pos += 1;

        if (delta == -128) {
            if (pos + 2 > size) break;
            delta = ReadI16(data + pos);
            pos += 2;
            if (delta == -32768) {
                if (pos + 4 > size) break;
                delta = ReadI32(data + pos);
                pos += 4;
                if (delta == INT32_MIN) {
                    if (pos + 8 > size) break;
                    delta = ReadI64(data + pos);
                    pos += 8;
                }
            }
        }

        value += delta;
        StorePixel(out + written, static_cast<int32_t>(value));
        written++;
    }

    StoreFence();

    if (written != n_pixels)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "miniCBF byte-offset stream ended after " + std::to_string(written)
                              + " of " + std::to_string(n_pixels) + " pixels");
}

namespace {

// Whether the file is gzip-wrapped, from its first two bytes rather than its name - a sweep is
// named consistently but the magic is what decides how to read it.
bool IsGzip(const std::string &path) {
    std::ifstream f(path, std::ios::binary);
    unsigned char m[2] = {0, 0};
    if (f)
        f.read(reinterpret_cast<char *>(m), sizeof(m));
    return m[0] == 0x1f && m[1] == 0x8b;
}

// The gzip path. EMBL Hamburg's beamlines write .cbf.gz by default, so a sweep from there is
// otherwise twice the disk and a decompression pass before anything can read a frame. The
// decompressed size is not in the stream, so the buffer grows in chunks instead of being sized
// up front the way the plain path can.
void SlurpGz(const std::string &path, size_t max_bytes, std::vector<uint8_t> &buf) {
    gzFile g = gzopen(path.c_str(), "rb");
    if (g == nullptr)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Cannot open gzipped CBF file " + path);
    buf.clear();
    constexpr size_t CHUNK = 1u << 20;
    while (buf.size() < max_bytes) {
        const size_t want = std::min(CHUNK, max_bytes - buf.size());
        const size_t at = buf.size();
        buf.resize(at + want);
        const int got = gzread(g, buf.data() + at, static_cast<unsigned>(want));
        if (got < 0) {
            int err = 0;
            const std::string msg = gzerror(g, &err);
            gzclose(g);
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Error decompressing " + path + ": " + msg);
        }
        buf.resize(at + static_cast<size_t>(got));
        if (static_cast<size_t>(got) < want)
            break;                      // end of the stream
    }
    gzclose(g);
}

void Slurp(const std::string &path, size_t max_bytes, std::vector<uint8_t> &buf) {
    // The plain path is kept exactly as it was rather than reading everything through zlib: a
    // gzFile over an uncompressed file works, but it copies every byte through zlib's own buffer,
    // and the overwhelming majority of sweeps are not compressed.
    if (IsGzip(path)) {
        SlurpGz(path, max_bytes, buf);
        return;
    }
    std::ifstream f(path, std::ios::binary);
    if (!f)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Cannot open CBF file " + path);
    f.seekg(0, std::ios::end);
    const auto file_size = static_cast<size_t>(f.tellg());
    f.seekg(0, std::ios::beg);
    buf.resize(std::min(file_size, max_bytes));
    f.read(reinterpret_cast<char *>(buf.data()), static_cast<std::streamsize>(buf.size()));
    buf.resize(static_cast<size_t>(f.gcount()));
}

// Enough to reach the separator on any header seen in the wild (the longest measured is ~6.3 kB).
// The first read is kept small because a sweep's startup reads every frame's header: on a spinning
// disk 256 kB from each of 2400 files took twice as long as 16 kB (12 s against 6 s). A header
// that does not end within it is read again with the large probe, so the result never depends on it.
constexpr size_t HEADER_FIRST_PROBE_BYTES = 16 * 1024;
constexpr size_t HEADER_PROBE_BYTES = 256 * 1024;

}  // namespace

Header ReadHeader(const std::string &path) {
    std::vector<uint8_t> buf;
    Slurp(path, HEADER_FIRST_PROBE_BYTES, buf);
    auto start = FindBinarySection(buf.data(), buf.size());
    if (!start.has_value()) {
        Slurp(path, HEADER_PROBE_BYTES, buf);
        start = FindBinarySection(buf.data(), buf.size());
    }
    if (!start.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No CBF binary section in " + path);
    return ParseHeader(reinterpret_cast<const char *>(buf.data()),
                       *start - sizeof(BINARY_SEPARATOR));
}

namespace {

// One read of the file, header parsed, dimensions checked. The caller supplies where the pixels go.
Header ReadCommon(const std::string &path, std::vector<uint8_t> &buf, size_t &binary_start) {
    Slurp(path, std::numeric_limits<size_t>::max(), buf);
    const auto start = FindBinarySection(buf.data(), buf.size());
    if (!start.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No CBF binary section in " + path);

    const Header h = ParseHeader(reinterpret_cast<const char *>(buf.data()),
                                 *start - sizeof(BINARY_SEPARATOR));
    if (!h.byte_offset)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Unsupported CBF compression in " + path + " (only x-CBF_BYTE_OFFSET)");
    if (h.nelem <= 0 || h.nx <= 0 || h.ny <= 0 || h.nelem != h.nx * h.ny)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Inconsistent image dimensions in " + path);
    // Without a pixel size there is no geometry at all - every resolution, every scattering vector
    // and the beam centre in millimetres all scale by it - and the default of 0 collapses all of
    // them silently. A byte-offset CBF carrying no "# Pixel_size" line is not a detector image from
    // this family at all; XDS writes correction files in exactly that shape. Refuse it here rather
    // than let it through with a geometry of zero.
    if (!(h.pixel_x_m > 0.0) || !(h.pixel_y_m > 0.0))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + " has no pixel size in its header (no '# Pixel_size' line); "
                              "it carries a CBF binary section but is not a detector image");

    binary_start = *start;
    return h;
}

}  // namespace

Header Read(const std::string &path, std::vector<int32_t> &out) {
    std::vector<uint8_t> buf;
    size_t start = 0;
    const Header h = ReadCommon(path, buf, start);
    out.resize(static_cast<size_t>(h.nelem));
    DecodeByteOffset(buf.data() + start, buf.size() - start, out.data(),
                     static_cast<size_t>(h.nelem));
    return h;
}

Header ReadInto(const std::string &path, int32_t *out, size_t capacity, std::vector<uint8_t> &scratch) {
    size_t start = 0;
    const Header h = ReadCommon(path, scratch, start);
    if (static_cast<size_t>(h.nelem) > capacity)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "CBF image does not fit the supplied buffer: " + path);
    DecodeByteOffset(scratch.data() + start, scratch.size() - start, out, static_cast<size_t>(h.nelem));
    return h;
}

}  // namespace minicbf