cargo-local 0.4.0

use cargo::core::Workspace as CargoWorkspace;
use cargo::ops::load_pkg_lockfile as load_cargo_lockfile;
use cargo::util::context::GlobalContext as CargoContext;
use cargo::util::CargoResult;
use cargo::util::cache_lock::CacheLockMode;
use clap::{Arg, ArgAction, Command, crate_version};

use cargo::core::global_cache_tracker::GlobalCacheTracker;

use std::collections::HashMap;
use std::env;
use std::path::{Path, PathBuf};
use std::process;

const DESCRIPTION: &str =
    "A third-party cargo extension that lists dependencies' source locations";

fn main() {
    let outer_matches = Command::new("cargo")
        .about(DESCRIPTION)
        // We have to lie about our binary name since this will be a third party
        // subcommand for cargo, this trick learned from cargo-outdated
        .bin_name("cargo")
        // We use a subcommand because parsed after `cargo` is sent to the third party plugin
        // which will be interpreted as a subcommand/positional arg by clap
        .subcommand(Command::new("local").about(DESCRIPTION)
                    .version(crate_version!())
                    .arg(Arg::new("PACKAGE")
                         .help("Individual packages to show the source locations of")
                         .action(ArgAction::Append)
                         .required(false))
                    .arg(Arg::new("only-names")
                         .short('n')
                         .long("only-names")
                         .help("Only list package names")
                         .action(ArgAction::SetTrue)))
        .subcommand_required(true)
        .get_matches();
    let arg_matches = outer_matches.subcommand_matches("local").unwrap();

    let src_dirs = match cargo_dirs() {
        Ok(Some(dirs)) => dirs,
        Ok(None) => {
            eprintln!("Error: Couldn't detect Cargo project in the current directory");
            process::exit(1);
        },
        Err(e) => {
            eprintln!("Error: {}", e);
            process::exit(1);
        },
    };

    if let Some(package_names) = arg_matches.get_many::<String>("PACKAGE") {
        for package_name in package_names {
            match src_dirs.get(package_name) {
                Some(dir) => {
                    if arg_matches.get_flag("only-names") {
                        println!("{}", package_name)
                    } else {
                        println!("{}", dir.display())
                    }
                },
                None => eprintln!("Warning: Couldn't find local dir for package: {}", package_name)
            }
        }
    } else {
        for (package_name, dir) in src_dirs {
            if arg_matches.get_flag("only-names") {
                println!("{}", package_name);
            } else {
                println!("{}", dir.display());
            }
        }
    }
}

fn cargo_dirs() -> CargoResult<Option<HashMap<String, PathBuf>>> {
    // Load the current project's dependencies from its Cargo manifest
    let manifest_path     = "Cargo.toml";
    let manifest_path     = Path::new(&manifest_path);
    let manifest_path_buf = absolutize(manifest_path.to_path_buf());
    let manifest_path     = manifest_path_buf.as_path();

    // The global context that is aware of where things are on the filesystem:
    let cargo_context        = CargoContext::default().expect("cargo_context");
    let _lock                = cargo_context.acquire_package_cache_lock(CacheLockMode::DownloadExclusive)?;
    let registry_source_path = cargo_context.registry_source_path().into_path_unlocked().to_owned();

    // Information about where source code is stored is located in an sqlite cache tracker,
    // generate a full map:
    let cache_tracker = GlobalCacheTracker::new(&cargo_context).unwrap();
    let package_src_map = cache_tracker.
        registry_src_all()?.
        into_iter().
        map(|(registry_src, _)| {
            let key = registry_src.package_dir.to_string();
            let value = registry_source_path.join(registry_src.encoded_registry_name.as_str()).join(&key);

            (key, value)
        }).
        collect::<HashMap<String, PathBuf>>();

    // Local package information is encapsulated in a Workspace:
    let workspace = CargoWorkspace::new(manifest_path, &cargo_context)?;
    let resolved = match load_cargo_lockfile(&workspace)? {
        Some(r) => r,
        None => return Ok(None),
    };

    let paths = resolved.iter().flat_map(|pkgid| {
        let package_key = format!("{}-{}", pkgid.name(), pkgid.version());

        let Some(path) = package_src_map.get(&package_key) else {
            return None;
        };

        if path.exists() {
            Some((pkgid.name().to_string(), path.clone()))
        } else {
            None
        }
    }).collect();

    Ok(Some(paths))
}

fn absolutize(pb: PathBuf) -> PathBuf {
    if pb.as_path().is_absolute() {
        pb
    } else {
        std::env::current_dir().expect("current_dir").join(&pb.as_path())
    }
}