#pragma once
#include "llama.h"
#include <cstdint>
#include <string>
#include <vector>
//
// suffix cache reuse (SCR)
//
// after an edit in the middle of a prompt, the cached attention KV of the text that survived is moved to its new
// positions (K re-rotated), the recurrent state of a hybrid model continues from a snapshot taken inside the last
// moved span, and only the rest is prefilled. the moved rows were computed over the pre-edit text, so the result
// approximates a fresh prefill of the edited prompt.
//
struct common_scr_params {
int32_t min_span = 256; // shortest span worth moving
int32_t max_runs = 64; // most spans moved per request
int32_t seed = 32; // n-gram length that seeds a match
};
// a span of the cached tokens [old_start, old_start + n) that reappears at [new_start, new_start + n)
struct common_scr_run {
int32_t new_start;
int32_t old_start;
int32_t n;
};
struct common_scr_plan {
int32_t n_prefix = 0; // leading tokens equal in both prompts
int32_t restore = 0; // position the prefill restarts from: a snapshot at or before n_prefix
int32_t adopt = -1; // old position of the snapshot the recurrent state continues from, -1 without one
std::vector<common_scr_run> runs; // in new order, all at or after n_prefix
bool empty() const { return runs.empty(); }
int32_t resume() const; // new position the final prefill starts at
int32_t n_reused() const;
std::string str() const;
};
// plans the reuse of a state that holds old[0, n_old) for the prompt cur.
// snapshots: sorted old positions q at which the recurrent state after old[0, q) can be restored (the live state
// counts as old.size()); with recurrent = false every position can be resumed from.
// LLAMA_TOKEN_NULL never matches. the plan always leaves at least one token of cur to prefill.
common_scr_plan common_scr_plan_build(
const std::vector<llama_token> & old,
const std::vector<llama_token> & cur,
const std::vector<int32_t> & snapshots,
bool recurrent,
const common_scr_params & params);
// rearranges the attention KV of seq_id for a plan: drops every cell at or after plan.restore that no run moves,
// then moves the runs. returns 1 when done, 0 when the runs cannot be moved in any order (nothing changed),
// -1 when a move failed after the cache was changed (the sequence must then be cleared).
int32_t common_scr_apply_kv(llama_context * ctx, llama_seq_id seq_id, const common_scr_plan & plan);