reimagine cascades rules?

memo_rule_apply.rs

#![feature(coroutines, coroutine_trait, coroutine_clone)]

use std::collections::VecDeque;
use std::{ops::Coroutine, sync::Arc};

#[derive(Clone, Copy, Debug)]
struct GroupId(usize);

impl std::fmt::Display for GroupId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "!{}", self.0)
    }
}

#[derive(Clone, Debug)]
enum Expr {
    Join(GroupId, GroupId),
    MultiJoin(Vec<GroupId>),
    Scan,
    Placeholder,
}

#[derive(Clone, Debug)]
enum Plan {
    Join(Arc<Plan>, Arc<Plan>),
    MultiJoin(Vec<Arc<Plan>>),
    Group(GroupId),
}

#[derive(Debug, Clone)]
enum Action {
    ExpandGroup(GroupId),
    InvokeSelf(Expr),
}

macro_rules! expand_group {
    ($x:ident) => {
        yield Action::ExpandGroup($x)
    };
}

macro_rules! recursive_call_self {
    ($x:ident) => {
        yield Action::InvokeSelf($x)
    };
}

fn join_assoc() -> impl Coroutine<Expr, Yield = Action, Return = Option<Plan>> + Clone {
    #[coroutine]
    |binding: Expr| {
        if let Expr::Join(left_group, c) = binding {
            let left = expand_group!(left_group);
            if let Expr::Join(a, b) = left {
                return Some(Plan::Join(
                    Plan::Group(a).into(),
                    Plan::Join(Plan::Group(b).into(), Plan::Group(c).into()).into(),
                ));
            }
        }
        None
    }
}

fn multi_join_loop() -> impl Coroutine<Expr, Yield = Action, Return = Option<Plan>> + Clone {
    #[coroutine]
    |mut binding: Expr| {
        let mut collected_groups = Vec::new();
        let mut last_right = None;
        while let Expr::Join(left, right) = binding {
            last_right = Some(right);
            binding = expand_group!(right);
            collected_groups.push(left);
        }
        if let Some(last_right) = last_right {
            collected_groups.push(last_right);
            if collected_groups.len() > 2 {
                let collected_groups = collected_groups
                    .into_iter()
                    .map(|group| Arc::new(Plan::Group(group)))
                    .collect();
                return Some(Plan::MultiJoin(collected_groups));
            }
        }
        None
    }
}

fn multi_join_recursive() -> impl Coroutine<Expr, Yield = Action, Return = Option<Plan>> + Clone {
    #[coroutine]
    |binding: Expr| {
        let Expr::Join(left_group, right_group) = binding else {
            return None;
        };
        let left = expand_group!(left_group);
        let right = expand_group!(right_group);
        let left_multi_join = if let Expr::Join(_, _) = left {
            let Expr::MultiJoin(lefts) = recursive_call_self!(left) else {
                return None;
            };
            lefts
        } else {
            vec![left_group]
        };
        let right_multi_join = if let Expr::Join(_, _) = right {
            let Expr::MultiJoin(rights) = recursive_call_self!(right) else {
                return None;
            };
            rights
        } else {
            vec![right_group]
        };
        let mut all = Vec::new();
        all.extend(left_multi_join);
        all.extend(right_multi_join);
        Some(Plan::MultiJoin(
            all.into_iter().map(|g| Arc::new(Plan::Group(g))).collect(),
        ))
    }
}

use std::{ops::CoroutineState, pin::Pin};

fn apply_group(
    memo: &[Vec<Expr>],
    group: GroupId,
    rule: impl Coroutine<Expr, Yield = Action, Return = Option<Plan>> + Clone + Unpin,
) -> Vec<Plan> {
    let mut result = Vec::new();
    for expr in memo[group.0].iter() {
        result.extend(apply(memo, expr.clone(), rule.clone()));
    }
    result
}

fn apply(
    memo: &[Vec<Expr>],
    expr: Expr,
    rule: impl Coroutine<Expr, Yield = Action, Return = Option<Plan>> + Clone + Unpin,
) -> Vec<Plan> {
    let mut states = VecDeque::new();
    let initial_state = rule.clone();
    states.push_back((expr, rule));
    let mut complete = Vec::<Plan>::new();
    while let Some((expr, state)) = states.pop_front() {
        let mut state = state.clone();
        let state_ = Pin::new(&mut state);
        match state_.resume(expr) {
            CoroutineState::Yielded(Action::ExpandGroup(group)) => {
                for expr in memo[group.0].iter() {
                    states.push_back((expr.clone(), state.clone()));
                }
            }
            CoroutineState::Yielded(Action::InvokeSelf(expr)) => {
                let results = apply(memo, expr, initial_state.clone());
                for result in results {
                    let expr = match result {
                        Plan::MultiJoin(groups) => Expr::MultiJoin(
                            groups
                                .iter()
                                .map(|g| {
                                    let Plan::Group(g) = g.as_ref() else {
                                        unreachable!()
                                    };
                                    g.clone()
                                })
                                .collect(),
                        ),
                        _ => unimplemented!(),
                    };
                    states.push_back((expr, state.clone()));
                }
            }
            CoroutineState::Complete(plan) => {
                if let Some(plan) = plan {
                    complete.push(plan);
                }
            }
        }
    }
    complete
}

