async_codegen/

util.rs

/*
 * Copyright © 2025 Anand Beh
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
use crate::common::SequenceConfig;
use crate::context::Context;
use crate::{IoOutput, Output, Writable, WritableSeq};
use std::convert::Infallible;
use std::marker::PhantomData;
use std::pin::pin;
use std::task::{Poll, Waker};

/// A wrapper intended to be used as a [SequenceConfig].
/// The tuple value should be a function that returns a writable for the data type.
pub struct WritableFromFunction<F>(pub F);

impl<F, T, W, O> SequenceConfig<T, O> for WritableFromFunction<F>
where
    O: Output,
    F: Fn(&T) -> W,
    W: Writable<O>,
{
    async fn write_datum(&self, datum: &T, output: &mut O) -> Result<(), O::Error> {
        let writable = (&self.0)(datum);
        writable.write_to(output).await
    }
}

/// An IO implementation that just writes to a string. The futures produced from methods
/// like [IoOutput::write] are complete instantly, and no errors are produced.
pub struct InMemoryIo<'s>(pub &'s mut String);

impl IoOutput for InMemoryIo<'_> {
    type Error = Infallible;

    async fn write(&mut self, value: &str) -> Result<(), Self::Error> {
        self.0.push_str(value);
        Ok(())
    }
}

/// An output implementation that just writes to a string. Produces no errors and the futures
/// complete instantly. This type is an analogue of [InMemoryIo].
pub struct InMemoryOutput<Ctx> {
    buf: String,
    context: Ctx,
}

impl<Ctx> InMemoryOutput<Ctx> {
    pub fn new(context: Ctx) -> Self {
        Self {
            buf: String::new(),
            context,
        }
    }
}

impl<Ctx> Output for InMemoryOutput<Ctx>
where
    Ctx: Context,
{
    type Io<'b>
        = InMemoryIo<'b>
    where
        Self: 'b;
    type Ctx = Ctx;
    type Error = Infallible;

    async fn write(&mut self, value: &str) -> Result<(), Self::Error> {
        self.buf.push_str(value);
        Ok(())
    }

    fn split(&mut self) -> (Self::Io<'_>, &Self::Ctx) {
        (InMemoryIo(&mut self.buf), &self.context)
    }

    fn context(&self) -> &Self::Ctx {
        &self.context
    }
}

impl<Ctx> InMemoryOutput<Ctx>
where
    Ctx: Context,
{
    /// Gets the string output of a single writable.
    ///
    /// Assumes that the only source of async is this [InMemoryOutput] itself, i.e. the writable
    /// will never return a pending future unless the output used with it returns a pending future.
    ///
    /// **Panics** if the writable returns a pending future (particularly, this may happen if
    /// a writable introduces its own async computations that do not come from the output)
    pub fn print_output<W>(context: Ctx, writable: &W) -> String
    where
        W: Writable<Self>,
    {
        Self::print_output_impl(context, writable).buf
    }

    fn print_output_impl<W>(context: Ctx, writable: &W) -> Self
    where
        W: Writable<Self>,
    {
        let mut output = Self {
            buf: String::new(),
            context,
        };
        let result = {
            let future = pin!(writable.write_to(&mut output));
            match future.poll(&mut std::task::Context::from_waker(Waker::noop())) {
                Poll::Pending => panic!("Expected a complete future"),
                Poll::Ready(result) => result,
            }
        };
        match result {
            Ok(()) => output,
            Err(e) => match e {}, // Unreachable
        }
    }
}

/// Turns a [WritableSeq] into an iterator over owned strings.
/// This calls [InMemoryOutput::print_output] for each writable produced by the sequence. It uses
/// Rust's async design in order to produce lazy iteration.
///
/// Iteration will **panic** if any writable returns a pending future (in particular, this may
/// happen if a writable introduces its own async computations that do not come from the output)
pub struct IntoStringIter<'f, Ctx, Seq>(
    string_iter::IntoStringIterStart<Ctx, Seq>,
    PhantomData<&'f ()>,
);

impl<'f, Ctx, Seq> IntoStringIter<'f, Ctx, Seq> {
    pub fn new(context: Ctx, sequence: Seq) -> Self {
        Self(
            string_iter::IntoStringIterStart { context, sequence },
            PhantomData,
        )
    }
}

impl<'f, Ctx, Seq> IntoIterator for IntoStringIter<'f, Ctx, Seq>
where
    Ctx: Context + 'f,
    Seq: WritableSeq<InMemoryOutput<Ctx>> + 'f,
{
    type Item = String;
    type IntoIter = ToStringIter<'f, Ctx, Seq>;
    fn into_iter(self) -> Self::IntoIter {
        ToStringIter {
            marker: PhantomData,
            progressor: string_iter::Progressor::from(self.0),
        }
    }
}

pub struct ToStringIter<'f, Ctx, Seq> {
    marker: PhantomData<(Ctx, Seq)>,
    progressor: string_iter::Progressor<'f>,
}

impl<'f, Ctx, Seq> Iterator for ToStringIter<'f, Ctx, Seq>
where
    Ctx: Context,
    Seq: WritableSeq<InMemoryOutput<Ctx>>,
{
    type Item = String;
    fn next(&mut self) -> Option<Self::Item> {
        self.progressor.next()
    }
}

mod string_iter {
    use crate::util::InMemoryOutput;
    use crate::{Output, SequenceAccept, Writable, WritableSeq};
    use std::cell::UnsafeCell;
    use std::convert::Infallible;
    use std::future::poll_fn;
    use std::marker::PhantomData;
    use std::pin::Pin;
    use std::ptr::NonNull;
    use std::task::{Context, Poll, Waker};
    use std::{mem, ptr};

    struct AdaptSeq<O, Seq>(PhantomData<O>, Seq);

    impl<O, Seq> AdaptSeq<O, Seq>
    where
        O: Output<Error = Infallible>,
        Seq: WritableSeq<O>,
    {
        fn make_future<S>(self, sink: S) -> impl Future<Output = ()> + use<O, Seq, S>
        where
            S: SequenceAccept<O>,
        {
            async move {
                let seq = self.1;
                let mut sink = sink;
                let res = WritableSeq::for_each(&seq, &mut sink).await;
                match res {
                    Ok(()) => (),
                    Err(e) => match e {},
                }
            }
        }
    }

    pub struct IntoStringIterStart<Ctx, Seq> {
        pub context: Ctx,
        pub sequence: Seq,
    }

    pub struct Progressor<'f> {
        // We use a manual pointer to share data in multiple places
        buffer: NonNull<ItemBuffer>,
        future: Pin<Box<dyn Future<Output = ()> + 'f>>,
        finished: bool,
    }

    impl<'f> Drop for Progressor<'f> {
        fn drop(&mut self) {
            unsafe {
                // SAFETY
                // We are not dropping the shared data anywhere else
                ptr::drop_in_place(self.buffer.as_ptr());
            }
        }
    }

    impl<'f, Ctx, Seq> From<IntoStringIterStart<Ctx, Seq>> for Progressor<'f>
    where
        Ctx: super::Context + 'f,
        Seq: WritableSeq<InMemoryOutput<Ctx>> + 'f,
    {
        fn from(value: IntoStringIterStart<Ctx, Seq>) -> Self {
            let buffer = Box::new(ItemBuffer::default());
            let buffer = unsafe {
                // SAFETY -- one of Rust's API mistakes was not returning NonNull
                NonNull::new_unchecked(Box::into_raw(buffer))
            };
            let adapt_seq = AdaptSeq(PhantomData, value.sequence);
            let seq_accept = SeqAccept {
                context: Some(value.context),
                buffer,
            };
            let future = adapt_seq.make_future(seq_accept);
            Self {
                buffer,
                future: Box::pin(future),
                finished: false,
            }
        }
    }

    impl<'f> Iterator for Progressor<'f> {
        type Item = String;
        fn next(&mut self) -> Option<Self::Item> {
            if self.finished {
                return None;
            }
            let poll = self
                .future
                .as_mut()
                .poll(&mut Context::from_waker(Waker::noop()));

            let buffer = unsafe {
                // SAFETY
                // We touch this buffer in one other place, inside the future
                &*self.buffer.as_ptr()
            };
            let extract = buffer.extract();

            match poll {
                Poll::Pending => {
                    // The future stalling tells us that there is an item in the buffer
                    // So, this should always be Some
                    assert!(
                        extract.is_some(),
                        "Extraneous async computations (writable should complete regularly)"
                    );
                    extract
                }
                Poll::Ready(()) => {
                    // All done!
                    self.finished = true;
                    // Can be None if we are fully empty
                    extract
                }
            }
        }
    }

    struct ItemBuffer {
        current: UnsafeCell<Option<String>>,
    }

    impl Default for ItemBuffer {
        fn default() -> Self {
            Self {
                current: UnsafeCell::new(None),
            }
        }
    }

    impl ItemBuffer {
        fn has_space(&self) -> bool {
            unsafe {
                // SAFETY
                // We only read and don't modify from anywhere else
                let ptr = self.current.get();
                (&*ptr).is_none()
            }
        }

        fn set_new(&self, value: String) {
            unsafe {
                // SAFETY
                // We don't modify from anywhere else
                let ptr = self.current.get();
                *ptr = Some(value);
            }
        }

        fn extract(&self) -> Option<String> {
            unsafe {
                // SAFETY
                // We don't modify from anywhere else
                let ptr = self.current.get();
                mem::replace(&mut *ptr, None)
            }
        }
    }

    struct SeqAccept<Ctx> {
        // This is always Some() except during a panic
        context: Option<Ctx>,
        buffer: NonNull<ItemBuffer>,
    }

    impl<Ctx> SequenceAccept<InMemoryOutput<Ctx>> for SeqAccept<Ctx>
    where
        Ctx: super::Context,
    {
        async fn accept<W>(&mut self, writable: &W) -> Result<(), Infallible>
        where
            W: Writable<InMemoryOutput<Ctx>>,
        {
            poll_fn(|_| {
                let buffer = unsafe {
                    // SAFETY
                    // We touch this buffer in one other place, outside the future
                    &*self.buffer.as_ptr()
                };
                if !buffer.has_space() {
                    return Poll::Pending;
                }
                unsafe {
                    // SAFETY
                    // We take the Ctx, then put it back, without modifying it
                    let ctx_store = &mut self.context;
                    let context = mem::take(ctx_store).unwrap_unchecked();
                    // Panic safety: upon panic, Ctx is dropped and the None prevents double-free
                    let InMemoryOutput {
                        buf: string,
                        context,
                    } = InMemoryOutput::print_output_impl(context, writable);
                    *ctx_store = Some(context);

                    // Then, we put the string into the storage
                    buffer.set_new(string);
                }
                Poll::Ready(Ok(()))
            })
            .await
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::common::{NoOpSeq, StrArrSeq};
    use crate::context::EmptyContext;
    use crate::util::IntoStringIter;

    #[test]
    fn sequence_iterator() {
        let sequence = StrArrSeq(&["One", "Two", "Three"]);
        let iterator = IntoStringIter::new(EmptyContext, sequence);
        let iterator = iterator.into_iter();
        let expected = &["One", "Two", "Three"].map(String::from);
        assert_eq!(iterator.collect::<Vec<_>>(), Vec::from(expected));
    }

    #[test]
    fn sequence_iterator_empty() {
        let sequence = NoOpSeq;
        let iterator = IntoStringIter::new(EmptyContext, sequence);
        let iterator = iterator.into_iter();
        assert!(iterator.collect::<Vec<_>>().is_empty());
    }
}