format-bytes-macros 0.4.0

// Copyright 2020, Raphaël Gomès <rgomes@octobus.net>

use proc_macro2::TokenStream;
use quote::quote;
use syn::{
    parse::{Parse, ParseStream},
    Token,
};

/// Defines the arguments to the `format_bytes!` macro.
#[derive(Debug)]
pub(crate) struct WriteBytesArgs {
    pub(crate) output: syn::Expr,
    pub(crate) format_string: syn::LitByteStr,
    pub(crate) positional_args: Vec<syn::Expr>,
    // TODO named arguments
}

impl Parse for WriteBytesArgs {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let output = input.parse()?;
        input.parse::<Token![,]>()?;
        let format_string = input.parse()?;
        let mut positional_args = vec![];

        while !input.is_empty() {
            input.parse::<Token![,]>()?;
            if input.is_empty() {
                break;
            }

            positional_args.push(input.parse()?);
        }

        Ok(Self {
            output,
            format_string,
            positional_args,
        })
    }
}

/// Returns the positions of the formatting brackets (or `{}` literally) in the
/// format bytestring.
pub(crate) fn extract_formatting_positions(
    input: &[u8],
) -> Result<Vec<usize>, ()> {
    let mut positions = vec![];

    let mut i = 0;
    while i < input.len() {
        if input[i] == b'{' {
            if i < input.len() - 1 {
                match input[i + 1] {
                    b'}' => {
                        positions.push(i);
                        i = i + 2;
                    }
                    b'{' => i = i + 2,
                    // Unclosed formatting delimiter
                    _ => return Err(()),
                }
            } else if i == input.len() - 1 {
                // Unclosed formatting delimiter
                return Err(());
            }
        } else {
            i = i + 1;
        }
    }
    Ok(positions)
}

/// Returns each call to `write_bytes` or `display_bytes`
pub(crate) fn generate_write_calls(
    args: &WriteBytesArgs,
    input: &[u8],
    positions: &[usize],
) -> Vec<TokenStream> {
    assert_eq!(positions.len(), args.positional_args.len());

    let mut calls = vec![];

    let mut i = 0;
    for (&position, arg) in positions.iter().zip(&args.positional_args) {
        if position > i {
            let literal =
                proc_macro2::Literal::byte_string(&input[i..position]);
            calls.push(quote! {
                ::std::io::Write::write_all(output, #literal)
            });
        }
        calls.push(quote! {
            format_bytes::DisplayBytes::display_bytes(&#arg, output)
        });
        i = position + 2; // skip the `{}`
    }

    // If there is still more format bytestring after the last replace
    if i < input.len() {
        let literal = proc_macro2::Literal::byte_string(&input[i..]);
        calls.push(quote! {
            ::std::io::Write::write_all(output, #literal)
        });
    }

    calls
}

/// Take tokens for any number of expressions that evaluate to `Result<(), E>`,
/// and return a single expression of the same type that evaluates them in order
/// until the first error.
pub(crate) fn combine_result_expressions(
    results: Vec<TokenStream>,
) -> TokenStream {
    let mut results = results.into_iter();
    let first = if let Some(first) = results.next() {
        first
    } else {
        // We are in `write_bytes!(output, "")`, so do nothing:
        return quote! {
            ::std::io::Result::Ok(())
        };
    };
    // Instead of chaining `Result::and_then` calls we could use the `?` operator
    // in a `try {}` block once those are stable https://github.com/rust-lang/rust/issues/31436
    // or in an immediately-called-closure.
    let mut combined = first;
    combined.extend(results.map(|result| {
        quote! {
            .and_then(|()| #result)
        }
    }));
    combined
}