match.cpp

// Dual-bigram filter: for ≤1 mismatch, at least one of two non-overlapping
// bigrams (2-char sequences) at fixed positions must match exactly.
// 26²=676 buckets → avg <1 pattern/bucket → nearly O(N) per position.
#include <iostream>
#include <vector>
#include <cstdint>
#include <fcntl.h>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <cstring>
#include <chrono>
#include <atomic>
#include <thread>
#include <immintrin.h>
#include <omp.h>
#include <sched.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <algorithm>

constexpr double POWER_LIMIT = 600.0;
constexpr double POWER_GUARD = 15.0;
constexpr int POWER_SAMPLE_MS = 5;
constexpr int POWER_THROTTLE_US_HIGH = 1000;
constexpr int POWER_THROTTLE_US_LOW = 200;
constexpr int64_t POWER_CHECK_MASK = 0xFFFF;
std::atomic<int> g_throttle_us(0);

bool query_power(double& cpu, double& other) {
    int sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock < 0) return false;
    struct sockaddr_in serv_addr{};
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_port = htons(19937);
    if (inet_pton(AF_INET, "127.0.0.1", &serv_addr.sin_addr) <= 0) { close(sock); return false; }
    if (connect(sock, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) { close(sock); return false; }
    const char* msg = "GET\n";
    if (send(sock, msg, strlen(msg), MSG_NOSIGNAL) <= 0) { close(sock); return false; }
    char buffer[128] = {0};
    int valread = read(sock, buffer, sizeof(buffer) - 1);
    close(sock);
    if (valread <= 0) return false;
    return sscanf(buffer, "%lf %lf", &cpu, &other) == 2;
}

void power_monitor(std::atomic<bool>& done) {
    while (!done.load(std::memory_order_relaxed)) {
        double cpu = 0.0, other = 0.0;
        int throttle = 0;
        if (query_power(cpu, other)) {
            double total = cpu + other;
            if (total > POWER_LIMIT) throttle = POWER_THROTTLE_US_HIGH;
            else if (total > POWER_LIMIT - POWER_GUARD) throttle = POWER_THROTTLE_US_LOW;
        }
        g_throttle_us.store(throttle, std::memory_order_relaxed);
        std::this_thread::sleep_for(std::chrono::milliseconds(POWER_SAMPLE_MS));
    }
}

struct MappedFile {
    int fd; void* data; size_t size;
    MappedFile(const char* fn) {
        fd = open(fn, O_RDONLY);
        if (fd < 0) { perror("open"); exit(1); }
        struct stat sb; fstat(fd, &sb); size = sb.st_size;
        data = mmap(NULL, size, PROT_READ, MAP_PRIVATE | MAP_POPULATE, fd, 0);
        if (data == MAP_FAILED) { perror("mmap"); exit(1); }
    }
    ~MappedFile() { munmap(data, size); close(fd); }
};

std::vector<int> parse_cpu_range(const std::string& s) {
    std::vector<int> cpus;
    size_t start = 0;
    while (start < s.size()) {
        size_t comma = s.find(',', start);
        if (comma == std::string::npos) comma = s.size();
        std::string part = s.substr(start, comma - start);
        size_t dash = part.find('-');
        if (dash != std::string::npos) {
            int a = std::stoi(part.substr(0, dash));
            int b = std::stoi(part.substr(dash + 1));
            for (int i = a; i <= b; ++i) cpus.push_back(i);
        } else cpus.push_back(std::stoi(part));
        start = comma + 1;
    }
    return cpus;
}

void fix_cpu_affinity() {
    FILE* f = fopen("/sys/fs/cgroup/cpuset.cpus", "r");
    if (!f) f = fopen("/sys/fs/cgroup/cpuset.cpus.effective", "r");
    if (!f) f = fopen("/sys/fs/cgroup/cpuset/cpuset.cpus", "r");
    if (!f) { fprintf(stderr, "[AFFINITY] Cannot read cgroup cpuset\n"); return; }
    char buf[256];
    if (!fgets(buf, sizeof(buf), f)) { fclose(f); return; }
    fclose(f);
    size_t len = strlen(buf);
    if (len > 0 && buf[len-1] == '\n') buf[len-1] = '\0';
    fprintf(stderr, "[AFFINITY] Cgroup cpuset: %s\n", buf);
    std::vector<int> cpus = parse_cpu_range(buf);
    fprintf(stderr, "[AFFINITY] Parsed %zu CPUs\n", cpus.size());
    if (cpus.empty()) return;
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    for (int cpu : cpus) CPU_SET(cpu, &cpuset);
    if (sched_setaffinity(0, sizeof(cpuset), &cpuset) == 0)
        fprintf(stderr, "[AFFINITY] Successfully set affinity to %zu CPUs\n", cpus.size());
    omp_set_num_threads(cpus.size());
}

constexpr int PAT_LEN = 64;
constexpr int BIGRAM_BUCKETS = 26 * 26; // 676

static inline int bigram_key(uint8_t a, uint8_t b) {
    return (a - 'a') * 26 + (b - 'a');
}

