thin-fetch 0.1.0

Low-level Git crate to fetch objects with high granularity
Documentation
use std::{fs::create_dir_all, io::Error, path::PathBuf};

use clap::{Arg, Command};
use thin_fetch::{HashKind, Reference, Repo, object::Object, utils::as_hex};

macro_rules! bail {
    ($msg:literal) => {{
        eprintln!($msg);
        std::process::exit(1);
    }};
    ($msg:literal, $($arg:tt)*) => {{
        eprintln!($msg, $($arg)*);
        std::process::exit(1);
    }}
}

fn mkdir(path: &PathBuf) {
    if let Err(err) = create_dir_all(path) {
        bail!("could not create {}: {err}", path.display());
    }
}

fn main() {
    let args = Command::new("shallow_clone")
        .arg(Arg::new("url").required(true).help("Repository URL"))
        .arg(Arg::new("output-dir").help("Output directory (default: .)"))
        .get_matches_from(std::env::args());

    let url = args.get_one::<String>("url").unwrap();
    let output = args
        .get_one::<String>("output-dir")
        .cloned()
        .unwrap_or(".".to_owned());

    let path = PathBuf::from(output);

    if path.exists() {
        bail!("path already exists!");
    }

    let local_repo = LocalRepo { path };

    let mut repo = Repo::new_discover(url).unwrap_or_else(|err| {
        bail!("could not discover repo: {err}");
    });

    let head = repo
        .get_ref("HEAD")
        .unwrap_or_else(|err| bail!("error getting refs: {err}"))
        .unwrap_or_else(|| bail!("head not found"));

    if let Err(err) = repo.fetch_shallow(head.target()) {
        bail!("error fetching: {err}");
    }

    let config = build_config(&head, url, repo.get_hash_kind());

    local_repo.write_config(&config).unwrap_or_else(|err| {
        bail!("could not write git config: {err}");
    });

    let branch_name = head
        .symref_target
        .as_deref()
        .map(|bref| bref.strip_prefix("refs/heads/").unwrap_or("main"));

    let head_file = if let Some(branch) = &head.symref_target {
        format!("ref: {branch}")
    } else {
        head.target().to_string()
    };

    local_repo.write_head(&head_file).unwrap_or_else(|err| {
        bail!("could not write git HEAD: {err}");
    });

    if let Some(name) = branch_name {
        local_repo
            .write_remote_ref("origin", name, head.target())
            .unwrap_or_else(|err| {
                bail!("could not write .git/refs/remotes/origin/{name}: {err}");
            });

        local_repo
            .write_local_head(name, head.target())
            .unwrap_or_else(|err| {
                bail!("could not write .git/refs/heads/{name}: {err}");
            });
    }

    for object in repo.objects() {
        local_repo.write_object(object);
    }
}

fn docstring(s: &str) -> String {
    let lines = s
        .trim_end()
        .split('\n')
        .filter(|l| !l.is_empty())
        .collect::<Vec<_>>();
    let indent = match lines.first() {
        Some(line) => {
            let idx = line.find(|c| !char::is_whitespace(c)).unwrap();
            &line[..idx]
        }
        None => "",
    };
    lines
        .into_iter()
        .map(|l| l.strip_prefix(indent).unwrap_or(l))
        .collect::<Vec<_>>()
        .join("\n")
}

fn build_config(head: &Reference, url: &str, hash_kind: HashKind) -> String {
    let mut config = docstring(&format!(
        r#"
        [core]
            bare = false
            repositoryformatversion = 1
            filemode = true
            logallrefupdates = true
        [remote "origin"]
            url = {url}
            fetch = +refs/heads/*:refs/remotes/origin/*
    "#
    ));

    if let Some(branch_ref) = &head.symref_target {
        let branch_name = branch_ref.strip_prefix("refs/heads/").unwrap_or("main");

        config += "\n";
        config += &docstring(&format!(
            r#"
            [branch "{branch_name}"]
                remote = origin
                merge = {branch_ref}
        "#
        ));
    }

    if hash_kind == HashKind::Sha256 {
        config += "\n";
        config += &docstring(
            r"
            [extensions]
	            objectformat = sha256
        ",
        );
    }

    config
}

struct LocalRepo {
    path: PathBuf,
}

impl LocalRepo {
    fn write_config(&self, config: &str) -> Result<(), Error> {
        let git_dir = self.path.join(".git");
        mkdir(&git_dir);

        let path = git_dir.join("config");
        std::fs::write(&path, config)
    }

    fn write_head(&self, head: &str) -> Result<(), Error> {
        let git_dir = self.path.join(".git");
        mkdir(&git_dir);

        let path = git_dir.join("HEAD");
        std::fs::write(&path, head)
    }

    fn write_local_head(&self, branch: &str, target: &str) -> Result<(), Error> {
        let dir = self.path.join(".git").join("refs").join("heads");
        mkdir(&dir);

        let file = dir.join(branch);
        std::fs::write(file, target)
    }

    fn write_remote_ref(&self, remote: &str, branch: &str, target: &str) -> Result<(), Error> {
        let dir = self
            .path
            .join(".git")
            .join("refs")
            .join("remotes")
            .join(remote);
        mkdir(&dir);

        let file = dir.join(branch);

        std::fs::write(file, target)
    }

    fn write_object(&self, object: &Object) {
        let dir = self.path.join(".git").join("objects");
        mkdir(&dir);

        let hash = as_hex(object.hash().as_ref());
        let (dir_name, file_name) = hash.split_at(2);

        let dir = dir.join(dir_name);
        mkdir(&dir);

        let file_path = dir.join(file_name);
        let mut file = std::fs::File::create(&file_path).unwrap_or_else(|err| {
            bail!("could not open object: {}: {err}", file_path.display());
        });

        let mut comp = flate2::write::ZlibEncoder::new(&mut file, flate2::Compression::best());
        if let Err(err) = object.write_to(&mut comp) {
            bail!("could not serialize object: {}: {err}", file_path.display());
        }

        if let Err(err) = comp.finish() {
            bail!("could not compress object: {err}");
        }
    }
}