use super::args::{self, Args};
use super::notify::{self, MISS, class};
use super::table::Spec;
use yo_common::{Code, Error};
use yo_kv::{Db, Kind};
const READS: &[&str] = &[
"string",
"bitmap",
"hyperloglog",
"list",
"hash",
"set",
"zset",
"geo",
"array",
"stream",
"graph",
"keyspace",
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum At {
Gone,
Fine,
Wrong,
}
impl At {
fn goes_on(self) -> bool {
self != At::Wrong
}
}
struct Probe<'a> {
db: &'a Db,
on: usize,
accepts: Option<&'static [Kind]>,
}
impl Probe<'_> {
fn peek(&self, key: &[u8]) -> Option<Kind> {
self.db.hold(key).kind_of(key)
}
fn read(&self, key: &[u8]) -> At {
match self.peek(key) {
None => {
notify::fire(self.on, class::KEY_MISS, MISS, key);
At::Gone
}
Some(kind) => match self.accepts {
Some(kinds) if !kinds.contains(&kind) => At::Wrong,
_ => At::Fine,
},
}
}
}
pub(super) fn report(db: &Db, on: usize, spec: &Spec, args: Args<'_>) {
if !notify::wanted(class::KEY_MISS) {
return;
}
let probe = Probe {
db,
on,
accepts: accepts(spec.name),
};
reads(spec, args, &probe);
}
fn accepts(name: &str) -> Option<&'static [Kind]> {
match name {
"sinter" | "sinterstore" | "sintercard" | "sunion" | "sunionstore" | "sunioncard"
| "sdiff" | "sdiffstore" | "sdiffcard" => Some(&[Kind::Set]),
"zunion" | "zunionstore" | "zinter" | "zinterstore" | "zintercard" | "zdiff"
| "zdiffstore" => Some(&[Kind::Set, Kind::Zset]),
"bitop" | "pfcount" | "pfmerge" => Some(&[Kind::String]),
"xread" | "xreadgroup" => Some(&[Kind::Stream]),
"tdigest.merge" | "cms.merge" => Some(&[Kind::Foreign]),
_ => None,
}
}
pub(super) fn undo(spec: &Spec, e: &Error) {
if spec.flags.contains(&"module") {
return;
}
if matches!(e.code(), Code::Invalid | Code::Unsupported) {
notify::unsay_misses();
}
}
fn reads(spec: &Spec, args: Args<'_>, probe: &Probe<'_>) -> bool {
match spec.name {
"getdel" | "getex" | "getset" | "copy" | "delex" | "sort" | "xack" | "xackdel"
| "xnack" | "xclaim" | "xautoclaim" | "georadius" | "georadiusbymember" => {
one(args, 1, probe)
}
"set" => {
if (3..args.len()).any(|i| args::is(args.get(i), b"get")) {
one(args, 1, probe)
} else {
true
}
}
"bitfield" => {
let writes = (2..args.len())
.any(|i| args::is(args.get(i), b"set") || args::is(args.get(i), b"incrby"));
if writes { true } else { one(args, 1, probe) }
}
"bitop" | "sinterstore" | "sunionstore" | "sdiffstore" => span(spec, args, 1, probe),
"zunionstore" | "zinterstore" | "zdiffstore" => counted(args, 2, probe),
"zrangestore" | "geosearchstore" => one(args, 2, probe),
"pfmerge" => span(spec, args, 0, probe),
"sintercard" | "sunioncard" | "sdiffcard" | "zdiff" | "zunion" | "zinter"
| "zintercard" => counted(args, 1, probe),
"xread" => streams(args, probe) && streams(args, probe),
"xreadgroup" => streams(args, probe),
"xinfo" => one(args, 2, probe),
"migrate" => migrated(args, probe),
"ts.info" | "cf.compact" => true,
"json.debug" => debugged(args, probe),
"tdigest.merge" | "cms.merge" => merged(spec.name, args, probe),
"FT.SUGGET" | "FT.SUGLEN" => one(args, 1, probe),
_ if spec.flags.contains(&"readonly")
&& (READS.contains(&spec.group)
|| (spec.flags.contains(&"module") && spec.group != "search")) =>
{
span(spec, args, 0, probe)
}
_ => true,
}
}
fn debugged(args: Args<'_>, probe: &Probe<'_>) -> bool {
if !args::is(args.get(1), b"memory") {
return true;
}
one(args, 2, probe)
}
fn merged(name: &str, args: Args<'_>, probe: &Probe<'_>) -> bool {
let Some(dest) = args.opt(1) else {
return true;
};
if name == "cms.merge" {
if probe.peek(dest).is_none() {
return true;
}
} else if !probe.read(dest).goes_on() {
return true;
}
let Ok(count) = args.int(2) else {
return true;
};
let Ok(count) = usize::try_from(count) else {
return true;
};
if 3 + count > args.len() {
return true;
}
for i in 0..count {
if probe.read(args.get(3 + i)) != At::Fine {
return false;
}
}
true
}
fn one(args: Args<'_>, at: usize, probe: &Probe<'_>) -> bool {
match args.opt(at) {
Some(key) => probe.read(key).goes_on(),
None => true,
}
}
fn span(spec: &Spec, args: Args<'_>, skip: i32, probe: &Probe<'_>) -> bool {
if spec.first_key <= 0 {
return true;
}
let step = spec.step.max(1);
let argc = i32::try_from(args.len()).unwrap_or(i32::MAX);
let last = if spec.last_key < 0 {
argc + spec.last_key
} else {
spec.last_key
};
let mut at = spec.first_key + skip * step;
while at <= last && at < argc {
if !probe.read(args.get(at as usize)).goes_on() {
return false;
}
at += step;
}
true
}
fn counted(args: Args<'_>, at: usize, probe: &Probe<'_>) -> bool {
let Ok(count) = args.int(at) else {
return true;
};
let Ok(count) = usize::try_from(count) else {
return true;
};
if at + 1 + count > args.len() {
return true;
}
for i in 0..count {
if !probe.read(args.get(at + 1 + i)).goes_on() {
return false;
}
}
true
}
fn migrated(args: Args<'_>, probe: &Probe<'_>) -> bool {
match args.opt(3) {
Some(key) if !key.is_empty() => probe.read(key).goes_on(),
_ => {
let Some(at) = (6..args.len()).find(|&i| args::is(args.get(i), b"keys")) else {
return true;
};
for i in at + 1..args.len() {
if !probe.read(args.get(i)).goes_on() {
return false;
}
}
true
}
}
}
fn streams(args: Args<'_>, probe: &Probe<'_>) -> bool {
let Some(at) = (1..args.len()).find(|&i| args::is(args.get(i), b"streams")) else {
return true;
};
let rest = args.len() - at - 1;
if rest == 0 || !rest.is_multiple_of(2) {
return true;
}
for i in 0..rest / 2 {
if !probe.read(args.get(at + 1 + i)).goes_on() {
return false;
}
}
true
}