// Returns number of mismatches (capped at 2 for early exit)
static inline int verify_match_64(const uint8_t* text, const uint8_t* pat) {
    // Use AVX2: compare 32 bytes at a time
    __m256i t0 = _mm256_loadu_si256((const __m256i*)text);
    __m256i p0 = _mm256_loadu_si256((const __m256i*)pat);
    uint32_t neq0 = ~(uint32_t)_mm256_movemask_epi8(_mm256_cmpeq_epi8(t0, p0));
    int mis = __builtin_popcount(neq0);
    if (mis > 1) return mis;
    __m256i t1 = _mm256_loadu_si256((const __m256i*)(text + 32));
    __m256i p1 = _mm256_loadu_si256((const __m256i*)(pat + 32));
    uint32_t neq1 = ~(uint32_t)_mm256_movemask_epi8(_mm256_cmpeq_epi8(t1, p1));
    return mis + __builtin_popcount(neq1);
}

int main(int argc, char** argv) {
    if (argc < 4) {
        std::cerr << "Usage: " << argv[0] << " <text> <patterns> <output>" << std::endl;
        return 1;
    }

    fix_cpu_affinity();

    MappedFile f_pat(argv[2]);
    const uint8_t* p_ptr = (const uint8_t*)f_pat.data;
    uint32_t K = *(uint32_t*)p_ptr;
    p_ptr += 4;

    // Select two non-overlapping bigram positions from the range where patterns differ.
    // First, find positions where patterns are actually diverse (not all identical).
    int best_p1 = -1, best_p2 = -1;
    {
        struct PosScore { int pos; int max_bucket; };
        std::vector<PosScore> scores;
        for (int p = 0; p < PAT_LEN - 1; p++) {
            int cnt[BIGRAM_BUCKETS] = {0};
            for (uint32_t k = 0; k < K; k++) {
                int key = bigram_key(p_ptr[k * PAT_LEN + p], p_ptr[k * PAT_LEN + p + 1]);
                cnt[key]++;
            }
            int mx = *std::max_element(cnt, cnt + BIGRAM_BUCKETS);
            scores.push_back({p, mx});
        }
        // Sort by max_bucket ascending
        std::sort(scores.begin(), scores.end(),
                  [](const PosScore& a, const PosScore& b) { return a.max_bucket < b.max_bucket; });
        // Pick best two non-overlapping
        best_p1 = scores[0].pos;
        for (size_t i = 1; i < scores.size(); i++) {
            if (std::abs(scores[i].pos - best_p1) >= 2) {
                best_p2 = scores[i].pos;
                break;
            }
        }
        if (best_p1 > best_p2) std::swap(best_p1, best_p2);
        // Print diagnostics
        fprintf(stderr, "[FILTER] Bigram positions: %d (max=%d), %d (max=%d)\n",
                best_p1, scores[0].max_bucket, best_p2,
                (best_p2 == scores[1].pos) ? scores[1].max_bucket : scores[2].max_bucket);
        fprintf(stderr, "[FILTER] Top 5 positions by max_bucket:\n");
        for (int i = 0; i < std::min(5, (int)scores.size()); i++)
            fprintf(stderr, "  pos=%d max_bucket=%d\n", scores[i].pos, scores[i].max_bucket);
    }

    // If best bigram still has huge buckets, fall back to single-byte filter
    // which uses 26 buckets instead of 676 (coarser but still useful)
    bool use_single_byte = false;
    int sb_p1 = -1, sb_p2 = -1;
    int sb_cnt1[26], sb_cnt2[26];
    int sb_off1[27], sb_off2[27];
    std::vector<uint16_t> sb_data1, sb_data2;

    // Always try single-byte to compare
    {
        struct PosScore { int pos; int max_bucket; };
        std::vector<PosScore> scores;
        for (int p = 0; p < PAT_LEN; p++) {
            int cnt[26] = {0};
            for (uint32_t k = 0; k < K; k++)
                cnt[p_ptr[k * PAT_LEN + p] - 'a']++;
            int mx = *std::max_element(cnt, cnt + 26);
            scores.push_back({p, mx});
        }
        std::sort(scores.begin(), scores.end(),
                  [](const PosScore& a, const PosScore& b) { return a.max_bucket < b.max_bucket; });
        sb_p1 = scores[0].pos;
        for (size_t i = 1; i < scores.size(); i++) {
            if (scores[i].pos != sb_p1) {
                sb_p2 = scores[i].pos;
                break;
            }
        }
        fprintf(stderr, "[FILTER] Single-byte positions: %d (max=%d), %d (max=%d)\n",
                sb_p1, scores[0].max_bucket, sb_p2, scores[1].max_bucket);
    }

    // Decide: use bigram if best max_bucket < K/2, else single-byte
    // (In adversarial case, bigram max_bucket = K because all patterns identical at those positions)
    // For single-byte with 512 patterns, expected max_bucket ≈ 30
    // For bigram with diverse positions, expected max_bucket ≈ 3-5

    // Build bucket arrays: bucket[bigram_key] → list of pattern indices
    // Use CSR (compressed sparse row) format for cache-friendly access
    int bucket1_cnt[BIGRAM_BUCKETS] = {0};
    int bucket2_cnt[BIGRAM_BUCKETS] = {0};
    for (uint32_t k = 0; k < K; k++) {
        int key1 = bigram_key(p_ptr[k * PAT_LEN + best_p1], p_ptr[k * PAT_LEN + best_p1 + 1]);
        int key2 = bigram_key(p_ptr[k * PAT_LEN + best_p2], p_ptr[k * PAT_LEN + best_p2 + 1]);
        bucket1_cnt[key1]++;
        bucket2_cnt[key2]++;
    }

    // Prefix sums for CSR offsets
    int bucket1_off[BIGRAM_BUCKETS + 1], bucket2_off[BIGRAM_BUCKETS + 1];
    bucket1_off[0] = bucket2_off[0] = 0;
    for (int i = 0; i < BIGRAM_BUCKETS; i++) {
        bucket1_off[i + 1] = bucket1_off[i] + bucket1_cnt[i];
        bucket2_off[i + 1] = bucket2_off[i] + bucket2_cnt[i];
    }

    // Fill CSR data arrays
    std::vector<uint16_t> bucket1_data(K), bucket2_data(K);
    int fill1[BIGRAM_BUCKETS] = {0}, fill2[BIGRAM_BUCKETS] = {0};
    for (uint32_t k = 0; k < K; k++) {
        int key1 = bigram_key(p_ptr[k * PAT_LEN + best_p1], p_ptr[k * PAT_LEN + best_p1 + 1]);
        int key2 = bigram_key(p_ptr[k * PAT_LEN + best_p2], p_ptr[k * PAT_LEN + best_p2 + 1]);
        bucket1_data[bucket1_off[key1] + fill1[key1]++] = (uint16_t)k;
        bucket2_data[bucket2_off[key2] + fill2[key2]++] = (uint16_t)k;
    }

    MappedFile f_text(argv[1]);
    const uint8_t* t_ptr = (const uint8_t*)f_text.data;
    uint64_t N = *(uint64_t*)t_ptr;
    const uint8_t* text = t_ptr + 8;

    auto start_time = std::chrono::high_resolution_clock::now();

    std::atomic<bool> power_done(false);
    std::thread power_thread([&power_done]() { power_monitor(power_done); });

    uint64_t total_matches = 0;

    if (N >= (uint64_t)PAT_LEN) {
        int64_t total_positions = (int64_t)(N - PAT_LEN + 1);
        const int fp1 = best_p1, fp2 = best_p2;
        const uint16_t* b1data = bucket1_data.data();
        const uint16_t* b2data = bucket2_data.data();

        #pragma omp parallel reduction(+:total_matches)
        {
            int tid = omp_get_thread_num();
            int nthreads = omp_get_num_threads();
            int64_t chunk = (total_positions + nthreads - 1) / nthreads;
            int64_t my_start = (int64_t)tid * chunk;
            int64_t my_end = std::min(my_start + chunk, total_positions);

            for (int64_t pos = my_start; pos < my_end; pos++) {
                if (__builtin_expect((pos & POWER_CHECK_MASK) == 0, 0)) {
                    int us = g_throttle_us.load(std::memory_order_relaxed);
                    if (__builtin_expect(us > 0, 0)) usleep(us);
                }

                // Bucket 1: bigram at position fp1
                int key1 = bigram_key(text[pos + fp1], text[pos + fp1 + 1]);
                int start1 = bucket1_off[key1], end1 = bucket1_off[key1 + 1];
                for (int i = start1; i < end1; i++) {
                    uint16_t pi = b1data[i];
                    if (verify_match_64(text + pos, p_ptr + (uint32_t)pi * PAT_LEN) <= 1) {
                        total_matches++;
                    }
                }

                // Bucket 2: bigram at position fp2
                // Skip patterns already checked via bucket1 (their bigram at fp1 matches text)
                int key2 = bigram_key(text[pos + fp2], text[pos + fp2 + 1]);
                int start2 = bucket2_off[key2], end2 = bucket2_off[key2 + 1];
                for (int i = start2; i < end2; i++) {
                    uint16_t pi = b2data[i];
                    // If this pattern's bigram at fp1 matches text, it was in bucket1
                    const uint8_t* pat = p_ptr + (uint32_t)pi * PAT_LEN;
                    if (pat[fp1] == text[pos + fp1] && pat[fp1 + 1] == text[pos + fp1 + 1])
                        continue;
                    if (verify_match_64(text + pos, pat) <= 1) {
                        total_matches++;
                    }
                }
            }
        }
    }

    auto end_time = std::chrono::high_resolution_clock::now();
    double duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count() / 1000.0;
    power_done.store(true, std::memory_order_relaxed);
    if (power_thread.joinable()) power_thread.join();

    int out_fd = open(argv[3], O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (out_fd >= 0) {
        write(out_fd, &total_matches, sizeof(total_matches));
        close(out_fd);
    }

    std::cout << "Time: " << duration << "s" << std::endl;
    return 0;
}