Skip to content

Commit

Permalink
Merge pull request 01mf02#239 from 01mf02/general-fold
Browse files Browse the repository at this point in the history 
Unify folding operators.
  • Loading branch information
@01mf02
01mf02 authored Nov 22, 2024
2 parents 743b88c + 7914629 commit 95b6a30
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 60 deletions.
31 changes: 17 additions & 14 deletions jaq-core/src/compile.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -117,10 +117,7 @@ pub(crate) enum Term<T = TermId> {
TryCatch(T, T),
/// If-then-else (`if f then g else h end`)
Ite(T, T, T),
/// `reduce`, `for`, and `foreach`
///
/// The first field indicates whether to yield intermediate results
/// (`false` for `reduce` and `true` for `foreach`).
/// `reduce` and `foreach`
///
/// Assuming that `xs` evaluates to `x0`, `x1`, ..., `xn`,
/// `reduce xs as $x (init; f)` evaluates to
Expand All @@ -132,20 +129,25 @@ pub(crate) enum Term<T = TermId> {
/// | xn as $x | f
/// ~~~
///
/// and `for xs as $x (init; f)` evaluates to
/// and `foreach xs as $x (init; f; project)` evaluates to
///
/// ~~~ text
/// init
/// | ., (x0 as $x | f
/// | ...
/// | ., (xn as $x | f)...)
/// init |
/// x0 as $x | f | project,
/// ...
/// xn as $x | f | project,
/// empty
/// ~~~
Reduce(T, Pattern<T>, T, T),
Foreach(T, Pattern<T>, T, T, Option<T>),

Fold(T, Pattern<T>, T, T, Fold<T>),
Path(T, crate::path::Path<T>),
}

#[derive(Clone, Debug)]
pub(crate) enum Fold<T> {
Reduce,
Foreach(Option<T>),
}

impl<T> Default for Term<T> {
fn default() -> Self {
Self::Id
Expand Down Expand Up @@ -660,6 +662,7 @@ impl<'s, F> Compiler<&'s str, F> {
Term::Label(self.lut.insert_term(tc))
}
Fold(name, xs, pat, args) => {
use self::Fold::{Foreach, Reduce};
let arity = args.len();
let mut args = args.into_iter();
let (init, update) = match (args.next(), args.next()) {
Expand All @@ -673,10 +676,10 @@ impl<'s, F> Compiler<&'s str, F> {
let update = self.with_vars(&vars, |c| c.iterm(update));

match (name, args.next(), args.next()) {
("reduce", None, None) => Term::Reduce(xs, pat, init, update),
("reduce", None, None) => Term::Fold(xs, pat, init, update, Reduce),
("foreach", proj, None) => {
let proj = proj.map(|p| self.with_vars(&vars, |c| c.iterm_tr(p, tr)));
Term::Foreach(xs, pat, init, update, proj)
Term::Fold(xs, pat, init, update, Foreach(proj))
}
_ => self.fail(name, Undefined::Filter(arity)),
}
Expand Down
46 changes: 18 additions & 28 deletions jaq-core/src/filter.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use crate::box_iter::{self, box_once, flat_map_then, flat_map_then_with, flat_map_with, map_with};
use crate::compile::{Lut, Pattern, Tailrec, Term as Ast};
use crate::fold::{fold, Fold};
use crate::compile::{Fold, Lut, Pattern, Tailrec, Term as Ast};
use crate::fold::fold;
use crate::val::{ValT, ValX, ValXs};
use crate::{exn, rc_lazy_list, Bind, Ctx, Error, Exn};
use alloc::boxed::Box;
Expand Down Expand Up @@ -102,7 +102,7 @@ where
F: Fn(T, V) -> ValXs<'a, V> + 'a,
{
let xs = rc_lazy_list::List::from_iter(xs);
Box::new(fold(false, xs, Fold::Input(init), f))
Box::new(fold(xs, init, f, |_| (), |_, _| None, Some))
}

fn label_skip<'a, V: 'a>(ys: ValXs<'a, V>, skip: usize) -> ValXs<'a, V> {
Expand Down Expand Up @@ -285,29 +285,20 @@ impl<F: FilterT<F>> FilterT<F> for Id {
Self::cartesian(l, r, lut, cv).map(|(x, y)| Ok(Self::V::from(op.run(&x?, &y?)))),
),

Ast::Reduce(xs, pat, init, update) => {
Ast::Fold(xs, pat, init, update, fold_type) => {
let xs = rc_lazy_list::List::from_iter(run_and_bind(xs, lut, cv.clone(), pat));
let init = init.run(lut, cv.clone());
let update = |ctx, v| update.run(lut, (ctx, v));
Box::new(fold(false, xs, Fold::Output(init), update))
}
Ast::Foreach(xs, pat, init, update, project) => {
let inputs = cv.0.inputs;
let xs = rc_lazy_list::List::from_iter(run_and_bind(xs, lut, cv.clone(), pat));
let init = init.run(lut, cv.clone());
let update = |ctx, v| update.run(lut, (ctx, v));
match project {
Some(proj) => flat_map_then_with(init, xs, move |init, xs| {
let init = Fold::Input((Ctx::new([], inputs), init));
let cvs = fold(true, xs, init, move |x, (_, acc)| {
map_with(update(x.clone(), acc), x, |y, x| Ok((x, y?)))
});
flat_map_then(cvs, |cv| proj.run(lut, cv))
}),
None => flat_map_then_with(init, xs, move |i, xs| {
Box::new(fold(true, xs, Fold::Input(i), update))
}),
}
let inner = |_, y: &Self::V| Some(y.clone());
let inner_proj = |ctx, y: &Self::V| Some((ctx, y.clone()));
flat_map_then_with(init, xs, move |i, xs| match fold_type {
Fold::Reduce => Box::new(fold(xs, i, update, |_| (), |_, _| None, Some)),
Fold::Foreach(None) => Box::new(fold(xs, i, update, |_| (), inner, |_| None)),
Fold::Foreach(Some(proj)) => flat_map_then(
fold(xs, i, update, |ctx| ctx.clone(), inner_proj, |_| None),
|(ctx, y)| proj.run(lut, (ctx, y)),
),
})
}

Ast::Var(v, skip) => match cv.0.vars.get(*v).unwrap() {
Expand Down Expand Up @@ -363,11 +354,10 @@ impl<F: FilterT<F>> FilterT<F> for Id {
Ast::Arr(_) | Ast::ObjEmpty | Ast::ObjSingle(..) => err,
Ast::Neg(_) | Ast::Logic(..) | Ast::Math(..) | Ast::Cmp(..) => err,
Ast::Update(..) | Ast::UpdateMath(..) | Ast::UpdateAlt(..) | Ast::Assign(..) => err,

// these are up for grabs to implement :)
Ast::TryCatch(..) | Ast::Label(_) | Ast::Reduce(..) | Ast::Foreach(..) => {
unimplemented!("updating with this operator is not supported yet")
}
// jq implements updates on `try ... catch` and `label`, but
// I do not see how to implement this in jaq
// folding, however, could be done, even if jq does not support it
Ast::TryCatch(..) | Ast::Label(_) | Ast::Fold(..) => err,

Ast::Id => f(cv.1),
Ast::Path(l, path) => {
Expand Down
42 changes: 24 additions & 18 deletions jaq-core/src/fold.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,45 +3,51 @@
use crate::box_iter::Results;
use alloc::vec::Vec;

pub(crate) enum Fold<'a, U, E> {
enum Fold<'a, X, Y, E> {
/// things to be processed
Input(U),
Input(Y),
/// things to be output, then to be input
Output(Results<'a, U, E>),
Output(X, Results<'a, Y, E>),
}

// if `inner` is true, output intermediate results
pub(crate) fn fold<'a, T: Clone + 'a, U: Clone + 'a, E: Clone + 'a>(
inner: bool,
pub(crate) fn fold<'a, T: 'a, TC: Clone + 'a, U: 'a, UC: 'a, E: 'a>(
xs: impl Iterator<Item = Result<T, E>> + Clone + 'a,
init: Fold<'a, U, E>,
init: U,
f: impl Fn(T, U) -> Results<'a, U, E> + 'a,
) -> impl Iterator<Item = Result<U, E>> + 'a {
let mut stack = Vec::from([(xs, init)]);
tc: impl Fn(&T) -> TC + 'a,
inner: impl Fn(TC, &U) -> Option<UC> + 'a,
outer: impl Fn(U) -> Option<UC> + 'a,
) -> impl Iterator<Item = Result<UC, E>> + 'a {
let mut stack = Vec::from([(xs, Fold::<TC, U, E>::Input(init))]);
core::iter::from_fn(move || loop {
let (mut xs, fold) = stack.pop()?;
match fold {
Fold::Output(mut ys) => match ys.next() {
Fold::Output(x, mut ys) => match ys.next() {
None => continue,
Some(y) => {
// do not grow the stack if the output is empty
if ys.size_hint() != (0, Some(0)) {
stack.push((xs.clone(), Fold::Output(ys)));
stack.push((xs.clone(), Fold::Output(x.clone(), ys)));
}
match y {
Ok(y) if inner => {
stack.push((xs, Fold::Input(y.clone())));
return Some(Ok(y));
Ok(y) => {
let inner = inner(x, &y);
stack.push((xs, Fold::Input(y)));
if let Some(inner) = inner {
return Some(Ok(inner));
}
}
Ok(y) => stack.push((xs, Fold::Input(y))),
Err(e) => return Some(Err(e)),
}
}
},
Fold::Input(y) => match xs.next() {
None if inner => continue,
None => return Some(Ok(y)),
Some(Ok(x)) => stack.push((xs, Fold::Output(f(x, y)))),
None => {
if let Some(outer) = outer(y) {
return Some(Ok(outer));
}
}
Some(Ok(x)) => stack.push((xs, Fold::Output(tc(&x), f(x, y)))),
Some(Err(e)) => return Some(Err(e)),
},
}
Expand Down

0 comments on commit 95b6a30

Please sign in to comment.