Nested polonius-the-crab

nested_polonius.rs

use ::polonius_the_crab::prelude::*;

struct DummyStream {
    buffer: Vec<u8>,
    frame_count: u32,
}

impl DummyStream {
    fn new(width: usize, height: usize) -> Self {
        Self {
            buffer: vec![0; width * height],
            frame_count: 0,
        }
    }

    // Returns a reference to a camera buffer owned by the camera driver.
    fn next(&mut self) -> &[u8] {
        // Update the buffer with a new "frame"
        for (i, byte) in self.buffer.iter_mut().enumerate() {
            *byte = ((i as u32 + self.frame_count) % 256) as u8;
        }
        self.frame_count += 1;
        &self.buffer
    }
}

struct TestCameraMono {
    left_stream: DummyStream,
}

impl TestCameraMono {
    fn new() -> Self {
        Self {
            left_stream: DummyStream::new(640, 480), // Example resolution
        }
    }

    fn next_frame(&mut self) -> &[u8] {
        // {
        //     let l_buf = self.left_stream.next();

        //     // Sometimes, we want to skip this frame based on some condition that
        //     // can only be calculated when we have access to the buffer.
        //     // In the camera case, this will be when a frame is too old compared to the other camera.
        //     let delta = l_buf[0] - 127;

        //     let condition = delta < 80;
        //     if condition {
        //         return l_buf;
        //     }
        // }
        let mut this = self;
        polonius!(|this| -> &'polonius [u8] {
            let l_buf = this.left_stream.next();

            // Sometimes, we want to skip this frame based on some condition that
            // can only be calculated when we have access to the buffer.
            // In the camera case, this will be when a frame is too old compared to the other camera.
            let delta = l_buf[0] - 127;

            let condition = delta < 80;
            if condition {
                polonius_return!(l_buf);
            }
        });

        this.left_stream.next()
    }
}

struct TestCameraStereo {
    left_stream: DummyStream,
    right_stream: DummyStream,
}

enum DeltaRes {
    LeftLagging,
    RightLagging,
    Synchronized,
}

impl DeltaRes {
    pub fn from_delta(delta: u8) -> Self {
        if delta < 80 {
            DeltaRes::LeftLagging
        } else if delta < 255 - 80 {
            DeltaRes::Synchronized
        } else {
            DeltaRes::RightLagging
        }
    }
}

impl TestCameraStereo {
    fn new() -> Self {
        Self {
            left_stream: DummyStream::new(640, 480), // Example resolution
            right_stream: DummyStream::new(640, 480), // Example resolution
        }
    }

    fn next_frame<'r>(&'r mut self) -> (&'r [u8], &'r [u8]) {
        use ::polonius_the_crab::*;

        let mut left_stream = &mut self.left_stream;
        let mut right_stream = &mut self.right_stream;

        match polonius::<_, _, ForLt![<'left> = (&'left [u8], &'r [u8])]>(
            &mut self.left_stream,
            |left_stream| match polonius::<_, _, ForLt![<'right> = (Option<&'_ [u8]>, &'right [u8])]>(
                //                                                          'left
                &mut self.right_stream,
                |right_stream| {
                    let l_buf = left_stream.next();
                    let r_buf = right_stream.next();

                    // Sometimes, we want to skip this frame based on some condition that
                    // can only be calculated when we have access to the buffer.
                    // In the camera case, this will be when a frame is too old compared to the other camera.
                    let delta = l_buf[0] - r_buf[0];

                    match DeltaRes::from_delta(delta)
                    // polonius logic: we are within the *inner* `polonius()` call, to determine
                    // whether we are outputting anything `right_stream`-dependent, such as `r_buf`,
                    // in which case we have to be `PoloniusResult::Borrowing` (and allowed to use
                    // `r_buf`); otherwise we return `PoloniusResult::Owned` with everything but
                    // `r_buf`/`right_stream`-related stuff, and the *caller* will receive the
                    // original `right_stream` in the `input_borrow` spot.
                    {
                        DeltaRes::Synchronized => {
                            // Cameras are synchronized, and we do not need to update any of them
                            // (l_buf, r_buf)
                            // \-> mentions r_buf: `Borrowing`
                            PoloniusResult::Borrowing((Some(l_buf), r_buf))
                        }
                        DeltaRes::LeftLagging => {
                            // Left is lagging behind, update it
                            // (self.left_stream.next(), r_buf)
                            // \-> mentions r_buf: `Borrowing`
                            //   + but we'll also need to convey, to the caller, that
                            //     we're not interested in `l_buf` anymore.
                            //     Hence the `Option`.
                            PoloniusResult::Borrowing((None, r_buf))
                        }
                        DeltaRes::RightLagging => {
                            // Right is lagging behind, update it
                            // (l_buf, self.right_stream.next())
                            // \-> we want to use `right_stream`, so we skip any further mentions
                            //     of `r_buf`, and return something `Owned`.
                            PoloniusResult::Owned {
                                input_borrow: Placeholder, // <- please give me back `right_stream`
                                value: l_buf, // "pass-through" companion payload.
                            }
                        }
                    }
                }
            )        
            // polonius logic: we are within the *outer* `polonius()` call, receiving
            // the output of the inner `polonius`: we'll either have a `return`-able `r_buf`
            // (friendly lifetime in it), or access to the `input_borrow` to the `right_stream`,
            // so as to query again.
            // We still have to determine whether we are outputting anything `left_stream`-dependent,
            // such as `l_buf`;
            /* .match */ {
                /*Inner*/PoloniusResult::Borrowing((Some(l_buf), r_buf)) => {
                    /*Outer*/PoloniusResult::Borrowing(
                        (l_buf, r_buf)
                    )
                },
                // (self.left_stream.next(), r_buf)
                /*Inner*/PoloniusResult::Borrowing((None, r_buf)) => {
                    /*Outer*/PoloniusResult::Owned {
                        input_borrow: Placeholder, // <- please give me back `left_stream`
                        value: r_buf,
                    }
                },
                // (l_buf, self.right_stream.next())
                /*Inner*/PoloniusResult::Owned {
                    input_borrow: right_stream, // <- we got `right_stream` back!
                    value: l_buf,
                } => /*Outer*/PoloniusResult::Borrowing(
                    (l_buf, right_stream.next())
                ),
            }
        )
        // Outside the outer polonius: time to just inject `left_stream.next()` wherever we
        // receive the `input_borrow: left_stream` back:
        /* .match */ {
            PoloniusResult::Borrowing((l_buf, r_buf)) => (
                l_buf, r_buf,
            ),

            // (self.left_stream.next(), r_buf)
            PoloniusResult::Owned { value: r_buf, input_borrow: left_stream } => (
                left_stream.next(), r_buf,
            ),
        }
    }
}

fn main() {
    let mut mono_cam = TestCameraMono::new();

    let l_buf = mono_cam.next_frame();
    println!("mono_cam: {}", l_buf[10]);

    let mut stereo_cam = TestCameraStereo::new();
    let (l_buf, r_buf) = stereo_cam.next_frame();
    println!("stereo_cam: {}, {}", l_buf[10], r_buf[10]);
}