v_escape_codegen 0.1.9

#![doc = include_str!("../README.md")]
use std::{
    collections::BTreeMap,
    fmt::Debug,
    fs,
    path::{Path, PathBuf},
    str,
};

use proc_macro2::{Ident, Span, TokenStream};
use serde::Serialize;
use tests::build_tests;
use toml::Value;

use clap::Parser;

use v_escape_codegen_base::generate as generate_base;
mod tests;

fn ident(s: &str) -> Ident {
    Ident::new(s, Span::call_site())
}

/// V_escape codegen - A tool for generating escape functions
#[derive(Parser, Debug)]
#[clap(
    author,
    version,
    about = "Generate escape functions from template files",
    long_about = "A tool for generating SIMD-optimized escape functions from template files. 
    
Creates a new crate with escape_fmt and escape_string functions based on character mappings defined in src/_lib.rs.

Example usage:
  mkdir my_escape
  cd my_escape
  cargo init --lib
  cat <<EOF > src/_lib.rs
  new!(
      '&' -> \"&amp;\",
      '/' -> \"&#x2f;\",
      '<' -> \"&lt;\",
      '>' -> \"&gt;\",
      '\"' -> \"&quot;\",
      '\\'' -> \"&#x27;\"
  );
  EOF
  v_escape-codegen -i .",
    after_help = "For more information, see: https://github.com/zzau13/v_escape"
)]
struct Args {
    /// Input directory containing the crate to generate
    #[clap(short, long, default_value = "./", value_name = "DIR")]
    pub input_dir: PathBuf,
}

#[derive(Serialize)]
struct Dep {
    workspace: bool,
}

fn generate(dir: impl AsRef<Path>) -> anyhow::Result<()> {
    _generate(dir.as_ref())
}

fn read_cargo(p: &Path) -> anyhow::Result<(Value, String)> {
    let cargo_src =
        fs::read_to_string(p).map_err(|e| anyhow::anyhow!("Failed to read Cargo.toml: {}", e))?;
    let mut cargo_value = cargo_src
        .parse::<Value>()
        .map_err(|e| anyhow::anyhow!("Failed to parse Cargo.toml: {}", e))?;

    let doc: Value = toml::from_str(
        r"
    [docs.rs]
    all-features = true
    ",
    )
    .map_err(|e| anyhow::anyhow!("Failed to parse TOML: {}", e))?;

    let cargo_mut = cargo_value
        .as_table_mut()
        .ok_or_else(|| anyhow::anyhow!("Expected a table for cargo_value"))?;

    cargo_mut
        .get_mut("package")
        .ok_or_else(|| anyhow::anyhow!("Expected a package section in Cargo.toml"))?
        .as_table_mut()
        .ok_or_else(|| anyhow::anyhow!("Expected a table for package"))?
        .insert("metadata".to_string(), doc);

    let mut features = BTreeMap::new();
    features.insert("default", vec!["std", "string", "fmt", "bytes"]);
    features.insert("std", vec!["v_escape-base/std", "alloc"]);
    features.insert("alloc", vec!["v_escape-base/alloc"]);
    features.insert("string", vec!["v_escape-base/string"]);
    features.insert("fmt", vec!["v_escape-base/fmt"]);
    features.insert("bytes", vec!["v_escape-base/bytes"]);

    cargo_mut.insert("features".into(), Value::from(features));

    if !cargo_mut.contains_key("dependencies") {
        cargo_mut.insert(
            "dependencies".into(),
            Value::from(BTreeMap::<String, Value>::new()),
        );
    }
    let dependencies = cargo_mut
        .get_mut("dependencies")
        .ok_or_else(|| anyhow::anyhow!("Expected a table for dependencies"))?;
    dependencies
        .as_table_mut()
        .ok_or_else(|| anyhow::anyhow!("Expected a table for dependencies"))?
        .insert(
            "v_escape-base".into(),
            Value::try_from(Dep { workspace: true })?,
        );

    let package_name = cargo_value
        .as_table()
        .ok_or_else(|| anyhow::anyhow!("Expected a table for cargo_value"))?
        .get("package")
        .ok_or_else(|| anyhow::anyhow!("Expected a package section in Cargo.toml"))?
        .as_table()
        .ok_or_else(|| anyhow::anyhow!("Expected a table for package"))?
        .get("name")
        .ok_or_else(|| anyhow::anyhow!("Expected a name in package section"))?
        .as_str()
        .ok_or_else(|| anyhow::anyhow!("Expected a name as str"))?;

    let package_name = package_name.to_string();
    Ok((cargo_value, package_name))
}

fn _generate(dir: &Path) -> anyhow::Result<()> {
    let head =
        "//! autogenerated by v_escape_codegen@".to_string() + env!("CARGO_PKG_VERSION") + "\n";

    if !dir.is_dir() {
        anyhow::bail!("input_dir should be a directory");
    }

    // Modify Cargo.toml
    // TODO: should use a pretty toml
    let cargo = dir.join("Cargo.toml");
    let (cargo_value, name) = read_cargo(&cargo)?;

    // Check directories
    let src = dir.join("src");
    let test = dir.join("tests");
    if !test.exists() {
        fs::create_dir(&test)?;
    }
    // Read template
    let template = src.join("_lib.rs");
    let template_src = fs::read_to_string(&template)?;

    // Generate code
    let (code, (escapes, escaped)) = generate_base(
        template_src
            .parse::<TokenStream>()
            .map_err(|e| anyhow::anyhow!("Failed to parse template source: {}", e))?,
        "v_escape_base",
    )?;

    // Prettify code
    let code_pretty = prettyplease::unparse(
        &syn::parse2(code)
            .map_err(|e| anyhow::anyhow!("Failed to parse code to TokenStream: {}", e))?,
    );
    // Generate tests
    let code_test = build_tests(&ident(&name), &escapes, &escaped);
    let code_test_pretty = prettyplease::unparse(
        &syn::parse2(code_test)
            .map_err(|e| anyhow::anyhow!("Failed to parse code to TokenStream: {}", e))?,
    );

    // Write files
    fs::write(&cargo, toml::to_string_pretty(&cargo_value)?)?;
    fs::write(src.join("lib.rs"), head.clone() + &code_pretty)?;
    fs::write(test.join("lib.rs"), head.clone() + &code_test_pretty)?;

    Ok(())
}

fn main() -> anyhow::Result<()> {
    let args = Args::parse();
    let dir: PathBuf = args.input_dir;

    generate(dir)
}