// Suffix cache reuse (SCR): after an edit in the middle of a prompt, the K/V of the text that survived the edit is
// moved to its new positions instead of being read again. Ported from llama-cpp-bonsai's common/scr.cpp.
//
// The moved cells were computed over the text before the edit, so the result approximates a fresh read. The GDN
// layers have no cells to move: their state continues from a conversation checkpoint at the end of the last moved
// run (the plan's `adopt`), and everything after that run is read from it. Strata's indexer pools 4 cells into one
// key and keeps only the pooled keys, so a run starts and ends on a block boundary of the OLD prompt and moves by a
// multiple of 4: up to 3 pad tokens are read before a run to make its shift one (DeepSeek V4.1's recipe pads to an
// even shift the same way). The moved blocks then keep their cells, and their pooled keys only turn by the shift.
#pragma once
#include <algorithm>
#include <cstdint>
#include <string>
#include <unordered_map>
#include <vector>
namespace strata::core {
struct ScrParams {
int64_t min_span = 256; // shortest run worth moving
int64_t max_runs = 64; // most runs moved per request
int64_t seed = 32; // n-gram length that seeds a match
int64_t block = 4; // the indexer's cells per pooled key: read parts start and end on these
};
// the cells old[old_start, old_start + n) reappear at [new_start, new_start + n) of the prompt, after `pad` pad tokens
// that the read inserts right before the run; `at` is where the run lands in the padded (physical) prompt
struct ScrRun {
int64_t new_start = 0, old_start = 0, n = 0;
int64_t pad = 0, at = 0;
int64_t shift() const { return at - old_start; }
};
struct ScrPlan {
int64_t n_prefix = 0; // leading tokens equal in both prompts
int64_t restore = 0; // the checkpoint the read restarts from (at or before n_prefix)
int64_t adopt = -1; // old length of the checkpoint the GDN state continues from after the last run
std::vector<ScrRun> runs; // in new order, at or after restore
bool empty() const { return runs.empty(); }
// the prompt position the final read starts at, and the same in the padded prompt
int64_t resume() const { return runs.empty() ? restore : runs.back().new_start + runs.back().n; }
int64_t resume_at() const { return runs.empty() ? restore : runs.back().at + runs.back().n; }
int64_t pads() const {
int64_t n = 0;
for (const ScrRun& r : runs) n += r.pad;
return n;
}
int64_t reused() const {
int64_t n = 0;
for (const ScrRun& r : runs) n += r.n;
return n;
}
std::string str() const {
std::string s = "prefix " + std::to_string(n_prefix) + ", restore " + std::to_string(restore) + ", adopt " +
std::to_string(adopt);
for (const ScrRun& r : runs)
s += ", R " + std::to_string(r.new_start) + " <- " + std::to_string(r.old_start) + " x " + std::to_string(r.n) +
(r.pad ? " pad " + std::to_string(r.pad) : "");
return s;
}
};
namespace scr_detail {
// greedy maximal spans of cur[from, ...) found in old[from, ...), seeded by n-gram hashes, each old token used once
template <class A, class B>
std::vector<ScrRun> find_spans(const std::vector<A>& old, const std::vector<B>& cur, int64_t from, const ScrParams& p) {
std::vector<ScrRun> spans;
const int64_t n_old = (int64_t) old.size(), n_cur = (int64_t) cur.size(), k = p.seed;
if (n_old - from < p.min_span || n_cur - from < p.min_span) return spans;
auto gram = [k](const auto& v, int64_t i) {
uint64_t h = 1469598103934665603ULL;
for (int64_t j = 0; j < k; ++j) h = (h ^ (uint32_t) v[(size_t) (i + j)]) * 1099511628211ULL;
return h;
};
constexpr size_t max_bucket = 16; // the first few old positions of every seed n-gram
std::unordered_map<uint64_t, std::vector<int64_t>> index;
index.reserve((size_t) (n_old - from));
for (int64_t o = from; o + k <= n_old; ++o) {
auto& b = index[gram(old, o)];
if (b.size() < max_bucket) b.push_back(o);
}
std::vector<bool> used((size_t) n_old, false);
for (int64_t q = from; q + k <= n_cur;) {
const auto it = index.find(gram(cur, q));
int64_t best_o = -1, best_n = 0;
if (it != index.end())
for (const int64_t o : it->second) {
int64_t n = 0;
while (q + n < n_cur && o + n < n_old && !used[(size_t) (o + n)] &&
(int64_t) cur[(size_t) (q + n)] == (int64_t) old[(size_t) (o + n)])
++n;
if (n > best_n) { best_n = n; best_o = o; }
}
if (best_n >= p.min_span) {
spans.push_back({q, best_o, best_n});
std::fill(used.begin() + best_o, used.begin() + best_o + best_n, true);
q += best_n;
} else {
++q;
}
}
return spans;
}
inline int64_t round_up(int64_t x, int64_t b) { return (x + b - 1) / b * b; }
inline int64_t round_down(int64_t x, int64_t b) { return x / b * b; }
} // namespace scr_detail
// Plans the reuse of a session that holds old[0, n_old) for the prompt `cur`. `checkpoints`: the old lengths whose
// GDN state can be restored (the live state counts as old.size()). A run starts on a block boundary of the new
// prompt; every run but the last also ends on one, and the last one ends at a checkpoint. The plan always leaves at
// least one token of `cur` to read.
template <class A, class B>
ScrPlan scr_plan(const std::vector<A>& old, const std::vector<B>& cur, std::vector<int64_t> checkpoints,
const ScrParams& p) {
using namespace scr_detail;
ScrPlan plan;
const int64_t n_cur = (int64_t) cur.size(), bl = std::max<int64_t>(p.block, 1);
std::sort(checkpoints.begin(), checkpoints.end());
int64_t a = 0;
while (a < (int64_t) old.size() && a < n_cur && (int64_t) old[(size_t) a] == (int64_t) cur[(size_t) a]) ++a;
plan.n_prefix = a;
for (const int64_t q : checkpoints)
if (q <= a) plan.restore = std::max(plan.restore, q);
auto spans = find_spans(old, cur, a, p);
// a run starts on a block boundary of the old prompt and (but the last) ends on one
for (ScrRun& s : spans) {
const int64_t cut = round_up(s.old_start, bl) - s.old_start;
s.new_start += cut;
s.old_start += cut;
s.n -= cut;
s.n = round_down(s.old_start + s.n, bl) - s.old_start;
}
spans.erase(std::remove_if(spans.begin(), spans.end(), [&](const ScrRun& s) { return s.n < p.min_span; }),
spans.end());
// the largest runs, back in prompt order
std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.n > y.n; });
if ((int64_t) spans.size() > p.max_runs) spans.resize((size_t) p.max_runs);
std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.new_start < y.new_start; });
// the last run ends at a checkpoint, which also ends the reuse: everything after it is read from that state.
// It was cut back to a block boundary, so a checkpoint up to block - 1 cells past it may still match.
while (!spans.empty()) {
ScrRun& last = spans.back();
int64_t end = -1;
for (const int64_t q : checkpoints) {
const int64_t n = q - last.old_start;
if (n < p.min_span || n >= last.n + bl || last.new_start + n >= n_cur || q <= end) continue;
if (n <= last.n || std::equal(old.begin() + last.old_start + last.n, old.begin() + q,
cur.begin() + last.new_start + last.n,
[](const A& x, const B& y) { return (int64_t) x == (int64_t) y; }))
end = q;
}
if (end >= 0) {
last.n = end - last.old_start;
plan.adopt = end;
break;
}
spans.pop_back();
}
// the pads: walking the runs in prompt order, each gets what makes its padded position a multiple of the block
// away from its old one
int64_t pads = 0;
for (ScrRun& s : spans) {
s.pad = ((s.old_start - (s.new_start + pads)) % bl + bl) % bl;
pads += s.pad;
s.at = s.new_start + pads;
}
plan.runs = std::move(spans);
return plan;
}
// The pads an earlier splice read (`pads`: their positions in `live`, sorted) that the prompt `cur` still shares with
// the session: the ones before the last token of `cur` that `live` holds at the same place once the pads are
// skipped. `out` is `cur` with those pads put back, which makes it start with `live` as far as they agree.
template <class A, class B>
std::vector<int64_t> scr_carry(const std::vector<B>& cur, const std::vector<A>& live, const std::vector<int64_t>& pads,
int64_t pad_token, std::vector<B>& out) {
size_t p = 0, i = 0, k = 0, kept = 0;
while (p < live.size()) {
if (k < pads.size() && pads[k] == (int64_t) p) {
if ((int64_t) live[p] != pad_token) break;
++k;
++p;
continue;
}
if (i + 1 >= cur.size() || (int64_t) live[p] != (int64_t) cur[i]) break;
++i;
++p;
kept = k;
}
std::vector<int64_t> carried(pads.begin(), pads.begin() + (long) kept);
out.clear();
out.reserve(cur.size() + carried.size());
size_t c = 0;
for (const int64_t at : carried) {
while ((int64_t) out.size() < at) out.push_back(cur[c++]);
out.push_back((B) pad_token);
}
out.insert(out.end(), cur.begin() + (long) c, cur.end());
return carried;
}
// The positions of the pads `plan` inserts, in the padded prompt.
inline std::vector<int64_t> scr_pad_positions(const ScrPlan& plan) {
std::vector<int64_t> at;
for (const ScrRun& r : plan.runs)
for (int64_t j = r.pad; j > 0; --j) at.push_back(r.at - j);
return at;
}
// The prompt the engine reads under `plan`: `cur` with each run's pads inserted right before it.
template <class B>
std::vector<B> scr_padded(const std::vector<B>& cur, const ScrPlan& plan, B pad_token) {
std::vector<B> out;
out.reserve(cur.size() + (size_t) plan.pads());
int64_t from = 0;
for (const ScrRun& r : plan.runs) {
out.insert(out.end(), cur.begin() + from, cur.begin() + r.new_start);
out.insert(out.end(), (size_t) r.pad, pad_token);
from = r.new_start;
}
out.insert(out.end(), cur.begin() + from, cur.end());
return out;
}
} // namespace strata::core