代码

#define DUCK 1

#ifdef DUCK
#include "router.h"
#endif

#include <cstdio>
#include <cstdlib>
#include <cstdint>
#include <cstddef>
#include <cstring>
#include <algorithm>
#include <arpa/inet.h>

#ifdef DUCK
#define DUCK_HTONL(x) htonl(x)
#else 
#define DUCK_HTONL(x) (x)
#endif

void print_ipv4(uint32_t addr)
{
    addr = htonl(addr);
    char str[INET6_ADDRSTRLEN];
    inet_ntop(AF_INET, &addr, str, sizeof(str));
    printf("%s", str);
}

struct rt_entry
{
    uint32_t net;
    uint32_t prefixlen;
    uint32_t nexthop;
};

struct rt_node
{
    const rt_entry *entry = nullptr;
    rt_node *prev = nullptr, *next[2] = { nullptr, nullptr };
    uint8_t nbit = 0; // which bit is used by next[]?
    uint32_t prefix = 0;
};

struct rt
{
    rt_node root;
};

static inline uint32_t first_diff(uint32_t a, uint32_t b)
{
    int i = 0;
    uint32_t c = a ^ b;
    if (!c) return -1;

    if ((c & 0xffff0000) == 0) { i += 16; } else { c >>= 16; }
    if ((c & 0xff00) == 0) { i += 8; } else { c >>= 8; }
    if ((c & 0xf0) == 0) { i += 4; } else { c >>= 4; }
    if ((c & 0xc) == 0) { i += 2; } else { c >>= 2; }
    if ((c & 0x2) == 0) { i += 1; }
    return i;
}

static inline uint32_t prefixlen2mask(uint32_t len)
{
    if (len == 0) return 0;
    return (uint32_t)-1 << (32 - len);
}

static rt_node *_pct_insert_as_prev(rt_node *node, const rt_entry *e)
{
    // insert a node into this path (between node->prev and node)
    uint32_t nbit = std::min(e->prefixlen, first_diff(node->prefix, DUCK_HTONL(e->net)));
    rt_node *prev = node->prev;
    //printf("case1 %u %u %u\n", prev->nbit, nbit, node->nbit);
    rt_node *&prev_next = node == prev->next[0] ? prev->next[0] : prev->next[1];
    rt_node *mid = new rt_node();
    prev_next = mid;
    mid->prev = prev;
    mid->nbit = nbit;
    mid->prefix = node->prefix & prefixlen2mask(nbit);
    int i = ((node->prefix << mid->nbit) & 0x80000000) != 0;
    mid->next[i] = node;
    node->prev = mid;
    return mid;
}

void rt_insert(rt_node *root, const rt_entry *e)
{
    //printf("+ ");
    //print_ipv4(e->net);
    //printf("/%u\n", e->prefixlen);
    rt_node *node = root;
    while (true)
    {
        if (node->nbit == e->prefixlen)
        {
            if (node->prefix == DUCK_HTONL(e->net))
            {
                break;
            }
            else
            {
                node = _pct_insert_as_prev(node, e);
            }
        }
        else if (node->nbit > e->prefixlen)
        {
            node = _pct_insert_as_prev(node, e);
        }
        else
        {
            int i = ((DUCK_HTONL(e->net) << node->nbit) & 0x80000000) != 0;
            rt_node *&next = node->next[i];
            if (!next)
            {
                next = new rt_node();
                next->prev = node;
                next->nbit = e->prefixlen;
                next->prefix = DUCK_HTONL(e->net);
                node = next;
                break;
            }
            else
            {
                uint32_t mask = prefixlen2mask(next->nbit);
                if ((next->prefix & mask) == (DUCK_HTONL(e->net) & mask))
                {
                    node = next;
                }
                else
                {
                    node = _pct_insert_as_prev(next, e);
                }
            }
        }
    }
    node->entry = e;
}

const rt_entry *rt_find(rt_node *root, uint32_t addr)
{
    rt_node *node = root;
    while (true)
    {
        // printf("* ");
        // print_ipv4(node->prefix);
        // printf("/%u\n", node->nbit);
        int i = ((addr << node->nbit) & 0x80000000) != 0;
        if (!node->next[i]) break;
        node = node->next[i];
    }
    while (node)
    {
        if (node->entry)
        {
            uint32_t mask = prefixlen2mask(node->entry->prefixlen);
            if ((DUCK_HTONL(node->entry->net) & mask) == (addr & mask))
            {
                return node->entry;
            }
        }
        node = node->prev;
    }
    return nullptr;
}

#ifdef DUCK
rt r;

void init(int n, int q, const RoutingTableEntry *a) {
    for (size_t i = 0; i < n; ++i)
    {
        rt_insert(&r.root, (const rt_entry *)&a[i]);
    }
}

unsigned query(unsigned addr) {
    const rt_entry *e = rt_find(&r.root, htonl(addr));
    if (e) return e->nexthop;
	return 0;
}
#else
void print_spaces(int n)
{
    for (int i = 0; i < n; ++i) printf(" ");
}

void print(rt_node *node, int level)
{
    if (!node)
    {
        print_spaces(level);
        printf("(null)\n");
        return;
    }
    print_spaces(level);
    printf("+ ");
    print_ipv4(node->prefix);
    printf("/%u\n", node->nbit);
    print_spaces(level);
    printf("| left:\n");
    print(node->next[0], level + 2);
    print_spaces(level);
    printf("|right:\n");
    print(node->next[1], level + 2);
}

int main()
{
    FILE *f = fopen("rt.bin", "rb");
    fseek(f, 0, SEEK_END);
    size_t size = ftell(f) / sizeof(rt_entry);
    fseek(f, 0, SEEK_SET);
    rt_entry *entries = new rt_entry[size];
    fread(entries, sizeof(rt_entry), size, f);
    fclose(f);
    rt rt;
    for (size_t i = 0; i < size; ++i)
    {
        rt_insert(&rt.root, &entries[i]);
    }
    //print(&rt.root, 0);
    while (true)
    {
        printf("> ");
        char ip[100];
        fgets(ip, sizeof(ip), stdin);
        char &c = ip[strlen(ip) - 1];
        if (c == '\n')
        {
            c = 0;
        }
        //printf("%s\n", ip);
        uint32_t addr;
        inet_pton(AF_INET, ip, &addr);
        addr = ntohl(addr);
        //printf("%08x\n", addr);
        const rt_entry *e = rt_find(&rt.root, addr);
        printf("Next hop is ");
        print_ipv4(e->nexthop);
        printf("\n");
    }
    return 0;
}
#endif

评测结果