fieldwork-borrow 0.1.0

Field-level borrow splitting for Rust
Documentation
//! e9: async split - socket.split() + &mut self method

use fieldwork_borrow::with_;
use std::marker::PhantomData;

struct WriteHalf<'a>(PhantomData<&'a ()>);
struct ReadHalf<'a>(PhantomData<&'a ()>);

#[derive(Debug, Clone, Default)]
struct Frontier {
    value: u32,
}

#[derive(Debug)]
struct ConnectionState {
    remote_frontier: Frontier,
}

struct TcpStream {
    data: u32,
}

impl TcpStream {
    pub fn split(&mut self) -> (ReadHalf<'_>, WriteHalf<'_>) {
        (ReadHalf(PhantomData), WriteHalf(PhantomData))
    }
}

struct BufReader<T> {
    inner: T,
    lines: Vec<String>,
    index: usize,
}

impl<T> BufReader<T> {
    pub fn new(inner: T, lines: Vec<String>) -> Self {
        BufReader { inner, lines, index: 0 }
    }

    pub async fn next_line(&mut self) -> Option<String> {
        if self.index < self.lines.len() {
            let line = self.lines[self.index].clone();
            self.index += 1;
            Some(line)
        } else {
            None
        }
    }
}

async fn send_message<W>(_stream: &mut W) {}

fn some_code(frontier: &mut Frontier) {
    frontier.value += 1;
}

pub struct Protocol<'socket> {
    socket: &'socket mut TcpStream,
    state: Option<ConnectionState>,
}

impl<'socket> Protocol<'socket> {
    // Original problem code (won't compile):
    // pub async fn problem(&mut self) {
    //     let (reader, mut writer) = self.socket.split();
    //     let mut reader = BufReader::new(reader, vec!["hello".into()]);
    //     loop {
    //         let line = reader.next_line();
    //         if let Some(line) = line.await {
    //             self.handle_message(line, &mut writer).await;
    //         } else {
    //             break;
    //         }
    //     }
    // }

    // The issue: socket.split() borrows socket, but handle_message needs &mut self
    // Solution: inline handle_message's logic using with_!

    pub async fn problem(&mut self) {
        let (reader, mut writer) = self.socket.split();
        let mut reader = BufReader::new(reader, vec!["hello".into(), "world".into()]);

        with_!(self, {
            methods {
                handle_message(mut state, _msg: String, writer: &mut WriteHalf<'_>) {
                    let conn_state = state.as_mut().unwrap();
                    some_code(&mut conn_state.remote_frontier);
                    // Can't actually await here since we're in a sync block
                    // But we can prepare the call
                    let _ = writer;
                }
            }
            body {
                loop {
                    let line = reader.next_line();
                    if let Some(line) = line.await {
                        self.handle_message(line, &mut writer);
                    } else {
                        break;
                    }
                }
            }
        });
    }

    async fn _handle_message(&mut self, _msg: String, writer: &mut WriteHalf<'_>) {
        let state = self.state.as_mut().unwrap();
        some_code(&mut state.remote_frontier);
        send_message(writer).await;
    }
}

#[tokio::test]
async fn test_e9() {
    let mut stream = TcpStream { data: 42 };
    let mut protocol = Protocol {
        socket: &mut stream,
        state: Some(ConnectionState {
            remote_frontier: Frontier { value: 0 },
        }),
    };

    protocol.problem().await;

    // handle_message was called twice (for "hello" and "world")
    assert_eq!(protocol.state.as_ref().unwrap().remote_frontier.value, 2);
}