coroutine_parser.rs

#![feature(coroutines, coroutine_trait)]

use snafu::prelude::*;
use std::{
    ops::{Coroutine, CoroutineState},
    pin::pin,
    str,
};

#[derive(Debug, Copy, Clone)]
enum Thing<T> {
    NeedsMoreInput,
    Data(T),
}

#[derive(Debug, Copy, Clone)]
enum Token<'b> {
    Plus(&'b str),
    Minus(&'b str),
    Number(&'b str),
}

impl<'b> Token<'b> {
    fn bytes(&self) -> &'b [u8] {
        use Token::*;

        match self {
            Plus(v) => v.as_bytes(),
            Minus(v) => v.as_bytes(),
            Number(v) => v.as_bytes(),
        }
    }

    fn len(&self) -> usize {
        self.bytes().len()
    }
}

#[derive(Debug, Snafu)]
enum Error {
    BadNumber,
}

fn make_parser_number<'a>(
) -> impl Coroutine<&'a [u8], Yield = Thing<Option<&'a str>>, Return = Result<&'a [u8], Error>> {
    use self::Thing::*;

    |buf| {
        let mut buf: &[u8] = buf;

        while buf.is_empty() {
            buf = yield NeedsMoreInput;
        }

        let l = buf.iter().take_while(|b| b.is_ascii_digit()).count();
        if l > 0 {
            // Safety: LOL this is toy code.
            let n = unsafe {
                let n = buf.get_unchecked(..l);
                str::from_utf8_unchecked(n)
            };
            buf = yield Data(Some(n));
        } else {
            buf = yield Data(None);
        }

        Ok(buf)
    }
}

fn make_parser<'a>(
) -> impl Coroutine<&'a [u8], Yield = Thing<Token<'a>>, Return = Result<(), Error>> {
    use self::Thing::*;
    use self::Token::*;

    move |buf| {
        let mut buf: &[u8] = buf;

        loop {
            while buf.is_empty() {
                buf = yield NeedsMoreInput;
            }

            if buf.starts_with(b"-") {
                // Safety: LOL this is toy code.
                let m = unsafe {
                    let m = buf.get_unchecked(..1);
                    str::from_utf8_unchecked(m)
                };
                buf = yield Data(Minus(m));
            } else if buf.starts_with(b"+") {
                let p = unsafe {
                    let p = buf.get_unchecked(..1);
                    str::from_utf8_unchecked(p)
                };
                buf = yield Data(Plus(p));
            } else {
                use CoroutineState::*;

                // Why is `Box` needed / `pin!` doesn't work?
                let mut number_parser = Box::pin(make_parser_number());

                loop {
                    match number_parser.as_mut().resume(buf) {
                        Yielded(NeedsMoreInput) => {
                            buf = yield NeedsMoreInput;
                        }
                        Yielded(Data(Some(n))) => {
                            buf = yield Data(Number(n));
                        }
                        Yielded(Data(None)) => {
                            return BadNumberSnafu.fail();
                        }
                        Complete(b) => {
                            buf = b?;
                            break;
                        }
                    }
                }
            }
        }
    }
}

#[snafu::report]
fn main() -> Result<(), Error> {
    let inputs: &[&[u8]] = &[b"1+", b"2+3", b"+4"];
    let mut inputs = inputs.into_iter();

    let mut head: &[u8] = &[];

    let mut parser = pin!(make_parser());

    loop {
        use CoroutineState::*;
        use Thing::*;

        match parser.as_mut().resume(&head) {
            Yielded(NeedsMoreInput) => match inputs.next() {
                Some(b) => head = b,
                None => break Ok(()),
            },
            Yielded(Data(t)) => {
                eprintln!("{t:?}");
                head = &head[t.len()..]
            }
            Complete(c) => break c,
        }
    }
}