Foundry/rosa/tests/test_rosa.c

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_rosa.c — Unit tests for suffix automaton
 */

#include "lib/test.h"
#include "rosa/rosa.h"
#include <stdlib.h>
#include <string.h>

/* ================================================================
 * Binary automaton tests
 * ================================================================ */

static void test_sa_new_free(void) {
    sa_automaton *sa = sa_new(16);
    ASSERT_NOT_NULL(sa);
    ASSERT_INT_EQ(1, sa->n_states);  /* root state */
    ASSERT_INT_EQ(0, sa->seq_len);
    sa_free(sa);
}

static void test_sa_extend_single(void) {
    sa_automaton *sa = sa_new(16);
    sa_extend(sa, 0, 0);
    ASSERT_TRUE(sa->n_states >= 2);  /* root + at least 1 new */
    ASSERT_TRUE(sa->n_states <= 3);  /* at most root + new + clone */
    ASSERT_INT_EQ(1, sa->seq_len);
    sa_free(sa);
}

static void test_sa_extend_sequence(void) {
    sa_automaton *sa = sa_new(32);
    /* Extend with "01101" */
    int seq[] = {0, 1, 1, 0, 1};
    int n = 5;
    for (int i = 0; i < n; i++)
        sa_extend(sa, seq[i], i);

    ASSERT_INT_EQ(n, sa->seq_len);
    ASSERT_TRUE(sa->n_states <= 2 * n + 1);  /* SA guarantee */
    sa_free(sa);
}

static void test_sa_query_exact(void) {
    sa_automaton *sa = sa_new(32);
    /* Build automaton over "01101" */
    int seq[] = {0, 1, 1, 0, 1};
    for (int i = 0; i < 5; i++)
        sa_extend(sa, seq[i], i);

    /* Query for "1101" — should find a match ending at position 4 */
    int query[] = {1, 1, 0, 1};
    int state = 0, match_len = 0;
    int endpos = -1;
    for (int i = 0; i < 4; i++)
        endpos = sa_query_step(sa, query[i], &state, &match_len);

    ASSERT_TRUE(endpos >= 0);        /* found a match */
    ASSERT_TRUE(match_len >= 4);     /* matched all 4 symbols */
    sa_free(sa);
}

static void test_sa_query_no_match(void) {
    sa_automaton *sa = sa_new(16);
    /* Build over "000" */
    for (int i = 0; i < 3; i++)
        sa_extend(sa, 0, i);

    /* Query for "1" — no 1s in the automaton */
    int state = 0, match_len = 0;
    int endpos = sa_query_step(sa, 1, &state, &match_len);

    /* After querying 1 on an automaton with only 0s, we should be back at root */
    ASSERT_INT_EQ(0, match_len);
    (void)endpos;  /* may be -1 or root endpos */
    sa_free(sa);
}

static void test_sa_query_partial(void) {
    sa_automaton *sa = sa_new(32);
    /* Build over "0110" */
    int seq[] = {0, 1, 1, 0};
    for (int i = 0; i < 4; i++)
        sa_extend(sa, seq[i], i);

    /* Query "10110" — only "0110" suffix should match (len 4) */
    int query[] = {1, 0, 1, 1, 0};
    int state = 0, match_len = 0;
    int endpos = -1;
    for (int i = 0; i < 5; i++)
        endpos = sa_query_step(sa, query[i], &state, &match_len);

    ASSERT_TRUE(endpos >= 0);
    /* The match should be a suffix of the query that exists in "0110" */
    ASSERT_TRUE(match_len > 0);
    ASSERT_TRUE(match_len <= 4);
    sa_free(sa);
}

static void test_sa_query_successor(void) {
    sa_automaton *sa = sa_new(32);
    /* Build "01010" */
    int seq[] = {0, 1, 0, 1, 0};
    for (int i = 0; i < 5; i++)
        sa_extend(sa, seq[i], i);

    /* Query "01" — should match, endpos tells us where "01" last ended */
    int state = 0, match_len = 0;
    sa_query_step(sa, 0, &state, &match_len);
    int endpos = sa_query_step(sa, 1, &state, &match_len);

    ASSERT_TRUE(endpos >= 0);
    ASSERT_INT_EQ(2, match_len);
    /* Successor position would be endpos + 1 */
    ASSERT_TRUE(endpos + 1 < 5);  /* valid successor exists */
    sa_free(sa);
}

