Deconstructing awaits

await_deconstructed.rs

use std::rc::Rc;
use std::thread::sleep;
use std::time::{Duration, SystemTime};

#[derive(Clone)]
enum AwaitState<T> {
    Uninitialized,
    Running,
    Complete(T),
}

trait Awaitable {
    type Item;

    fn next(&mut self) -> AwaitState<Self::Item>;
}

fn blocking_wait<T>(mut awaitable: impl Awaitable<Item=T>) -> T {
    loop {
        if let AwaitState::Complete(res) = awaitable.next() {
            return res;
        } else {
            sleep(Duration::from_millis(10));
        }
    }
}

struct LongRunning<T> {
    duration: Duration,
    result: Rc<T>,
    state: AwaitState<Rc<T>>,
    start: Option<SystemTime>,
}

impl<T> LongRunning<T> {
    fn new(duration: Duration, result: T) -> Self {
        return Self {
            duration,
            result: Rc::new(result),
            state: AwaitState::Uninitialized,
            start: None,
        };
    }
}

impl<T> Awaitable for LongRunning<T> {
    type Item = Rc<T>;

    fn next(&mut self) -> AwaitState<Self::Item> {
        match &self.state {
            AwaitState::Uninitialized => {
                self.state = AwaitState::Running;
                self.start = Some(SystemTime::now());
                AwaitState::Running
            }
            AwaitState::Running => {
                let elapsed = SystemTime::now().duration_since(self.start.unwrap()).unwrap();
                if elapsed > self.duration {
                    self.state = AwaitState::Complete(self.result.clone());
                }
                self.state.clone()
            }
            AwaitState::Complete(_) => self.state.clone()
        }
    }
}

struct Gather<A, B> {
    lhs: Box<dyn Awaitable<Item=Rc<A>>>,
    rhs: Box<dyn Awaitable<Item=Rc<B>>>,
}

impl<A, B> Gather<A, B> {
    fn new(lhs: Box<dyn Awaitable<Item=Rc<A>>>, rhs: Box<dyn Awaitable<Item=Rc<B>>>) -> Self {
        return Self {
            lhs,
            rhs,
        };
    }
}

impl<A, B> Awaitable for Gather<A, B> {
    type Item = (Rc<A>, Rc<B>);

    fn next(&mut self) -> AwaitState<Self::Item> {
        if let (AwaitState::Complete(lhs), AwaitState::Complete(rhs)) = (self.lhs.next(), self.rhs.next()) {
            return AwaitState::Complete((lhs.clone(), rhs.clone()));
        }
        return AwaitState::Running;
    }
}

struct MyFirstInnerFunction {
    // state machine
    function_state: usize,
    // arguments/results
    args: (SystemTime,),
    result: Rc<i64>,
    // local variables
    awaitable: LongRunning<i64>,
}

impl MyFirstInnerFunction {
    fn new(args: (SystemTime,)) -> Self {
        Self {
            args,
            function_state: 0,
            result: Rc::new(0),
            awaitable: LongRunning::new(Duration::from_millis(800), 42),
        }
    }

    fn function__0(&mut self) -> AwaitState<<Self as Awaitable>::Item> {
        let elapsed = SystemTime::now().duration_since(self.args.0).unwrap().as_millis();
        println!("First inner function started at {elapsed} ms");
        self.function_state = 1;
        AwaitState::Running
    }

    fn function__1(&mut self) -> AwaitState<<Self as Awaitable>::Item> {
        if let AwaitState::Complete(res) = self.awaitable.next() {
            self.function_state = 2;
            self.result = res;
        }
        AwaitState::Running
    }

    fn function__2(&mut self) -> AwaitState<<Self as Awaitable>::Item> {
        let elapsed = SystemTime::now().duration_since(self.args.0).unwrap().as_millis();
        let val = *self.result;
        println!("First inner completed at {elapsed} ms with result {val}");
        self.function_state = 3;
        AwaitState::Complete(self.result.clone())
    }
}

impl Awaitable for MyFirstInnerFunction {
    type Item = Rc<i64>;

    fn next(&mut self) -> AwaitState<Self::Item> {
        match self.function_state {
            0 => self.function__0(),
            1 => self.function__1(),
            2 => self.function__2(),
            _ => AwaitState::Complete(self.result.clone())
        }
    }
}

struct MySecondInnerFunction {
    // state machine
    function_state: usize,
    // arguments/results
    args: (SystemTime,),
    result: Rc<i64>,
    // local variables
    awaitable: LongRunning<i64>,
}

impl MySecondInnerFunction {
    fn new(args: (SystemTime,)) -> Self {
        Self {
            args,
            function_state: 0,
            result: Rc::new(0),
            awaitable: LongRunning::new(Duration::from_millis(400), 137),
        }
    }

    fn function__0(&mut self) -> AwaitState<<Self as Awaitable>::Item> {
        let elapsed = SystemTime::now().duration_since(self.args.0).unwrap().as_millis();
        println!("Second inner function started at {elapsed} ms");
        self.function_state = 1;
        AwaitState::Running
    }

    fn function__1(&mut self) -> AwaitState<<Self as Awaitable>::Item> {
        if let AwaitState::Complete(res) = self.awaitable.next() {
            self.function_state = 2;
            self.result = res;
        }
        AwaitState::Running
    }

