196: Delimited Continuations

// Delimited Continuations via Simulation
//
// OCaml 5 effects give `shift` and `reset` natively.
// Rust has neither, but we can simulate the classic use case:
// nondeterminism (list-of-successes / Prolog-style search).
//
// `reset` marks the delimiter; `shift` captures the continuation up to
// the nearest `reset` and passes it to a function that decides what to do.
//
// The standard translation: shift/reset for nondeterminism ≅ the list monad.
// We implement:
//   1. A direct list-monad simulation (most idiomatic, same semantics)
//   2. A CPS-based shift/reset using Rc to handle multiple invocations

use std::rc::Rc;

// ── 1. List monad (direct translation of shift/reset nondeterminism) ──────

fn nd_return<A: Clone>(x: A) -> Vec<A> {
    vec![x]
}

fn nd_bind<A, B, F: Fn(A) -> Vec<B>>(xs: Vec<A>, f: F) -> Vec<B> {
    xs.into_iter().flat_map(f).collect()
}

/// "choose" = shift in the OCaml example
fn choose<A: Clone>(xs: Vec<A>) -> Vec<A> {
    xs
}

/// "guard" — zero or one result based on a predicate
fn nd_guard(cond: bool) -> Vec<()> {
    if cond { vec![()] } else { vec![] }
}

// ── 2. CPS-based shift/reset ────────────────────────────────────────────────
//
// We use Rc so continuations can be invoked multiple times (needed for
// nondeterminism where the continuation runs once per branch).
//
// type DelimComp<A, R> = Rc<dyn Fn(Rc<dyn Fn(A) -> R>) -> R>
// reset(m)  = m(Rc::new(|x| x))
// shift(g)  = Rc::new(move |k| g(k))

type Cont<A, R> = Rc<dyn Fn(A) -> R>;
type DelimComp<A, R> = Rc<dyn Fn(Cont<A, R>) -> R>;

fn dc_return<A: Clone + 'static, R: 'static>(x: A) -> DelimComp<A, R> {
    Rc::new(move |k: Cont<A, R>| k(x.clone()))
}

fn dc_bind<A: Clone + 'static, B: Clone + 'static, R: Clone + 'static>(
    m: DelimComp<A, R>,
    f: Rc<dyn Fn(A) -> DelimComp<B, R>>,
) -> DelimComp<B, R> {
    Rc::new(move |k: Cont<B, R>| {
        let f2 = f.clone();
        let k2 = k.clone();
        m(Rc::new(move |a: A| {
            let mb = f2(a);
            mb(k2.clone())
        }))
    })
}

fn dc_reset<A: Clone + 'static>(m: DelimComp<A, A>) -> A {
    m(Rc::new(|x| x))
}

/// shift: g receives the current continuation
fn dc_shift<A: 'static, R: 'static, G>(g: G) -> DelimComp<A, R>
where
    G: Fn(Cont<A, R>) -> R + 'static,
{
    Rc::new(move |k: Cont<A, R>| g(k))
}

// ── 3. Nondeterminism via CPS shift/reset ─────────────────────────────────

fn choose_cps<A: Clone + 'static>(xs: Vec<A>) -> DelimComp<A, Vec<A>> {
    Rc::new(move |k: Cont<A, Vec<A>>| {
        xs.iter().flat_map(|x| k(x.clone())).collect()
    })
}

fn nd_return_cps<A: Clone + 'static>(x: A) -> DelimComp<A, Vec<A>> {
    dc_return(x)
}

fn nd_reset_cps<A: Clone + 'static>(m: DelimComp<A, Vec<A>>) -> Vec<A> {
    m(Rc::new(|x| vec![x]))
}

fn main() {
    // ── List monad (direct) ────────────────────────────────────────────
    println!("=== List monad nondeterminism ===");

    // Choose x from [1..3], y from [x..x+2], return (x,y) if x+y<=5
    let results = nd_bind(choose(vec![1, 2, 3]), |x| {
        nd_bind(choose(vec![x, x + 1, x + 2]), move |y| {
            nd_bind(nd_guard(x + y <= 5), move |_| nd_return((x, y)))
        })
    });
    println!("x in [1..3], y in [x..x+2], x+y<=5: {:?}", results);

    // Pythagorean triples up to 10
    let triples = nd_bind(choose((1..=10).collect()), |a| {
        nd_bind(choose((a..=10).collect()), move |b| {
            nd_bind(choose((b..=10).collect()), move |c| {
                nd_bind(nd_guard(a * a + b * b == c * c), move |_| nd_return((a, b, c)))
            })
        })
    });
    println!("Pythagorean triples ≤10: {:?}", triples);

    println!();

    // ── CPS shift/reset nondeterminism ────────────────────────────────
    println!("=== CPS shift/reset nondeterminism ===");
    let cps_results = nd_reset_cps(dc_bind(
        choose_cps(vec![1_i32, 2, 3]),
        Rc::new(|x| {
            dc_bind(
                choose_cps(vec![10_i32, 20]),
                Rc::new(move |y| nd_return_cps(x + y)),
            )
        }),
    ));
    println!("[1,2,3] × [10,20] sums: {:?}", cps_results);

    println!();

    // ── Classic "all permutations" via list monad ─────────────────────
    println!("=== Permutations of [1,2,3] ===");

    fn remove_one(x: i32, list: &[i32]) -> Vec<i32> {
        list.iter().filter(|&&y| y != x).cloned().collect()
    }

    fn perms(list: Vec<i32>) -> Vec<Vec<i32>> {
        if list.is_empty() {
            return vec![vec![]];
        }
        nd_bind(list.clone(), |x| {
            let rest = remove_one(x, &list);
            nd_bind(perms(rest), move |mut tail| {
                tail.insert(0, x);
                nd_return(tail)
            })
        })
    }

    let ps = perms(vec![1, 2, 3]);
    println!("permutations: {} total", ps.len());
    for p in &ps {
        println!("  {:?}", p);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_list_monad_combinations() {
        let results = nd_bind(choose(vec![1, 2]), |x| {
            nd_bind(choose(vec![10, 20]), move |y| nd_return(x + y))
        });
        assert_eq!(results.len(), 4);
        assert!(results.contains(&11));
        assert!(results.contains(&22));
    }

    #[test]
    fn test_pythagorean_triple_3_4_5() {
        let triples = nd_bind(choose((1..=10).collect()), |a| {
            nd_bind(choose((a..=10).collect()), move |b| {
                nd_bind(choose((b..=10).collect()), move |c| {
                    nd_bind(nd_guard(a * a + b * b == c * c), move |_| nd_return((a, b, c)))
                })
            })
        });
        assert!(triples.contains(&(3, 4, 5)));
        assert!(triples.contains(&(6, 8, 10)));
    }

    #[test]
    fn test_cps_shift_reset_sums() {
        let cps_results = nd_reset_cps(dc_bind(
            choose_cps(vec![1_i32, 2, 3]),
            Rc::new(|x| {
                dc_bind(
                    choose_cps(vec![10_i32, 20]),
                    Rc::new(move |y| nd_return_cps(x + y)),
                )
            }),
        ));
        assert_eq!(cps_results.len(), 6);
        assert!(cps_results.contains(&11));
        assert!(cps_results.contains(&23));
    }

    #[test]
    fn test_guard_filters_with_empty() {
        let results = nd_bind(choose(vec![1, 2, 3, 4, 5]), |x| {
            nd_bind(nd_guard(x % 2 == 0), move |_| nd_return(x))
        });
        assert_eq!(results, vec![2, 4]);
    }
}