代码

// This code is AI-generated. (AI 生成的代码)
// router32: longest-prefix-match for a BGP table.
//
// A 16/8/8 multibit trie.  Level 1 maps the top 16 bits; levels 2 and 3 are
// 256-entry tables for the next 8 and last 8 bits.  The level-2 and level-3
// tables are run-length compressed into (bitmask, prefix-sum, pointer) triples
// so a query only touches a few cache-friendly arrays instead of 48 MB of raw
// tables; vpcmpeqb-style popcount over the 32-bit run mask locates the entry.
#pragma GCC optimize("Ofast,unroll-loops")
#pragma GCC target("arch=skylake")
#include "router.h"
#include <arpa/inet.h>
#include <stdint.h>
#include <string.h>

typedef uint16_t u16;
typedef uint32_t u32;

#define T16_SIZE 65536
#define N24 32768
#define N32 16384
#define L2_RUNS (4u << 20)
#define L3_RUNS (2u << 20)

static u16 t16[T16_SIZE];
static u32 t24[N24][256];
static u32 t32[N32][256];
static u32 n24_cnt, n32_cnt;

static u16 l2_ptr[L2_RUNS];
static u32 l2_run_cnt;
static u32 l2_sum[N24][8];
static u32 l2_bit[N24][8];

static u16 l3_ptr[L3_RUNS];
static u32 l3_run_cnt;
static u32 l3_sum[N32][8];
static u32 l3_bit[N32][8];

static u32 nh_key[1 << 17];
static u16 nh_val[1 << 17];
static u32 nh[1 << 16];
static u32 nh_count = 1;

static u16 intern(u32 x) {
    u32 h = (x * 2654435761u) >> 15;
    while (nh_key[h] && nh_key[h] != x + 1) h = (h + 1) & 0x1ffffu;
    if (nh_key[h]) return nh_val[h];
    nh_key[h] = x + 1;
    nh_val[h] = (u16)nh_count;
    nh[nh_count] = x;
    return (u16)nh_count++;
}

static void fill16(u16 *d, u32 cnt, u16 v) {
    for (u32 i = 0; i < cnt; ++i) d[i] = v;
}
static void fill32(u32 *d, u32 cnt, u32 v) {
    for (u32 i = 0; i < cnt; ++i) d[i] = v;
}

void init(int n, int q, const RoutingTableEntry *tbl) {
    (void)q;
    memset(t16, 0, sizeof(t16));
    n24_cnt = n32_cnt = 0;
    l2_run_cnt = l3_run_cnt = 0;
    nh_count = 1;

    for (int i = 0; i < n; ++i) {
        u32 addr = htonl(tbl[i].addr);
        u32 len = tbl[i].len;
        u16 id = intern(tbl[i].nexthop);
        if (len <= 16) {
            fill16(t16 + (addr >> 16), 1u << (16 - len), id);
        } else if (len <= 24) {
            u32 top = addr >> 16, v = t16[top];
            u32 idx;
            if (!(v & 0x8000)) {
                idx = n24_cnt++;
                fill32(t24[idx], 256, v);
                t16[top] = (u16)(0x8000 | idx);
            } else idx = v & 0x7fff;
            fill32(t24[idx] + ((addr >> 8) & 255), 1u << (24 - len), id);
        } else {
            u32 top = addr >> 16, v = t16[top];
            u32 idx24;
            if (!(v & 0x8000)) {
                idx24 = n24_cnt++;
                fill32(t24[idx24], 256, v);
                t16[top] = (u16)(0x8000 | idx24);
            } else idx24 = v & 0x7fff;
            u32 mid = (addr >> 8) & 255, v2 = t24[idx24][mid];
            u32 idx32;
            if (!(v2 & 0x80000000u)) {
                idx32 = n32_cnt++;
                fill32(t32[idx32], 256, v2);
                t24[idx24][mid] = 0x80000000u | idx32;
            } else idx32 = v2 & 0x7fffffffu;
            fill32(t32[idx32] + (addr & 255), 1u << (32 - len), id);
        }
    }

    // compress level-2 tables
    for (u32 idx = 0; idx < n24_cnt; ++idx) {
        u16 *p = l2_ptr + l2_run_cnt;
        u16 cur = (t24[idx][0] >> 31) ? (u16)(0x8000 | (t24[idx][0] & 0x7fff)) : (u16)t24[idx][0];
        u32 cnt = 0;
        p[0] = cur;
        l2_sum[idx][0] = l2_run_cnt;
        for (u32 i = 1; i < 256; ++i) {
            if ((i & 31) == 0) l2_sum[idx][i >> 5] = l2_run_cnt + cnt;
            u16 e = (t24[idx][i] >> 31) ? (u16)(0x8000 | (t24[idx][i] & 0x7fff)) : (u16)t24[idx][i];
            if (e != cur) { l2_bit[idx][i >> 5] |= 1u << (i & 31); cur = e; p[++cnt] = e; }
        }
        l2_run_cnt += cnt + 1;
    }
    // compress level-3 tables
    for (u32 idx = 0; idx < n32_cnt; ++idx) {
        u16 *p = l3_ptr + l3_run_cnt;
        u16 cur = (u16)t32[idx][0];
        u32 cnt = 0;
        p[0] = cur;
        l3_sum[idx][0] = l3_run_cnt;
        for (u32 i = 1; i < 256; ++i) {
            if ((i & 31) == 0) l3_sum[idx][i >> 5] = l3_run_cnt + cnt;
            u16 e = (u16)t32[idx][i];
            if (e != cur) { l3_bit[idx][i >> 5] |= 1u << (i & 31); cur = e; p[++cnt] = e; }
        }
        l3_run_cnt += cnt + 1;
    }
}

unsigned query(unsigned addr_raw) {
    u32 addr = htonl(addr_raw);
    u32 v = t16[addr >> 16];
    if (v & 0x8000) {
        u32 idx = v & 0x7fff;
        u32 a = (addr >> 8) & 255;
        u32 chunk = a >> 5;
        u32 off = l2_sum[idx][chunk] + __builtin_popcount(l2_bit[idx][chunk] & ((2u << (a & 31)) - 1));
        u16 e = l2_ptr[off];
        if (e & 0x8000) {
            u32 idx2 = e & 0x7fff;
            u32 a2 = addr & 255;
            u32 chunk2 = a2 >> 5;
            u32 off2 = l3_sum[idx2][chunk2] + __builtin_popcount(l3_bit[idx2][chunk2] & ((2u << (a2 & 31)) - 1));
            return nh[l3_ptr[off2]];
        }
        return nh[e];
    }
    return nh[v];
}

评测结果