#include "scr.h"
#include "testing.h"
#include <algorithm>
#include <cstdint>
#include <vector>
static std::vector<llama_token> make_prompt(int n_blocks, bool reordered) {
std::vector<llama_token> out;
out.reserve(n_blocks * 11);
for (int i = 0; i < n_blocks; ++i) {
const int block = reordered ? (i * 37 + 1) % n_blocks : i;
out.push_back((llama_token) (100000 + block));
for (int j = 0; j < 10; ++j) {
out.push_back((llama_token) (block * 1000 + j + 1));
}
}
return out;
}
static common_scr_params fixture_params(int max_runs = -1) {
common_scr_params params;
params.min_span = 8;
params.seed = 4;
if (max_runs >= 0) {
params.max_runs = max_runs;
}
return params;
}
static void check_exact_runs(
testing & t,
const std::vector<llama_token> & old,
const std::vector<llama_token> & cur,
const common_scr_plan & plan) {
for (const auto & run : plan.runs) {
t.assert_true("run leaves one current token",
run.new_start + run.n < (int32_t) cur.size());
t.assert_true("run is in old bounds",
run.old_start >= 0 && run.old_start + run.n <= (int32_t) old.size());
t.assert_true("run is in current bounds",
run.new_start >= 0 && run.new_start + run.n <= (int32_t) cur.size());
for (int32_t i = 0; i < run.n; ++i) {
t.assert_equal("run matches both ranges",
old[run.old_start + i], cur[run.new_start + i]);
t.assert_true("null token is never reused", old[run.old_start + i] != LLAMA_TOKEN_NULL);
}
}
}
int main() {
testing t;
t.test("default cap keeps 32 disjoint spans", [](testing & t) {
const auto old = make_prompt(32, false);
const auto cur = make_prompt(32, true);
const auto plan = common_scr_plan_build(old, cur, {}, false, fixture_params());
t.assert_equal(32, (int) plan.runs.size());
check_exact_runs(t, old, cur, plan);
});
t.test("explicit cap six discards the rest", [](testing & t) {
const auto old = make_prompt(70, false);
const auto cur = make_prompt(70, true);
const auto plan = common_scr_plan_build(old, cur, {}, false, fixture_params(6));
t.assert_equal(6, (int) plan.runs.size());
check_exact_runs(t, old, cur, plan);
});
t.test("default cap is exactly 64", [](testing & t) {
const auto old = make_prompt(70, false);
const auto cur = make_prompt(70, true);
const auto plan = common_scr_plan_build(old, cur, {}, false, fixture_params());
t.assert_equal(64, (int) plan.runs.size());
check_exact_runs(t, old, cur, plan);
});
t.test("recurrent reuse ends at a real snapshot", [](testing & t) {
const auto old = make_prompt(32, false);
const auto cur = make_prompt(32, true);
std::vector<int32_t> snapshots;
for (int i = 0; i < 32; ++i) {
snapshots.push_back(i * 11 + 11);
}
const auto plan = common_scr_plan_build(old, cur, snapshots, true, fixture_params());
t.assert_true(!plan.runs.empty());
t.assert_true(std::find(snapshots.begin(), snapshots.end(), plan.adopt) != snapshots.end());
t.assert_equal(plan.adopt - plan.runs.back().old_start, plan.runs.back().n);
check_exact_runs(t, old, cur, plan);
});
t.test("no snapshot means no recurrent reuse", [](testing & t) {
const auto old = make_prompt(32, false);
const auto cur = make_prompt(32, true);
const auto plan = common_scr_plan_build(old, cur, {}, true, fixture_params());
t.assert_true(plan.runs.empty());
t.assert_equal(0, plan.n_reused());
});
t.test("null token breaks reuse", [](testing & t) {
auto old = make_prompt(4, false);
auto cur = make_prompt(4, true);
old[1] = LLAMA_TOKEN_NULL;
const auto plan = common_scr_plan_build(old, cur, {}, false, fixture_params());
check_exact_runs(t, old, cur, plan);
t.assert_true(!plan.runs.empty());
t.assert_true(plan.n_reused() < 4 * 11);
});
return t.summary();
}