/* ================================================================
 * Token-level automaton tests
 * ================================================================ */

static void test_sa_token_new_free(void) {
    sa_token_automaton *sa = sa_token_new(16, 65536);
    ASSERT_NOT_NULL(sa);
    ASSERT_INT_EQ(1, sa->n_states);
    ASSERT_INT_EQ(65536, sa->alphabet_size);
    sa_token_free(sa);
}

static void test_sa_token_extend_query(void) {
    sa_token_automaton *sa = sa_token_new(32, 65536);

    /* Build over token sequence [100, 200, 300, 100, 200] */
    int seq[] = {100, 200, 300, 100, 200};
    for (int i = 0; i < 5; i++)
        sa_token_extend(sa, seq[i], i);

    ASSERT_INT_EQ(5, sa->seq_len);
    ASSERT_TRUE(sa->n_states <= 2 * 5 + 1);

    /* Query [100, 200] — should match positions 0-1 and 3-4 */
    int state = 0, match_len = 0;
    sa_token_query_step(sa, 100, &state, &match_len);
    int endpos = sa_token_query_step(sa, 200, &state, &match_len);

    ASSERT_TRUE(endpos >= 0);
    ASSERT_INT_EQ(2, match_len);
    sa_token_free(sa);
}

static void test_sa_token_large_sequence(void) {
    sa_token_automaton *sa = sa_token_new(64, 65536);

    /* Build over 10K tokens with some repetition */
    int n = 10000;
    for (int i = 0; i < n; i++)
        sa_token_extend(sa, i % 1000, i);  /* 1000 unique tokens, repeating */

    ASSERT_INT_EQ(n, sa->seq_len);
    ASSERT_TRUE(sa->n_states <= 2 * n + 1);

    /* Query should work */
    int state = 0, match_len = 0;
    int endpos = sa_token_query_step(sa, 42, &state, &match_len);
    ASSERT_TRUE(endpos >= 0);

    sa_token_free(sa);
}

/* ================================================================
 * Position map tests
 * ================================================================ */

static void test_position_map(void) {
    sa_position_map *map = sa_position_map_new(4);
    ASSERT_NOT_NULL(map);
    ASSERT_INT_EQ(0, map->count);

    int pos0 = sa_position_map_append(map, 0, 0, 0);
    int pos1 = sa_position_map_append(map, 0, 1, 5);
    int pos2 = sa_position_map_append(map, 1, 0, 0);

    ASSERT_INT_EQ(0, pos0);
    ASSERT_INT_EQ(1, pos1);
    ASSERT_INT_EQ(2, pos2);
    ASSERT_INT_EQ(3, map->count);

    const sa_position_info *info = sa_position_map_get(map, 1);
    ASSERT_NOT_NULL(info);
    ASSERT_INT_EQ(0, info->conv_idx);
    ASSERT_INT_EQ(1, info->msg_idx);
    ASSERT_INT_EQ(5, info->token_offset);

    ASSERT_TRUE(sa_position_map_get(map, -1) == NULL);
    ASSERT_TRUE(sa_position_map_get(map, 99) == NULL);

    sa_position_map_free(map);
}

/* ================================================================
 * Main
 * ================================================================ */

int main(void) {
    RUN_TEST(sa_new_free);
    RUN_TEST(sa_extend_single);
    RUN_TEST(sa_extend_sequence);
    RUN_TEST(sa_query_exact);
    RUN_TEST(sa_query_no_match);
    RUN_TEST(sa_query_partial);
    RUN_TEST(sa_query_successor);
    RUN_TEST(sa_token_new_free);
    RUN_TEST(sa_token_extend_query);
    RUN_TEST(sa_token_large_sequence);
    RUN_TEST(position_map);
    TEST_REPORT();
    return __ft_failed > 0 ? 1 : 0;
}