async_codegen/

common.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.
 */

//!
//!
//! Provides common, re-usable types that work for any code generation purposes
//!
//! Some of the types in this module are conceptual and are related to the library itself, like
//! [NoOp] or [Combined]. Other types are syntax based, but their syntax is so general that
//! they are useful for code generation in different languages.
//!

use super::{Output, SequenceAccept, Writable, WritableSeq};

/// Empty writable
#[derive(Copy, Clone, Debug)]
pub struct NoOp;

impl<O> Writable<O> for NoOp
where
    O: Output,
{
    async fn write_to(&self, _: &mut O) -> Result<(), O::Error> {
        Ok(())
    }
}

/// Empty writable sequence
#[derive(Copy, Clone, Debug)]
pub struct NoOpSeq;

impl<O> WritableSeq<O> for NoOpSeq
where
    O: Output,
{
    async fn for_each<S: SequenceAccept<O>>(&self, _: &mut S) -> Result<(), O::Error> {
        Ok(())
    }
}

/// A combined writable that simply writes its tuple values out
#[derive(Clone, Debug)]
pub struct Combined<W1, W2>(pub W1, pub W2);

impl<O, W1, W2> Writable<O> for Combined<W1, W2>
where
    O: Output,
    W1: Writable<O>,
    W2: Writable<O>,
{
    async fn write_to(&self, output: &mut O) -> Result<(), O::Error> {
        self.0.write_to(output).await?;
        self.1.write_to(output).await
    }
}

/// A combined writable sequence that consists of two other sequences
#[derive(Clone, Debug)]
pub struct CombinedSeq<S1, S2>(pub S1, pub S2);

impl<O, S1, S2> WritableSeq<O> for CombinedSeq<S1, S2>
where
    O: Output,
    S1: WritableSeq<O>,
    S2: WritableSeq<O>,
{
    async fn for_each<S>(&self, sink: &mut S) -> Result<(), O::Error>
    where
        S: SequenceAccept<O>,
    {
        self.0.for_each(sink).await?;
        self.1.for_each(sink).await
    }
}

/// A singular converts a [`Writable`] into a single-element [`WritableSeq`]
/// Its implementation just writes the tuple value.
#[derive(Clone, Debug)]
pub struct SingularSeq<W>(pub W);

impl<O, W> WritableSeq<O> for SingularSeq<W>
where
    O: Output,
    W: Writable<O>,
{
    async fn for_each<S>(&self, sink: &mut S) -> Result<(), O::Error>
    where
        S: SequenceAccept<O>,
    {
        sink.accept(&self.0).await
    }
}

/// A struct that repeats a value multiple times in order to turn it into a [`WritableSeq`].
/// This struct is indeed a generalization of [`SingularSeq`] (it's equivalent when the count is 1).
#[derive(Clone, Debug)]
pub struct RepeatSeq<W>(pub usize, pub W);

impl<O, W> WritableSeq<O> for RepeatSeq<W>
where
    O: Output,
    W: Writable<O>,
{
    async fn for_each<S>(&self, sink: &mut S) -> Result<(), O::Error>
    where
        S: SequenceAccept<O>,
    {
        for _ in 0..self.0 {
            sink.accept(&self.1).await?;
        }
        Ok(())
    }
}

/// A sequence based on an array of writable elements
#[derive(Debug)]
pub struct ArrSeq<'a, W>(pub &'a [W]);

impl<'a, W> Clone for ArrSeq<'a, W> {
    fn clone(&self) -> Self {
        Self(self.0)
    }
}

impl<O, W> WritableSeq<O> for ArrSeq<'_, W>
where
    O: Output,
    W: Writable<O>,
{
    async fn for_each<S>(&self, sink: &mut S) -> Result<(), O::Error>
    where
        S: SequenceAccept<O>,
    {
        for elem in self.0 {
            sink.accept(elem).await?;
        }
        Ok(())
    }
}

/// Turns a string or a string reference into a writable value.
/// Can be used with [str], [String], or any type that implements `AsRef<str>`
#[derive(Copy, Clone, Debug)]
pub struct Str<STR>(pub STR);

impl<O, STR> Writable<O> for Str<STR>
where
    O: Output,
    STR: AsRef<str>,
{
    async fn write_to(&self, output: &mut O) -> Result<(), O::Error> {
        output.write(self.0.as_ref()).await
    }
}

/// Like [ArrSeq], but specialized for string so that you don't need to wrap each element in [Str]
#[derive(Clone, Debug)]
pub struct StrArrSeq<'a>(pub &'a [&'a str]);

impl<O> WritableSeq<O> for StrArrSeq<'_>
where
    O: Output,
{
    async fn for_each<S>(&self, sink: &mut S) -> Result<(), O::Error>
    where
        S: SequenceAccept<O>,
    {
        for elem in self.0 {
            sink.accept(&Str(*elem)).await?;
        }
        Ok(())
    }
}

/// A general purpose sequence acceptor that tracks whether any element is written and can
/// write a separator between each element, in addition to an optional intro at the beginning.
#[derive(Debug)]
pub struct SeparatedSeqAccept<'o, O, Intro, Sep> {
    /// The output.
    pub output: &'o mut O,
    /// The intro. Must be writable
    pub intro: Intro,
    /// The separator. Must be writable
    pub separator: Sep,
    /// Whether this instance wrote anything. Should be initially set to false.
    pub wrote_any: bool,
}

