comlexr/macro/src/pipe.rs

use quote::{quote, ToTokens};
use syn::{
    parse::{discouraged::Speculative, Parse},
    punctuated::Punctuated,
    spanned::Spanned,
    Error, Token,
};

use crate::cmd::Value;

pub struct Pipe {
    stdin: Option<Value>,
    commands: Punctuated<CommandExpr, Token![|]>,
}

impl Parse for Pipe {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let stdin = match input.cursor().ident() {
            Some((ident, _)) if ident == "stdin" && input.peek2(Token![=]) => {
                _ = input.parse::<syn::Ident>()?;
                _ = input.parse::<Token![=]>()?;
                let stdin = input.parse()?;
                _ = input.parse::<Token![;]>()?;
                Some(stdin)
            }
            _ => None,
        };
        let commands = Punctuated::parse_separated_nonempty(input)?;

        if commands.is_empty() {
            return Err(Error::new_spanned(
                commands,
                "At least one command is required",
            ));
        }

        Ok(Self { stdin, commands })
    }
}

impl ToTokens for Pipe {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let Self { stdin, commands } = self;
        let command_count = commands.len();
        let last_command_span = commands.last().unwrap().span();
        let mut commands_iter = commands.iter();
        let first_command = commands_iter.next().unwrap();

        let initial = quote! {
            let mut _c_0: #first_command;
            _c_0.stdout(::std::process::Stdio::piped());
            _c_0.stdin(::std::process::Stdio::piped());
            let mut _child_0 = _c_0.spawn()?;

            if let Some(stdin) = stdin {
                _child_0
                    .stdin
                    .as_mut()
                    .ok_or(::comlexr::ExecutorError::NoStdIn)?
                    .write_all(stdin.as_ref())?;
            }
        };

        let stdin = stdin.as_ref().map_or_else(
            || {
                quote! {
                    None::<&[u8]>
                }
            },
            |stdin| {
                quote! {
                    Some(#stdin)
                }
            },
        );

        let commands = commands_iter.enumerate().map(|(index, command)| {
            let previous_span = if index == 0 {
                first_command.span()
            } else {
                commands[index - 1].span()
            };
            let prev_com_ident = syn::Ident::new(&format!("_c_{index}"), previous_span);
            let prev_child_ident = syn::Ident::new(&format!("_child_{index}"), previous_span);
            let com_ident = syn::Ident::new(&format!("_c_{}", index + 1), command.span());
            let child_ident = syn::Ident::new(&format!("_child_{}", index + 1), command.span());

            quote! {
                let _output = #prev_child_ident.wait_with_output()?;

                if !_output.status.success() {
                    return Err(::comlexr::ExecutorError::FailedCommand{
                        command: std::boxed::Box::new(#prev_com_ident),
                        exit_code: _output.status.code().unwrap_or(1),
                    });
                }

                let mut #com_ident: #command;
                #com_ident.stdout(::std::process::Stdio::piped());
                #com_ident.stdin(::std::process::Stdio::piped());

                let mut #child_ident = #com_ident.spawn()?;
                #child_ident
                    .stdin
                    .as_mut()
                    .ok_or(::comlexr::ExecutorError::NoStdIn)?
                    .write_all(&_output.stdout)?;
            }
        });

        let last_child_ident =
            syn::Ident::new(&format!("_child_{}", command_count - 1), last_command_span);

        tokens.extend(quote! {
            ::comlexr::Executor::new(
                #stdin,
                |stdin| -> ::std::result::Result<
                    ::std::process::Child,
                    ::comlexr::ExecutorError,
                > {
                    use ::std::io::Write;
                    #initial
                    #(#commands)*
                    Ok(#last_child_ident)
                }
            )
        });
    }
}

enum CommandExpr {
    Macro(syn::ExprMacro),
    Reference(syn::ExprReference),
    Function(syn::ExprCall),
    Block(syn::ExprBlock),
}

impl Parse for CommandExpr {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let fork = input.fork();

        fork.parse()
            .map(|value| {
                input.advance_to(&fork);
                Self::Macro(value)
            })
            .or_else(|_| {
                let fork = input.fork();
                let value = fork.parse()?;
                input.advance_to(&fork);
                Ok(Self::Function(value))
            })
            .or_else(|_: syn::Error| {
                let fork = input.fork();
                let value = fork.parse()?;
                input.advance_to(&fork);
                Ok(Self::Block(value))
            })
            .or_else(|_: syn::Error| {
                let fork = input.fork();
                let value = fork.parse()?;
                input.advance_to(&fork);
                Ok(Self::Reference(value))
            })
            .map_err(|_: syn::Error| {
                syn::Error::new(
                    input.span(),
                    "Only references, function calls, macro calls, and blocks are allowed",
                )
            })
    }
}

impl ToTokens for CommandExpr {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        tokens.extend(match self {
            Self::Macro(macr) => quote! {
                ::std::process::Command = #macr
            },
            Self::Reference(refer) => quote! {
                &mut ::std::process::Command = #refer
            },
            Self::Function(fun) => quote! {
                ::std::process::Command = #fun
            },
            Self::Block(block) => quote! {
                ::std::process::Command = #block
            },
        });
    }
}