use std::collections::HashMap;
use std::collections::hash_map::Entry;
use yo_common::num::parse_f64;
use yo_common::parse_i64;
use yo_search::Index;
use yo_search::expr::{Expr, Value, seventeen};
use yo_search::field::Kind;
use yo_search::reduce;
use yo_search::token;
use super::cursor::{Kept, Made};
use super::{
AS_SHORT, Args, Asked, BAD_ARGS, COUNT_ONLY, DUPLICATE_END, DUPLICATE_PROP, GROUP_COUNT,
GROUP_SHORT, MAX_COUNT, MISSING_ARGS, MOST_SAMPLE, NO_AT, NO_AT_END, NO_AT_MID, NO_PROPERTY,
NO_REDUCER, NOT_A_NUMBER, NOT_LOADED, NOT_MAIN, NOT_THERE, OUT_OF_RANGE, PERCENTAGE, QUOTE_END,
REDUCE_BARE, RESOLUTION, RESOLUTION_ARG, Row, SAMPLE_BIG, SAMPLE_SIZE, SORT_BOUNDS, SORT_COUNT,
SORT_PROP, SORT_PROP_END, SORT_SHORT, SORT_TWICE, SORT_WAY, SORT_WAY_END, STEP_SHORT, Shows,
UNKNOWN, cursor, line, twelve,
};
use crate::dispatch::Server;
use crate::dispatch::args;
use crate::dispatch::indexing;
use crate::reply::Out;
const SORTED: usize = 10;
#[derive(Default)]
pub(super) struct Pipe<'a> {
pub(super) load: Vec<(&'a [u8], &'a [u8])>,
pub(super) all: bool,
pub(super) loader: bool,
pub(super) addscores: bool,
pub(super) sortkeys: bool,
pub(super) pending: bool,
pub(super) stepped: bool,
pub(super) base: Vec<(Box<[u8]>, Reads)>,
pub(super) steps: Vec<Step>,
pub(super) arrange: Option<usize>,
pub(super) stage: Option<Vec<Box<[u8]>>>,
}
pub(super) enum Reads {
Field(Box<[u8]>, Shape),
Score,
Made,
}
#[derive(Clone, Copy, PartialEq, Eq)]
pub(super) enum Shape {
Words,
Number,
Folded,
}
pub(super) enum Step {
Group(Group),
Apply {
at: usize,
name: Box<[u8]>,
expr: Expr,
},
Filter(Expr),
Sort(Sort),
}
#[derive(Default)]
pub(super) struct Sort {
keys: Vec<(usize, bool)>,
offset: usize,
count: Option<usize>,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Look {
Any,
Named,
Sorted,
}
pub(super) struct Group {
by: Vec<(usize, Box<[u8]>)>,
folds: Vec<Reducer>,
}
pub(super) struct Reducer {
kind: reduce::Kind,
of: Option<usize>,
by: Option<usize>,
name: Box<[u8]>,
}
pub(super) fn group<'a>(
args: Args<'a>,
at: usize,
asked: &mut Asked<'a>,
index: &Index,
) -> core::result::Result<usize, Vec<u8>> {
asked.pipe.stepped = true;
let mut at = at + 1;
let Some(count) = args.opt(at) else {
return Err(GROUP_SHORT.as_bytes().to_vec());
};
at += 1;
let Some(count) = parse_i64(count) else {
return Err(GROUP_COUNT.as_bytes().to_vec());
};
let Ok(count) = usize::try_from(count) else {
return Err(GROUP_SHORT.as_bytes().to_vec());
};
let mut by: Vec<(usize, Box<[u8]>)> = Vec::with_capacity(count);
let mut names: Vec<Box<[u8]>> = Vec::with_capacity(count);
for _ in 0..count {
let Some(word) = args.opt(at) else {
return Err(GROUP_SHORT.as_bytes().to_vec());
};
at += 1;
let Some(name) = word.strip_prefix(b"@") else {
let mut out = line(NO_AT, word, NO_AT_MID);
out.extend_from_slice(word);
out.extend_from_slice(NO_AT_END.as_bytes());
return Err(out);
};
if names.iter().any(|held| **held == *name) {
return Err(line(DUPLICATE_PROP, name, DUPLICATE_END));
}
let Some(from) = locate(&mut asked.pipe, index, name, Look::Any) else {
return Err(line(NO_PROPERTY, name, QUOTE_END));
};
by.push((from, name.into()));
names.push(name.into());
}
let mut folds: Vec<Reducer> = Vec::new();
while args.opt(at).is_some_and(|word| args::is(word, b"REDUCE")) {
let fold = reducer(args, &mut at, asked, index)?;
if names.contains(&fold.name) {
return Err(line(DUPLICATE_PROP, &fold.name, DUPLICATE_END));
}
names.push(fold.name.clone());
folds.push(fold);
}
asked.pipe.steps.push(Step::Group(Group { by, folds }));
asked.pipe.arrange = None;
asked.pipe.stage = Some(names);
Ok(at)
}
fn reducer<'a>(
args: Args<'a>,
at: &mut usize,
asked: &mut Asked<'a>,
index: &Index,
) -> core::result::Result<Reducer, Vec<u8>> {
*at += 1;
let Some(func) = args.opt(*at) else {
return Err(REDUCE_BARE.as_bytes().to_vec());
};
*at += 1;
let Some(count) = args.opt(*at) else {
return Err(line(BAD_ARGS, func, NOT_THERE));
};
*at += 1;
let Some(count) = parse_i64(count) else {
return Err(line(BAD_ARGS, func, NOT_A_NUMBER));
};
let Ok(count) = usize::try_from(count) else {
return Err(line(BAD_ARGS, func, OUT_OF_RANGE));
};
let mut words: Vec<&[u8]> = Vec::with_capacity(count);
for _ in 0..count {
let Some(word) = args.opt(*at) else {
return Err(line(BAD_ARGS, func, NOT_THERE));
};
*at += 1;
words.push(word);
}
let (kind, of, by) = fold(func, &words, asked, index)?;
let name = match args.opt(*at).is_some_and(|word| args::is(word, b"AS")) {
false => generated(func, &words),
true => {
*at += 1;
let Some(name) = args.opt(*at) else {
return Err(line(BAD_ARGS, b"AS", NOT_THERE));
};
*at += 1;
name.into()
}
};
Ok(Reducer { kind, of, by, name })
}
#[expect(clippy::too_many_lines, reason = "twelve reducers and their arguments")]
fn fold(
func: &[u8],
words: &[&[u8]],
asked: &mut Asked<'_>,
index: &Index,
) -> core::result::Result<(reduce::Kind, Option<usize>, Option<usize>), Vec<u8>> {
let first = |asked: &mut Asked<'_>| match words.first() {
None => Err(line(MISSING_ARGS, func, "")),
Some(word) => {
let name = word.strip_prefix(b"@").unwrap_or(word);
match locate(&mut asked.pipe, index, name, Look::Any) {
Some(from) => Ok(from),
None => Err(line(NOT_LOADED, name, QUOTE_END)),
}
}
};
if args::is(func, b"COUNT") {
if !words.is_empty() {
return Err(COUNT_ONLY.as_bytes().to_vec());
}
return Ok((reduce::Kind::Count, None, None));
}
for (name, kind) in [
(b"SUM".as_slice(), reduce::Kind::Sum),
(b"MIN", reduce::Kind::Min),
(b"MAX", reduce::Kind::Max),
(b"AVG", reduce::Kind::Avg),
(b"STDDEV", reduce::Kind::Stddev),
(b"TOLIST", reduce::Kind::ToList),
(b"COUNT_DISTINCT", reduce::Kind::Distinct),
(b"COUNT_DISTINCTISH", reduce::Kind::Distinctish),
] {
if args::is(func, name) {
return Ok((kind, Some(first(asked)?), None));
}
}
if args::is(func, b"QUANTILE") {
let of = first(asked)?;
let Some(want) = words.get(1) else {
return Err(line(MISSING_ARGS, func, ""));
};
let Some(want) = parse_f64(want) else {
return Err(line(BAD_ARGS, func, NOT_A_NUMBER));
};
if !(0.0..=1.0).contains(&want) {
return Err(PERCENTAGE.as_bytes().to_vec());
}
if let Some(fine) = words.get(2) {
let Some(fine) = parse_i64(fine) else {
return Err(line(RESOLUTION_ARG, b"", NOT_A_NUMBER));
};
if !(1..=MOST_SAMPLE).contains(&fine) {
return Err(RESOLUTION.as_bytes().to_vec());
}
}
if let Some(word) = words.get(3) {
return Err(inside(word, 3, func));
}
return Ok((reduce::Kind::Quantile(want), Some(of), None));
}
if args::is(func, b"RANDOM_SAMPLE") {
let of = first(asked)?;
let Some(size) = words.get(1) else {
return Err(line(SAMPLE_SIZE, b"", NOT_THERE));
};
let Some(size) = parse_i64(size) else {
return Err(line(SAMPLE_SIZE, b"", NOT_A_NUMBER));
};
if size < 0 {
return Err(line(SAMPLE_SIZE, b"", OUT_OF_RANGE));
}
if size > MOST_SAMPLE {
return Err(SAMPLE_BIG.as_bytes().to_vec());
}
let size = usize::try_from(size).unwrap_or(0);
return Ok((reduce::Kind::Sample(size), Some(of), None));
}
if args::is(func, b"FIRST_VALUE") {
let of = first(asked)?;
let mut order = reduce::Order {
by: false,
desc: false,
numeric: false,
};
let mut key = None;
if let Some(word) = words.get(1) {
if !args::is(word, b"BY") {
return Err(inside(word, 1, func));
}
let Some(word) = words.get(2) else {
return Err(line(MISSING_ARGS, func, ""));
};
let name = word.strip_prefix(b"@").unwrap_or(word);
let Some(from) = locate(&mut asked.pipe, index, name, Look::Any) else {
return Err(line(NOT_LOADED, name, QUOTE_END));
};
order.by = true;
order.numeric = index
.field(name)
.is_some_and(|f| matches!(f.kind, Kind::Numeric));
key = Some(from);
if let Some(word) = words.get(3) {
match () {
() if args::is(word, b"ASC") => {}
() if args::is(word, b"DESC") => order.desc = true,
() => return Err(inside(word, 3, func)),
}
}
if let Some(word) = words.get(4) {
return Err(inside(word, 4, func));
}
}
return Ok((reduce::Kind::First(order), Some(of), key));
}
Err(line(NO_REDUCER, func, ""))
}
fn inside(word: &[u8], position: usize, func: &[u8]) -> Vec<u8> {
let mut out = UNKNOWN.as_bytes().to_vec();
out.extend_from_slice(word);
out.extend_from_slice(NOT_MAIN.as_bytes());
out.extend_from_slice(position.to_string().as_bytes());
out.extend_from_slice(b" for ");
out.extend_from_slice(func);
out
}
fn generated(func: &[u8], words: &[&[u8]]) -> Box<[u8]> {
let mut out = b"__generated_alias".to_vec();
out.extend_from_slice(&func.to_ascii_lowercase());
for (at, word) in words.iter().enumerate() {
if at > 0 {
out.push(b',');
}
let word = word.strip_prefix(b"@").unwrap_or(word);
out.extend_from_slice(&word.to_ascii_lowercase());
}
out.into()
}
fn locate(pipe: &mut Pipe<'_>, index: &Index, name: &[u8], look: Look) -> Option<usize> {
if let Some(stage) = &pipe.stage {
return stage.iter().position(|held| **held == *name);
}
if let Some(at) = pipe.base.iter().position(|(held, _)| **held == *name) {
return Some(at);
}
if name == b"__score" {
if !pipe.addscores {
return None;
}
pipe.base.push((name.into(), Reads::Score));
return Some(pipe.base.len() - 1);
}
let (from, shape) = match index.field(name) {
Some(field) => {
if look == Look::Sorted && !field.sortable && !pipe.all {
return None;
}
let shape = match () {
() if field.kind == Kind::Numeric => Shape::Number,
() if field.sortable && !field.is_unf() => Shape::Folded,
() => Shape::Words,
};
(field.identifier.clone(), shape)
}
None if pipe.all && look != Look::Named => (name.into(), Shape::Words),
None => return None,
};
pipe.base.push((name.into(), Reads::Field(from, shape)));
Some(pipe.base.len() - 1)
}
fn makes(pipe: &mut Pipe<'_>, name: &[u8]) -> usize {
if let Some(stage) = &mut pipe.stage {
if let Some(at) = stage.iter().position(|held| **held == *name) {
return at;
}
stage.push(name.into());
return stage.len() - 1;
}
if let Some(at) = pipe.base.iter().position(|(held, _)| **held == *name) {
return at;
}
pipe.base.push((name.into(), Reads::Made));
pipe.base.len() - 1
}
fn built(src: &[u8], asked: &mut Asked<'_>, index: &Index) -> core::result::Result<Expr, Vec<u8>> {
let mut expr = Expr::parse(src)?;
expr.bind(&mut |name| locate(&mut asked.pipe, index, name, Look::Sorted))
.map_err(|missing| line(NOT_LOADED, &missing.0, QUOTE_END))?;
Ok(expr)
}
pub(super) fn apply<'a>(
args: Args<'a>,
at: usize,
asked: &mut Asked<'a>,
index: &Index,
) -> core::result::Result<usize, Vec<u8>> {
asked.pipe.stepped = true;
let Some(src) = args.opt(at + 1) else {
return Err(STEP_SHORT.as_bytes().to_vec());
};
let mut at = at + 2;
let name: Box<[u8]> = match args.opt(at).is_some_and(|word| args::is(word, b"AS")) {
false => src.into(),
true => {
at += 1;
let Some(name) = args.opt(at) else {
return Err(AS_SHORT.as_bytes().to_vec());
};
at += 1;
name.into()
}
};
let expr = built(src, asked, index)?;
let slot = makes(&mut asked.pipe, &name);
asked.pipe.steps.push(Step::Apply {
at: slot,
name,
expr,
});
Ok(at)
}
pub(super) fn keeps<'a>(
args: Args<'a>,
at: usize,
asked: &mut Asked<'a>,
index: &Index,
) -> core::result::Result<usize, Vec<u8>> {
asked.pipe.stepped = true;
let Some(src) = args.opt(at + 1) else {
return Err(STEP_SHORT.as_bytes().to_vec());
};
let expr = built(src, asked, index)?;
asked.pipe.steps.push(Step::Filter(expr));
Ok(at + 2)
}
pub(super) fn sorts<'a>(
args: Args<'a>,
at: usize,
asked: &mut Asked<'a>,
index: &Index,
) -> core::result::Result<usize, Vec<u8>> {
let Some(count) = args.opt(at + 1) else {
return Err(SORT_SHORT.as_bytes().to_vec());
};
let Some(count) = parse_i64(count) else {
return Err(SORT_COUNT.as_bytes().to_vec());
};
let Ok(count) = usize::try_from(count) else {
return Err(SORT_BOUNDS.as_bytes().to_vec());
};
let mut at = at + 2;
let mut keys: Vec<(&[u8], bool)> = Vec::new();
for _ in 0..count {
let Some(word) = args.opt(at) else {
return Err(SORT_SHORT.as_bytes().to_vec());
};
at += 1;
if let Some(name) = word.strip_prefix(b"@") {
keys.push((name, false));
continue;
}
let down = match () {
() if args::is(word, b"ASC") => false,
() if args::is(word, b"DESC") => true,
() => return Err(line(SORT_WAY, word, SORT_WAY_END)),
};
if let Some((_, way)) = keys.last_mut() {
*way = down;
}
}
let mut most = None;
if args.opt(at).is_some_and(|word| args::is(word, b"MAX")) {
let Some(number) = args.opt(at + 1).and_then(parse_i64).filter(|n| *n >= 0) else {
return Err(MAX_COUNT.as_bytes().to_vec());
};
most = usize::try_from(number).ok().filter(|n| *n > 0);
at += 2;
}
let spot = arranged(&mut asked.pipe);
if matches!(&asked.pipe.steps[spot], Step::Sort(sort) if !sort.keys.is_empty()) {
return Err(SORT_TWICE.as_bytes().to_vec());
}
let mut found = Vec::with_capacity(keys.len());
for (name, down) in keys {
let Some(from) = locate(&mut asked.pipe, index, name, Look::Named) else {
let mut out = SORT_PROP.as_bytes().to_vec();
out.extend_from_slice(name);
out.extend_from_slice(SORT_PROP_END.as_bytes());
return Err(out);
};
found.push((from, down));
}
asked.pipe.loader |= !found.is_empty();
if let Some(most) = most {
asked.rows.count = most;
}
if let Step::Sort(sort) = &mut asked.pipe.steps[spot] {
sort.keys = found;
if let Some(most) = most {
sort.count = Some(most);
}
}
Ok(at)
}
fn arranged(pipe: &mut Pipe<'_>) -> usize {
if let Some(at) = pipe.arrange {
return at;
}
pipe.steps.push(Step::Sort(Sort::default()));
pipe.arrange = Some(pipe.steps.len() - 1);
pipe.steps.len() - 1
}
pub(super) fn most(asked: &mut Asked<'_>) {
let Some(spot) = asked.pipe.arrange else {
return;
};
if let Step::Sort(sort) = &mut asked.pipe.steps[spot]
&& !sort.keys.is_empty()
&& sort.count.is_none()
{
sort.count = Some(SORTED);
}
}
pub(super) fn windows(asked: &mut Asked<'_>, offset: usize, count: usize) {
let spot = arranged(&mut asked.pipe);
if let Step::Sort(sort) = &mut asked.pipe.steps[spot] {
sort.offset = offset;
sort.count = Some(count);
}
}
struct Held {
values: Vec<Value>,
dropped: usize,
from: Option<usize>,
}
pub(super) fn piped(
server: &Server,
db: usize,
total: usize,
rows: &[Row],
asked: &Asked<'_>,
index: &[u8],
out: &mut Out,
) {
let pipe = &asked.pipe;
let mut names: Vec<Box<[u8]>> = pipe.base.iter().map(|(name, _)| name.clone()).collect();
let reads = pipe.all
|| pipe
.base
.iter()
.any(|(_, from)| matches!(from, Reads::Field(..)));
let mut table: Vec<Held> = Vec::with_capacity(rows.len());
let mut walk = vec![true; rows.len()];
let mut lost = 0;
for (at, row) in rows.iter().enumerate() {
let doc = match reads {
false => None,
true => match indexing::read(&server.dbs[db], &row.key) {
Some(doc) => Some(doc),
None => {
lost += 1;
walk[at] = false;
continue;
}
},
};
let pairs = doc.as_ref().map(indexing::Document::pairs);
let mut made: Vec<Value> = Vec::with_capacity(pipe.base.len());
for (_, from) in &pipe.base {
made.push(match from {
Reads::Made => Value::Missing,
Reads::Score => Value::Text(twelve(row.score).into_bytes().into()),
Reads::Field(id, shape) => {
match pairs.iter().flatten().find(|(held, _)| *held == &**id) {
Some((_, value)) => shaped(value, *shape),
None => Value::Missing,
}
}
});
}
if pipe.all {
for (field, value) in pairs.iter().flatten() {
if pipe.base.iter().any(|(held, _)| **held == **field) {
continue;
}
let spot = match names.iter().position(|held| **held == **field) {
Some(spot) => spot,
None => {
names.push((*field).into());
names.len() - 1
}
};
made.resize_with(made.len().max(spot + 1), || Value::Missing);
made[spot] = Value::Text((*value).into());
}
}
made.resize_with(names.len(), || Value::Missing);
table.push(Held {
values: made,
dropped: 0,
from: Some(at),
});
}
let mut start = total - lost;
let mut warning = None;
let mut grouped = false;
let mut ranked = false;
let mut gone = 0;
let mut sorted = None;
for step in &pipe.steps {
match step {
Step::Group(group) => {
table = folded(group, &table);
names = group
.by
.iter()
.map(|(_, name)| name.clone())
.chain(group.folds.iter().map(|fold| fold.name.clone()))
.collect();
start = table.len();
grouped = true;
sorted = None;
}
Step::Apply { at, name, expr } => {
match *at < names.len() {
true => names[*at] = name.clone(),
false => names.push(name.clone()),
}
let mut done = 0;
for held in &mut table {
match expr.eval(&held.values) {
Ok(value) => {
held.values.resize_with(names.len(), || Value::Missing);
held.values[*at] = value;
}
Err(bad) => {
warning = Some(bad);
break;
}
}
done += 1;
}
if warning.is_some() {
table.truncate(done);
}
}
Step::Filter(expr) => {
let mut over = 0;
let mut stop = false;
let seen = table.len();
table.retain_mut(|held| {
if stop {
return false;
}
match expr.eval(&held.values) {
Ok(value) if value.truth() => {
held.dropped += over;
true
}
Ok(_) => {
over += 1;
if let Some(at) = held.from {
walk[at] = false;
}
false
}
Err(bad) => {
warning = Some(bad);
stop = true;
false
}
}
});
gone += seen - table.len();
}
Step::Sort(sort) => {
if let Some((first, _)) = sort.keys.first() {
table.sort_by(|left, right| ordered(&sort.keys, left, right));
ranked = true;
sorted = Some(*first);
}
table.drain(..sort.offset.min(table.len()));
if let Some(count) = sort.count {
table.truncate(count);
}
}
}
if warning.is_some() {
break;
}
}
let shown: Vec<(Option<&Row>, &Vec<Value>)> = table
.iter()
.map(|held| (held.from.map(|at| &rows[at]), &held.values))
.collect();
if let (Some(bad), true) = (&warning, shown.is_empty()) {
out.error(bad);
return;
}
let settled = grouped || ranked;
let counted = counting(&table, start, gone, asked, settled, shown.len(), out);
if asked.cursor.is_some() {
drop(shown);
if settled {
walk = Vec::new();
}
let made = Made::Piped {
names,
sorted,
rows: table
.into_iter()
.map(|held| (held.from.map(|at| rows[at].clone()), held.values))
.collect(),
warning,
};
opened(server, index, made, walk, counted, asked, settled, out);
return;
}
writes(
counted,
&names,
&shown,
sorted,
asked.rolls(),
warning.as_deref(),
out,
);
}
#[allow(clippy::too_many_arguments)]
fn opened(
server: &Server,
index: &[u8],
made: Made,
walk: Vec<bool>,
counted: usize,
asked: &Asked<'_>,
settled: bool,
out: &mut Out,
) {
let Some(asks) = asked.cursor else {
return;
};
let want = window(asked);
let kept = Kept {
made,
walk,
shows: asked.rolls(),
total: counted,
whole: settled || asked.rows.count == 0 || super::buffered(asked),
loader: asked.pipe.loader,
offset: want.0,
window: want.1,
};
cursor::open(server, index, kept, asks, out);
}
fn window(asked: &Asked<'_>) -> (usize, usize) {
let mut found = None;
for step in &asked.pipe.steps {
if let Step::Sort(sort) = step {
found = Some((sort.offset, sort.count.unwrap_or(usize::MAX)));
}
}
found.unwrap_or((asked.rows.offset, asked.rows.count))
}
fn counting(
table: &[Held],
start: usize,
gone: usize,
asked: &Asked<'_>,
settled: bool,
shown: usize,
out: &Out,
) -> usize {
let want = &asked.rows;
let dropped = match table.first() {
Some(held) => held.dropped,
None => gone,
};
let count = start - dropped.min(start);
let deep = out.proto().is_resp3();
let whole = settled
|| want.count == 0
|| super::buffered(asked)
|| (asked.pipe.loader && want.offset == 0);
if whole {
return count;
}
let reached = match deep || asked.pipe.loader {
true => shown,
false => 1,
};
want.offset.saturating_add(reached).min(count)
}
fn ordered(keys: &[(usize, bool)], left: &Held, right: &Held) -> core::cmp::Ordering {
for (from, down) in keys {
let this = left.values.get(*from).unwrap_or(&Value::Missing);
let that = right.values.get(*from).unwrap_or(&Value::Missing);
let held = |value: &Value| !matches!(value, Value::Missing);
match (held(this), held(that)) {
(false, false) => continue,
(false, true) => return core::cmp::Ordering::Greater,
(true, false) => return core::cmp::Ordering::Less,
(true, true) => {}
}
let Some(way) = yo_search::expr::order(this, that) else {
continue;
};
if way != core::cmp::Ordering::Equal {
return match down {
true => way.reverse(),
false => way,
};
}
}
core::cmp::Ordering::Equal
}
fn shaped(value: &[u8], shape: Shape) -> Value {
match shape {
Shape::Words => Value::Text(value.into()),
Shape::Folded => Value::Text(token::fold(value).into()),
Shape::Number => match yo_common::num::parse_f64(value) {
Some(number) => Value::Number(number),
None => Value::Text(value.into()),
},
}
}
fn folded(group: &Group, table: &[Held]) -> Vec<Held> {
let mut order: Vec<Vec<Value>> = Vec::new();
let mut folds: Vec<Vec<reduce::Fold>> = Vec::new();
let mut seen: HashMap<Vec<u8>, usize> = HashMap::new();
let mut counts = vec![false; group.folds.len()];
for row in table {
let key: Vec<Value> = group
.by
.iter()
.map(|(from, _)| match &row.values[*from] {
Value::Missing => Value::Nil,
held => held.clone(),
})
.collect();
let at = match seen.entry(tagged(&key)) {
Entry::Occupied(held) => *held.get(),
Entry::Vacant(spot) => {
spot.insert(order.len());
order.push(key);
folds.push(
group
.folds
.iter()
.map(|fold| reduce::Fold::new(fold.kind.clone()))
.collect(),
);
order.len() - 1
}
};
for ((fold, what), numeric) in folds[at]
.iter_mut()
.zip(&group.folds)
.zip(counts.iter_mut())
{
let spelled;
let value = match what.of.map(|at| &row.values[at]) {
Some(Value::Text(text)) => Some(&**text),
Some(Value::Number(number)) => {
*numeric = true;
spelled = twelve(*number).into_bytes();
Some(&spelled[..])
}
_ => None,
};
let ordered;
let by = match what.by.map(|at| &row.values[at]) {
Some(Value::Text(text)) => Some(&**text),
Some(Value::Number(number)) => {
ordered = twelve(*number).into_bytes();
Some(&ordered[..])
}
_ => None,
};
fold.add(value, by);
}
}
order
.into_iter()
.zip(folds)
.map(|(key, folds)| Held {
values: key
.into_iter()
.chain(
folds
.into_iter()
.zip(&counts)
.map(|(fold, numeric)| answered(fold.done(), *numeric)),
)
.collect(),
dropped: 0,
from: None,
})
.collect()
}
fn tagged(values: &[Value]) -> Vec<u8> {
let mut out = Vec::new();
for value in values {
match value {
Value::Missing | Value::Nil => out.push(0),
Value::Text(text) => {
out.push(1);
out.extend_from_slice(&text.len().to_le_bytes());
out.extend_from_slice(text);
}
Value::Number(number) => {
out.push(3);
out.extend_from_slice(&number.to_le_bytes());
}
Value::List(list) => {
out.push(2);
out.extend_from_slice(&list.len().to_le_bytes());
for item in list {
out.extend_from_slice(&tagged(core::slice::from_ref(item)));
}
}
}
}
out
}
fn answered(answer: reduce::Answer, numeric: bool) -> Value {
let held = |text: Box<[u8]>| match numeric.then(|| yo_common::num::parse_f64(&text)).flatten() {
Some(number) => Value::Number(number),
None => Value::Text(text),
};
match answer {
reduce::Answer::Number(number) => Value::Number(number),
reduce::Answer::Text(text) => held(text),
reduce::Answer::List(list) => Value::List(list.into_iter().map(held).collect()),
reduce::Answer::Nil => Value::Nil,
}
}
pub(super) fn writes(
count: usize,
names: &[Box<[u8]>],
shown: &[(Option<&Row>, &Vec<Value>)],
sorted: Option<usize>,
shows: Shows,
warning: Option<&[u8]>,
out: &mut Out,
) {
let extras =
usize::from(shows.scores) + usize::from(shows.payloads) + usize::from(shows.sortkeys);
let scored = shows.addscores && !names.iter().any(|name| **name == *b"__score");
let scoring = |row: Option<&Row>| scored.then(|| row.map(|row| row.score)).flatten();
if out.proto().is_resp3() {
out.map(5);
out.simple(b"attributes");
out.array(0);
out.simple(b"format");
out.simple(b"STRING");
out.simple(b"results");
out.array(shown.len());
for (row, values) in shown {
out.map(1 + extras + usize::from(shows.fields));
if shows.scores {
out.simple(b"score");
out.double(row.map_or(0.0, |row| row.score));
}
if shows.payloads {
out.simple(b"payload");
match row.and_then(|row| row.payload.as_ref()) {
Some(payload) => out.bulk(payload),
None => out.nil(),
}
}
if shows.sortkeys {
out.simple(b"sortkey");
match keyed(values, sorted) {
Some(key) => out.bulk(&key),
None => out.nil(),
}
}
if shows.fields {
out.simple(b"extra_attributes");
mapped(names, values, scoring(*row), out);
}
out.simple(b"values");
out.array(0);
}
out.simple(b"total_results");
out.int(count as i64);
out.simple(b"warning");
match warning {
Some(warning) => {
out.array(1);
out.bulk(warning);
}
None => out.array(0),
}
return;
}
out.array(1 + shown.len() * (extras + usize::from(shows.fields)));
out.int(count as i64);
for (row, values) in shown {
if shows.scores {
out.double(row.map_or(0.0, |row| row.score));
}
if shows.payloads {
match row.and_then(|row| row.payload.as_ref()) {
Some(payload) => out.bulk(payload),
None => out.nil(),
}
}
if shows.sortkeys {
match keyed(values, sorted) {
Some(key) => out.bulk(&key),
None => out.nil(),
}
}
if shows.fields {
mapped(names, values, scoring(*row), out);
}
}
}
fn keyed(values: &[Value], sorted: Option<usize>) -> Option<Vec<u8>> {
let mut out = Vec::new();
match values.get(sorted?)? {
Value::Number(number) => {
out.push(b'#');
out.extend_from_slice(seventeen(*number).as_bytes());
}
Value::Text(text) => {
out.push(b'$');
out.extend_from_slice(text);
}
_ => return None,
}
Some(out)
}
fn mapped(names: &[Box<[u8]>], row: &[Value], score: Option<f64>, out: &mut Out) {
let held: Vec<(&Box<[u8]>, &Value)> = names
.iter()
.zip(row)
.filter(|(_, value)| !matches!(value, Value::Missing))
.collect();
out.map(held.len() + usize::from(score.is_some()));
if let Some(score) = score {
out.bulk(b"__score");
out.bulk(twelve(score).as_bytes());
}
for (name, value) in held {
out.bulk(name);
written(value, out);
}
}
fn written(value: &Value, out: &mut Out) {
match value {
Value::Text(text) => out.bulk(text),
Value::Number(number) => out.bulk(twelve(*number).as_bytes()),
Value::Missing | Value::Nil => out.nil(),
Value::List(list) => {
out.array(list.len());
for item in list {
written(item, out);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_reducer_with_no_name_gets_one_built_out_of_what_it_was_given() {
assert_eq!(&*generated(b"COUNT", &[]), b"__generated_aliascount");
assert_eq!(&*generated(b"Sum", &[b"@n"]), b"__generated_aliassumn");
assert_eq!(
&*generated(b"FIRST_VALUE", &[b"@n", b"BY", b"@n", b"Desc"]),
b"__generated_aliasfirst_valuen,by,n,desc".as_slice()
);
}
#[test]
fn two_groups_that_run_together_without_their_lengths_do_not_run_together_with_them() {
let one = [
Value::Text(b"ab".to_vec().into()),
Value::Text(b"c".to_vec().into()),
];
let two = [
Value::Text(b"a".to_vec().into()),
Value::Text(b"bc".to_vec().into()),
];
assert_ne!(tagged(&one), tagged(&two));
assert_ne!(
tagged(&[Value::Nil]),
tagged(&[Value::Text(Box::default())])
);
}
fn row(values: &[Value]) -> Held {
Held {
values: values.to_vec(),
dropped: 0,
from: None,
}
}
#[test]
fn a_row_without_the_property_sorts_last_whichever_way_the_key_runs() {
let here = row(&[Value::Number(1.0)]);
let gone = row(&[Value::Missing]);
for down in [false, true] {
assert_eq!(
ordered(&[(0, down)], &here, &gone),
core::cmp::Ordering::Less
);
assert_eq!(
ordered(&[(0, down)], &gone, &here),
core::cmp::Ordering::Greater
);
}
assert_eq!(
ordered(&[(0, false)], &gone, &gone),
core::cmp::Ordering::Equal
);
}
#[test]
fn a_null_is_a_value_of_its_own_and_turns_round_with_the_key() {
let null = row(&[Value::Nil]);
let held = row(&[Value::Number(-9.0)]);
assert_eq!(
ordered(&[(0, false)], &null, &held),
core::cmp::Ordering::Less
);
assert_eq!(
ordered(&[(0, true)], &null, &held),
core::cmp::Ordering::Greater
);
}
#[test]
fn the_keys_are_read_in_turn_until_one_of_them_settles_it() {
let first = row(&[Value::Text(b"aa".as_slice().into()), Value::Number(2.0)]);
let second = row(&[Value::Text(b"aa".as_slice().into()), Value::Number(1.0)]);
let keys = [(0, false), (1, false)];
assert_eq!(
ordered(&keys, &first, &second),
core::cmp::Ordering::Greater
);
let keys = [(0, false), (1, true)];
assert_eq!(ordered(&keys, &first, &second), core::cmp::Ordering::Less);
assert_eq!(ordered(&[], &first, &second), core::cmp::Ordering::Equal);
}
#[test]
fn two_values_are_ordered_the_way_a_comparison_orders_them() {
let number = row(&[Value::Number(10.0)]);
let words = row(&[Value::Text(b"9".as_slice().into())]);
assert_eq!(
ordered(&[(0, false)], &number, &words),
core::cmp::Ordering::Greater
);
let list = row(&[Value::List(vec![Value::Number(1.0)])]);
assert_eq!(
ordered(&[(0, false)], &list, &number),
core::cmp::Ordering::Less
);
}
#[test]
fn a_sort_key_says_which_of_the_two_kinds_of_value_it_holds() {
let row = [
Value::Number(-4.0),
Value::Text(b"blue".as_slice().into()),
Value::Missing,
Value::Nil,
];
assert_eq!(keyed(&row, Some(0)).as_deref(), Some(b"#-4".as_slice()));
assert_eq!(keyed(&row, Some(1)).as_deref(), Some(b"$blue".as_slice()));
assert_eq!(keyed(&row, Some(2)), None);
assert_eq!(keyed(&row, Some(3)), None);
assert_eq!(keyed(&row, Some(9)), None);
assert_eq!(keyed(&row, None), None);
}
}