impl<'o, O, Intro, Sep> SeparatedSeqAccept<'o, O, Intro, Sep> {
    pub fn new(output: &'o mut O, intro: Intro, separator: Sep) -> Self {
        Self {
            output,
            intro,
            separator,
            wrote_any: false,
        }
    }
}

impl<'o, O, Sep> SeparatedSeqAccept<'o, O, NoOp, Sep> {
    pub fn no_intro(output: &'o mut O, separator: Sep) -> Self {
        Self::new(output, NoOp, separator)
    }
}

impl<'o, O> SeparatedSeqAccept<'o, O, NoOp, Str<&'static str>> {
    /// Makes a sink for sequences that separates by `, `.
    pub fn comma_separated(output: &'o mut O) -> Self {
        Self::new(output, NoOp, Str(", "))
    }
}

impl<'o, O, Intro, Sep> SequenceAccept<O> for SeparatedSeqAccept<'o, O, Intro, Sep>
where
    O: Output,
    Intro: Writable<O>,
    Sep: Writable<O>,
{
    async fn accept<W>(&mut self, writable: &W) -> Result<(), O::Error>
    where
        W: Writable<O>,
    {
        if self.wrote_any {
            self.separator.write_to(self.output).await?;
        } else {
            self.intro.write_to(self.output).await?;
            self.wrote_any = true;
        }
        writable.write_to(self.output).await
    }
}

/// A general purpose sequence acceptor that surrounds every element written with the "before"
/// and "after" writable elements.
#[derive(Debug)]
pub struct SurroundingSeqAccept<'o, O, Before, After> {
    pub output: &'o mut O,
    pub before: Before,
    pub after: After,
}

impl<'o, O, Before, After> SurroundingSeqAccept<'o, O, Before, After> {
    pub fn new(output: &'o mut O, before: Before, after: After) -> Self {
        Self {
            output,
            before,
            after,
        }
    }
}

impl<'o, O, Before, After> SequenceAccept<O> for SurroundingSeqAccept<'o, O, Before, After>
where
    O: Output,
    Before: Writable<O>,
    After: Writable<O>,
{
    async fn accept<W>(&mut self, writable: &W) -> Result<(), O::Error>
    where
        W: Writable<O>,
    {
        self.before.write_to(self.output).await?;
        writable.write_to(self.output).await?;
        self.after.write_to(self.output).await
    }
}

/// A sequence implementation that takes elements in an array, and writes each element using
/// the config. The config should implement [SequenceConfig].
#[derive(Debug)]
pub struct SequenceViaConfig<'t, T, Config> {
    pub data: &'t [T],
    pub config: Config,
}

impl<'t, T, Config> Clone for SequenceViaConfig<'t, T, Config>
where
    Config: Clone,
{
    fn clone(&self) -> Self {
        Self {
            data: self.data,
            config: self.config.clone(),
        }
    }
}

/// A helper object that writes other data, and not the object itself, on demand
pub trait SequenceConfig<T, O>
where
    O: Output,
{
    /// Writes an individual data point.
    async fn write_datum(&self, datum: &T, output: &mut O) -> Result<(), O::Error>;
}

impl<'t, T, Config, O> WritableSeq<O> for SequenceViaConfig<'t, T, Config>
where
    O: Output,
    Config: SequenceConfig<T, O>,
{
    async fn for_each<S>(&self, sink: &mut S) -> Result<(), O::Error>
    where
        S: SequenceAccept<O>,
    {
        struct WriteSingle<'t, T, Config> {
            datum: &'t T,
            config: &'t Config,
        }
        impl<'t, T, Config, O> Writable<O> for WriteSingle<'t, T, Config>
        where
            O: Output,
            Config: SequenceConfig<T, O>,
        {
            async fn write_to(&self, output: &mut O) -> Result<(), O::Error> {
                self.config.write_datum(self.datum, output).await
            }
        }
        for datum in self.data {
            sink.accept(&WriteSingle {
                datum,
                config: &self.config,
            })
            .await?;
        }
        Ok(())
    }
}

/// Runs the given writable or sequence, but places the future inside [Box::pin].
/// This layer of indirection helps to break recursion issues in async code.
///
/// (This struct can be used as either a writable or a sequence, depending on its contents)
pub struct BoxedIndirection<W>(pub W);

impl<O, W> Writable<O> for BoxedIndirection<W>
where
    O: Output,
    W: Writable<O>,
{
    async fn write_to(&self, output: &mut O) -> Result<(), O::Error> {
        Box::pin(self.0.write_to(output)).await
    }
}

impl<O, W> WritableSeq<O> for BoxedIndirection<W>
where
    O: Output,
    W: WritableSeq<O>,
{
    async fn for_each<S>(&self, sink: &mut S) -> Result<(), O::Error>
    where
        S: SequenceAccept<O>,
    {
        Box::pin(self.0.for_each(sink)).await
    }
}