    fn function__2(&mut self) -> AwaitState<<Self as Awaitable>::Item> {
        let elapsed = SystemTime::now().duration_since(self.args.0).unwrap().as_millis();
        let val = *self.result;
        println!("Second inner completed at {elapsed} ms with result {val}");
        self.function_state = 3;
        AwaitState::Complete(self.result.clone())
    }
}


impl Awaitable for MySecondInnerFunction {
    type Item = Rc<i64>;

    fn next(&mut self) -> AwaitState<Self::Item> {
        match self.function_state {
            0 => self.function__0(),
            1 => self.function__1(),
            2 => self.function__2(),
            _ => AwaitState::Complete(self.result.clone())
        }
    }
}


struct MyOuterFunction {
    // state machine
    function_state: usize,
    // arguments/results
    args: (SystemTime,),
    result: Rc<f64>,
    // local variables
    awaitable: Gather<i64, i64>,
    pi_computer: LongRunning<f64>,
    fn_results: (Rc<i64>, Rc<i64>),
    value_of_pi: Rc<f64>,
}

impl MyOuterFunction {
    fn new(args: (SystemTime,)) -> Self {
        Self {
            args,
            function_state: 0,
            result: Rc::new(0.),
            awaitable: Gather::new(Box::new(MyFirstInnerFunction::new(args)), Box::new(MySecondInnerFunction::new(args))),
            pi_computer: LongRunning::new(Duration::from_millis(500), 3.14),
            fn_results: (Rc::new(0), Rc::new(0)),
            value_of_pi: Rc::new(0.),
        }
    }

    fn function__0(&mut self) -> AwaitState<<Self as Awaitable>::Item> {
        let elapsed = SystemTime::now().duration_since(self.args.0).unwrap().as_millis();
        println!("My outer function started at {elapsed}");
        self.function_state = 1;
        AwaitState::Running
    }

    fn function__1(&mut self) -> AwaitState<<Self as Awaitable>::Item> {
        if let AwaitState::Complete(results) = self.awaitable.next() {
            self.function_state = 2;
            self.fn_results = results;
        }
        AwaitState::Running
    }

    fn function__2(&mut self) -> AwaitState<<Self as Awaitable>::Item> {
        let elapsed = SystemTime::now().duration_since(self.args.0).unwrap().as_millis();
        println!("My outer function continues at {elapsed}, computing pi");
        self.function_state = 3;
        AwaitState::Running
    }

    fn function__3(&mut self) -> AwaitState<<Self as Awaitable>::Item> {
        if let AwaitState::Complete(res) = self.pi_computer.next() {
            self.function_state = 4;
            self.value_of_pi = res;
        }
        AwaitState::Running
    }

    fn function__4(&mut self) -> AwaitState<<Self as Awaitable>::Item> {
        let elapsed = SystemTime::now().duration_since(self.args.0).unwrap().as_millis();
        println!("My outer function continues at {elapsed}, returning value");
        self.result = Rc::new((*self.fn_results.0 as f64) * (*self.fn_results.1 as f64) * (*self.value_of_pi));
        self.function_state = 5;
        return AwaitState::Complete(self.result.clone());
    }
}


impl Awaitable for MyOuterFunction {
    type Item = Rc<f64>;

    fn next(&mut self) -> AwaitState<Self::Item> {
        match self.function_state {
            0 => self.function__0(),
            1 => self.function__1(),
            2 => self.function__2(),
            3 => self.function__3(),
            4 => self.function__4(),
            _ => AwaitState::Complete(self.result.clone())
        }
    }
}


fn main() {
    let start = SystemTime::now();
    let function = MyOuterFunction::new((start,));
    let result = blocking_wait(function);
    let ms = SystemTime::now().duration_since(start).unwrap().as_millis();
    println!("Got: {result}, Took {ms} ms");
}

Same as

async long_running<T>(delay: Duration, result: T) -> T {
    task::sleep(delay).await;
    result
}

async fn first_inner_function(start) -> i64 {
    let elapsed = SystemTime::now().duration_since(start).unwrap().as_millis();
    println!("First inner function started at {elapsed} ms");
    let result: i64 = long_running(Duration::from_millis(800), 42).await;
    let elapsed = SystemTime::now().duration_since(start).unwrap().as_millis();
    println!("First inner completed at {elapsed} ms with result {result}");
    return result;
}

async fn second_inner_function(start) -> i64 {
    let elapsed = SystemTime::now().duration_since(start).unwrap().as_millis();
    println!("Second inner function started at {elapsed} ms");
    let result: i64 = long_running(Duration::from_millis(400), 137).await;
    let elapsed = SystemTime::now().duration_since(start).unwrap().as_millis();
    println!("Second inner completed at {elapsed} ms with result {result}");
    return result;
}

async fn outer_function(start) -> f64 {
    let elapsed = SystemTime::now().duration_since(start).unwrap().as_millis();
    println!("My outer function started at {elapsed}");
    let (r1, r2) = join!(first_inner_function(start), second_inner_function(start)).await;
    let elapsed = SystemTime::now().duration_since(start).unwrap().as_millis();
    println!("My outer function continues at {elapsed}, computing pi");
    let pi: f64 = long_running(Duration::from_millis(500), 3.14).await;
    let elapsed = SystemTime::now().duration_since(start).unwrap().as_millis();
    println!("My outer function continues at {elapsed}, returning value");
    return (r1 as f64) * (r2 as f64) * pi;
}