cargo-pgo 0.3.0

use std::path::PathBuf;

use cargo_pgo::bolt::instrument::{BoltInstrumentArgs, bolt_instrument};
use cargo_pgo::bolt::optimize::{BoltOptimizeArgs, bolt_optimize};
use cargo_pgo::build::CargoCommand;
use cargo_pgo::check::environment_info;
use cargo_pgo::clean::clean_artifacts;
use cargo_pgo::get_cargo_ctx;
use cargo_pgo::pgo::instrument::{PgoInstrumentArgs, PgoInstrumentShortcutArgs, pgo_instrument};
use cargo_pgo::pgo::optimize::{PgoOptimizeArgs, pgo_optimize};
use clap::Parser;
use env_logger::Env;

#[derive(clap::Parser, Debug)]
#[clap(author, version, about)]
#[clap(bin_name("cargo"))]
#[clap(disable_help_subcommand(true))]
enum Args {
    #[clap(subcommand)]
    #[clap(author, version, about)]
    Pgo(Subcommand),
}

#[derive(clap::Subcommand, Debug)]
enum Subcommand {
    /// Display information about your environment. Can be used to check whether it is prepared for
    /// PGO and BOLT.
    Info,
    /// Execute a `cargo` command to create PGO-instrumented artifact(s).
    /// After the artifacts are executed, they will produce profiles that can be later used in the
    /// `optimize` step.
    Instrument(PgoInstrumentArgs),
    /// Execute `cargo build` to create a PGO-instrumented binary. When executed, the binary will produce
    /// profiles that can be later used in the `optimize` step.
    Build(PgoInstrumentShortcutArgs),
    /// Execute `cargo test` to produce PGO profiles from test execution, which can be later used
    /// in the `optimize` step.
    Test(PgoInstrumentShortcutArgs),
    /// Execute `cargo run` to produce PGO profiles from binary execution, which can be later used
    /// in the `optimize` step.
    Run(PgoInstrumentShortcutArgs),
    /// Execute `cargo bench` to produce PGO profiles from benchmark execution, which can be later
    /// used in the `optimize` step.
    Bench(PgoInstrumentShortcutArgs),
    /// Build an optimized version of a binary using generated PGO profiles.
    Optimize(PgoOptimizeArgs),
    /// Optimization using BOLT.
    #[clap(subcommand)]
    Bolt(BoltArgs),
    /// Clean PGO and BOLT artifacts from the disk.
    Clean(CleanArgs),
}

#[derive(clap::Subcommand, Debug)]
enum BoltArgs {
    /// Run `cargo build` with instrumentation to prepare for BOLT optimization.
    Build(BoltInstrumentArgs),
    /// Built an optimized version of a binary using generated BOLT profiles.
    Optimize(BoltOptimizeArgs),
}

impl BoltArgs {
    fn cargo_args(&self) -> &[String] {
        match self {
            BoltArgs::Build(args) => args.cargo_args(),
            BoltArgs::Optimize(args) => args.cargo_args(),
        }
    }
}

#[derive(clap::Parser, Debug)]
struct CleanArgs {
    /// Override the PGO profile directory.
    #[clap(long)]
    profiles_dir: Option<PathBuf>,
}

impl Args {
    fn cargo_args(&self) -> &[String] {
        match self {
            Args::Pgo(args) => match args {
                Subcommand::Info => &[],
                Subcommand::Instrument(args) => args.cargo_args(),
                Subcommand::Build(args)
                | Subcommand::Run(args)
                | Subcommand::Test(args)
                | Subcommand::Bench(args) => args.cargo_args(),
                Subcommand::Optimize(args) => args.cargo_args(),
                Subcommand::Bolt(args) => args.cargo_args(),
                Subcommand::Clean(..) => &[],
            },
        }
    }

    fn profiles_dir(&self) -> Option<PathBuf> {
        match self {
            Args::Pgo(args) => match args {
                Subcommand::Info => None,
                Subcommand::Instrument(args) => args.profiles_dir().to_owned(),
                Subcommand::Build(args)
                | Subcommand::Run(args)
                | Subcommand::Test(args)
                | Subcommand::Bench(args) => args.profiles_dir().to_owned(),
                Subcommand::Optimize(args) => args.profiles_dir().to_owned(),
                Subcommand::Bolt(..) => None,
                Subcommand::Clean(CleanArgs { profiles_dir }) => profiles_dir.to_owned(),
            },
        }
    }
}

fn run() -> anyhow::Result<()> {
    let args = Args::parse();

    let cargo_args = args.cargo_args();
    let profiles_dir = args.profiles_dir();
    let ctx = get_cargo_ctx(cargo_args, profiles_dir)?;

    let Args::Pgo(args) = args;
    match args {
        Subcommand::Info => environment_info(),
        Subcommand::Instrument(args) => pgo_instrument(ctx, args),
        Subcommand::Build(args) => pgo_instrument(ctx, args.into_full_args(CargoCommand::Build)),
        Subcommand::Test(args) => pgo_instrument(ctx, args.into_full_args(CargoCommand::Test)),
        Subcommand::Run(args) => pgo_instrument(ctx, args.into_full_args(CargoCommand::Run)),
        Subcommand::Bench(args) => pgo_instrument(ctx, args.into_full_args(CargoCommand::Bench)),
        Subcommand::Optimize(args) => pgo_optimize(ctx, args),
        Subcommand::Bolt(BoltArgs::Build(args)) => bolt_instrument(ctx, args),
        Subcommand::Bolt(BoltArgs::Optimize(args)) => bolt_optimize(ctx, args),
        Subcommand::Clean(..) => clean_artifacts(ctx),
    }
}

fn main() {
    env_logger::Builder::from_env(Env::default().default_filter_or("cargo_pgo=info")).init();

    if let Err(error) = run() {
        eprintln!("{}", format!("{error:?}").trim_end_matches('\n'));
        std::process::exit(1);
    }
}