fn main() {
    // group 5: Join 4, 1
    // group 4: Join 2, 3
    // group 1, 2, 3: Scan
    let memo = vec![
        vec![],
        vec![Expr::Scan],
        vec![Expr::Scan],
        vec![Expr::Scan],
        vec![
            Expr::Join(GroupId(2), GroupId(3)),
            Expr::Join(GroupId(3), GroupId(2)),
        ],
        vec![
            Expr::Join(GroupId(1), GroupId(4)),
            Expr::Join(GroupId(4), GroupId(1)),
        ],
    ];

    for plan in apply_group(&memo, GroupId(5), join_assoc()) {
        println!("plan = {:?}", plan);
    }
    println!("---");
    for plan in apply_group(&memo, GroupId(5), multi_join_loop()) {
        println!("plan = {:?}", plan);
    }
    println!("---");
    for plan in apply_group(&memo, GroupId(5), multi_join_recursive()) {
        println!("plan = {:?}", plan);
    }
}

simple_no_recursion.rs

#![feature(coroutines, coroutine_trait, coroutine_clone)]

use std::collections::VecDeque;
use std::{ops::Coroutine, sync::Arc};

#[derive(Clone, Copy, Debug)]
struct GroupId(usize);

impl std::fmt::Display for GroupId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "!{}", self.0)
    }
}

#[derive(Clone, Debug)]
enum Expr {
    Join(GroupId, GroupId),
    MultiJoin(Vec<GroupId>),
    Scan,
    Placeholder,
}

#[derive(Clone, Debug)]
enum Plan {
    Join(Arc<Plan>, Arc<Plan>),
    MultiJoin(Vec<Arc<Plan>>),
    Group(GroupId),
}

macro_rules! expand_group {
    ($x:ident) => {
        yield $x
    };
}

fn join_assoc() -> impl Coroutine<Expr, Yield = GroupId, Return = Option<Plan>> + Clone {
    #[coroutine]
    |binding: Expr| {
        if let Expr::Join(left_group, c) = binding {
            let left = expand_group!(left_group);
            if let Expr::Join(a, b) = left {
                return Some(Plan::Join(
                    Plan::Group(a).into(),
                    Plan::Join(Plan::Group(b).into(), Plan::Group(c).into()).into(),
                ));
            }
        }
        None
    }
}

fn multi_join() -> impl Coroutine<Expr, Yield = GroupId, Return = Option<Plan>> + Clone {
    #[coroutine]
    |mut binding: Expr| {
        let mut collected_groups = Vec::new();
        let mut last_right = None;
        while let Expr::Join(left, right) = binding {
            last_right = Some(right);
            binding = expand_group!(right);
            collected_groups.push(left);
        }
        if let Some(last_right) = last_right {
            collected_groups.push(last_right);
            if collected_groups.len() > 2 {
                let collected_groups = collected_groups
                    .into_iter()
                    .map(|group| Arc::new(Plan::Group(group)))
                    .collect();
                return Some(Plan::MultiJoin(collected_groups));
            }
        }
        None
    }
}

use std::{ops::CoroutineState, pin::Pin};

fn apply(
    groups: Vec<Vec<Expr>>,
    group: GroupId,
    rule: impl Coroutine<Expr, Yield = GroupId, Return = Option<Plan>> + Clone + Unpin,
) -> Vec<Plan> {
    let mut states = VecDeque::new();
    states.push_back((group, rule)); // initial state, let join_assoc explore group 5
    let mut complete = Vec::<Plan>::new();
    while let Some((group, state)) = states.pop_front() {
        for expr in groups[group.0].iter() {
            let mut state = state.clone();
            let state_ = Pin::new(&mut state);
            match state_.resume(expr.clone()) {
                CoroutineState::Yielded(group) => {
                    states.push_back((group, state));
                }
                CoroutineState::Complete(plan) => {
                    if let Some(plan) = plan {
                        complete.push(plan);
                    }
                }
            }
        }
    }
    complete
}

fn main() {
    // group 5: Join 4, 1
    // group 4: Join 2, 3
    // group 1, 2, 3: Scan
    let groups = vec![
        vec![],
        vec![Expr::Scan],
        vec![Expr::Scan],
        vec![Expr::Scan],
        vec![
            Expr::Join(GroupId(2), GroupId(3)),
            Expr::Join(GroupId(3), GroupId(2)),
        ],
        vec![Expr::Join(GroupId(1), GroupId(4)), Expr::Join(GroupId(4), GroupId(1))],
    ];

    for plan in apply(groups.clone(), GroupId(5), join_assoc()) {
        println!("plan = {:?}", plan);
    }

    for plan in apply(groups.clone(), GroupId(5), multi_join()) {
        println!("plan = {:?}", plan);
    }
}