Shared via Rust Playground

playground.rs

#![feature(fnbox)]
use std::boxed::FnBox;
use std::marker::PhantomData;
use std::mem::replace;
use std::fmt::{Debug, Formatter};
use std::fmt;
use std::collections::BTreeMap;

struct World {
    magic: BTreeMap<usize, String>,
}

struct Thunk<A>(Box<Work<End=A>>);

#[derive(Debug)]
enum Step<A> {
    Done(A),
    More(Thunk<A>),
}

impl<A> Step<A> {
    fn try_step(self, world: &mut World) -> Option<Self> {
        match self {
            Step::Done(_) => None,
            Step::More(thunk) => Some(thunk.0.step(world)),
        }
    }
}

impl<A> Debug for Thunk<A> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "<thunk>")
    }
}

impl<A: 'static + Send> Thunk<A> {
    fn new(a: A) -> Self {
        Thunk(Box::new(Done(a)))
    }
    fn map<B, F>(self, f: F) -> Thunk<B>
        where F: 'static + Send + FnOnce(A, &mut World) -> B
    {
        Thunk(Box::new(Map {
            work: self.0,
            f: f,
        }))
    }
    fn and_then<B, F>(self, f: F) -> Thunk<B>
        where F: 'static + Send + FnOnce(A, &mut World) -> Thunk<B>
    {
        Thunk(Box::new(FlatMap {
            work: self.0,
            f: |a: A, world: &mut World| {f(a, world).0},
        }))
    }
    fn run(self, world: &mut World) -> A {
        let mut work = Some(self.0.erase());
        loop {
            let current_work = replace(&mut work, None).unwrap();
            let step = current_work.step(world);
            match step {
                Step::Done(a) => return a,
                Step::More(w) => {
                    replace(&mut work, Some(w.0));
                }
            }
        }
    }
    
    fn run_log(self, world: &mut World) -> A
        where A: Debug
    {
        let mut work = Some(self.0.erase());
        loop {
            let current_work = replace(&mut work, None).unwrap();
            let step = current_work.step(world);
            println!("Step: {:?}", step);
            match step {
                Step::Done(a) => return a,
                Step::More(w) => {
                    replace(&mut work, Some(w.0));
                }
            }
        }
    }
}

trait Work: Send {
    type End;
    fn step(self: Box<Self>, world: &mut World) -> Step<Self::End>;
    fn erase(self: Box<Self>) -> Box<Work<End=Self::End>>;
}

struct Done<A>(A);
impl<A: 'static + Send> Work for Done<A> {
    type End = A;
    fn erase(self: Box<Self>) -> Box<Work<End=A>> {
        Box::new(*self)
    }
    fn step(self: Box<Self>, world: &mut World) -> Step<A> {
        Step::Done(self.0)
    }
}

struct Map<A, F> {
    work: Box<Work<End=A>>,
    f: F,
}
impl<A: 'static + Send, B, F: 'static + Send + FnOnce(A, &mut World) -> B>
    Work for Map<A, F>
{
    type End = B;
    fn erase(self: Box<Self>) -> Box<Work<End=B>> {
        Box::new(*self)
    }
    fn step(self: Box<Self>, world: &mut World) -> Step<B> {
        let s = *self;
        let work = s.work;
        let f = s.f;
        match work.step(world) {
            Step::Done(a) => Step::Done(f(a, world)),
            Step::More(w) => Step::More(Thunk(Box::new(
                Map {
                    work: w.0,
                    f: f,
                }
            ))),
        }
    }
}

struct FlatMap<A, F> {
    work: Box<Work<End=A>>,
    f: F,
}
impl<A: 'static + Send, B, F: 'static + Send + FnOnce(A, &mut World) -> Box<Work<End = B>>>
    Work for FlatMap<A, F>
{
    type End = B;
    fn erase(self: Box<Self>) -> Box<Work<End=B>> {
        Box::new(*self)
    }
    fn step(self: Box<Self>, world: &mut World) -> Step<B> {
        let s = *self;
        let work = s.work;
        let f = s.f;
        match work.step(world) {
            Step::Done(a) => Step::More(Thunk(f(a, world))),
            Step::More(w) => Step::More(Thunk(Box::new(
                FlatMap {
                    work: w.0,
                    f: f,
                }
            ))),
        }
    }
}

fn read(key: usize) -> Thunk<Option<String>> {
    Thunk::new(()).map(move |_, world| {
        world.magic.remove(&key)
    })
}

fn put(key: usize, value: String) -> Thunk<()> {
    Thunk::new(()).map(move |_, world| {
        world.magic.insert(key, value);
    })
}

fn transpose<A: 'static + Send>(mut thunks: Vec<Thunk<A>>) -> Thunk<Vec<A>> {
    if let Some(t) = thunks.pop() {
        t.and_then(move |a, _| {
            let rest_thunk = transpose(thunks);
            rest_thunk.map(move |mut rest, _| {
                rest.push(a);
                rest
            })
        })
    } else {
        Thunk::new(Vec::new())
    }
}

fn main() {
    let thunk = Thunk::new(12).map(|i, _| {
        let mut keys = Vec::new();
        for j in 0..i {
            keys.push(j);
        }
        keys
    }).map(|keys, _| {
        let puts = keys.into_iter().map(|k| {
            let value = format!("Value for Key {}", k);
            put(k, value)
        }).collect();
        puts
    }).and_then(|inserts, _| {
        transpose(inserts)
    }).and_then(|_,_| {
        let mut reads = Vec::new();
        for i in 0..13 {
            reads.push(read(i));
        }
        transpose(reads)
    });
    let mut world = World { magic: BTreeMap::new() };
    println!("Thunk result: {:?}", thunk.run(&mut world));
    println!("Final world: {:?}", world.magic);
}