swh-graph 12.1.2

/*
 * Copyright (C) 2024-2026  The Software Heritage developers
 * See the AUTHORS file at the top-level directory of this distribution
 * License: GNU General Public License version 3, or any later version
 * See top-level LICENSE file for more information
 */

use std::path::PathBuf;

use anyhow::{anyhow, bail, ensure, Context, Result};
use clap::{Parser, Subcommand};
use mmap_rs::Mmap;

use swh_graph::cli::MphAlgorithm;
use swh_graph::java_compat::mph::gov::GOVMPH;
use swh_graph::map::Node2SWHID;
use swh_graph::map::{MappedPermutation, Permutation};
#[cfg(feature = "pthash")]
use swh_graph::mph::SwhidPthash;
use swh_graph::mph::{LoadableSwhidMphf, SwhidFmphgo, SwhidMphf};
#[cfg(feature = "pthash")]
use swh_graph::person::PersonPthash;
use swh_graph::person::{LoadablePersonMphf, PersonFmphgo, PersonHasher};
use swh_graph::{OutOfBoundError, SWHID};

#[derive(Parser, Debug)]
/// Runs any of swh-graph's Minimal Perfect Hash functions
///
/// Lines in stdin are hashed one by one, and a decimal-encoded 64-bits integer is
/// written on the output for each of them.
///
/// If any of the input lines was not in the dataset used to build the MPH, then the result
/// will either be a silent hash collision or cause a non-zero exit.
struct Args {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand, Debug)]
enum Commands {
    Swhids {
        #[arg(long)]
        num_nodes: usize,
        #[arg(long)]
        mph_algo: MphAlgorithm,
        #[arg(long)]
        mph: PathBuf,
        #[arg(long)]
        permutation: PathBuf,
        #[arg(long)]
        /// If given, uses a `.node2swhid.bin` table to check for hash collisions.
        /// This turns all silent hash collisions into a hard error.
        node2swhid: Option<PathBuf>,
    },
    /// Hash pseudonymized persons (ie. sha256(fullname)), base64-encoded
    Persons {
        #[arg(long)]
        mph_algo: MphAlgorithm,
        #[arg(long)]
        /// The 2024-08-23 graph used the Labels MPH algo to hash persons; this works around this
        /// bug.
        workaround_2024_08_23: bool,
        #[arg(long)]
        mph: PathBuf,
    },
    /// Hash persons' fullnames
    PersonFullnames {
        #[arg(long)]
        mph_algo: MphAlgorithm,
        #[arg(long)]
        /// Whether the fullnames are given as a base64-encoded string.
        ///
        /// This is recommended for batch scripts, as not all fullnames are UTF8.
        base64: bool,
        #[arg(long)]
        mph: PathBuf,
    },
}

pub fn main() -> Result<()> {
    let args = Args::parse();

    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();

    match args.command {
        Commands::Swhids {
            num_nodes,
            mph_algo,
            mph,
            permutation,
            node2swhid,
        } => {
            log::info!("Loading permutation...");
            let permutation = MappedPermutation::load(num_nodes, permutation.as_path())
                .with_context(|| format!("Could not load permutation {}", permutation.display()))?;

            log::info!("Loading node2swhid...");
            let node2swhid = node2swhid
                .as_ref()
                .map(Node2SWHID::<Mmap>::load)
                .transpose()
                .with_context(|| {
                    format!(
                        "Could not load node2swhid from {}",
                        node2swhid.unwrap().display()
                    )
                })?;

            match mph_algo {
                MphAlgorithm::Fmphgo => hash_swhids::<SwhidFmphgo>(mph, permutation, node2swhid),
                MphAlgorithm::Pthash => {
                    #[cfg(not(feature = "pthash"))]
                    bail!(
                        "Cannot load MPHF {} because pthash support is disabled. Recompile swh-graph with --features pthash.",
                        mph.display()
                    );
                    #[cfg(feature = "pthash")]
                    hash_swhids::<SwhidPthash>(mph, permutation, node2swhid)
                }
                MphAlgorithm::Cmph => hash_swhids::<GOVMPH>(mph, permutation, node2swhid),
            }
        }
        Commands::Persons {
            mph_algo,
            mph,
            workaround_2024_08_23,
        } => match mph_algo {
            MphAlgorithm::Fmphgo => {
                ensure!(
                    !workaround_2024_08_23,
                    "--workaround-2024-08-23 is only meant for pthash."
                );
                hash_pseudonymized_persons::<PersonFmphgo>(mph)
            }
            MphAlgorithm::Pthash => {
                #[cfg(not(feature = "pthash"))]
                bail!("pthash is not supported. Recompile with --features phtash");
                #[cfg(feature = "pthash")]
                if workaround_2024_08_23 {
                    hash_pseudonymized_persons_pthash_2024_08_23(mph)
                } else {
                    hash_pseudonymized_persons::<PersonPthash>(mph)
                }
            }
            MphAlgorithm::Cmph => {
                ensure!(
                    !workaround_2024_08_23,
                    "--workaround-2024-08-23 is only meant for pthash."
                );
                hash_pseudonymized_persons_cmph(mph)
            }
        },
        Commands::PersonFullnames {
            mph_algo,
            base64,
            mph,
        } => match mph_algo {
            MphAlgorithm::Fmphgo => hash_person_fullnames::<PersonFmphgo>(mph, base64),
            MphAlgorithm::Pthash => {
                #[cfg(not(feature = "pthash"))]
                bail!("pthash is not supported. Recompile with --features phtash");
                #[cfg(feature = "pthash")]
                hash_person_fullnames::<PersonPthash>(mph, base64)
            }
            MphAlgorithm::Cmph => {
                bail!("'--mph-algo cmph' is not supported for non-pseudonymized graphs")
            }
        },
    }
}

fn hash_swhids<MPHF: LoadableSwhidMphf>(
    mph: PathBuf,
    permutation: MappedPermutation,
    node2swhid: Option<Node2SWHID<Mmap>>,
) -> Result<()> {
    log::info!("Loading MPH function...");
    let mph =
        MPHF::load(&mph).with_context(|| format!("Could not load MPH from {}", mph.display()))?;

    log::info!("Hashing input...");

    for (i, line) in std::io::stdin().lines().enumerate() {
        let line = line.with_context(|| format!("Could not read input line {i}"))?;
        let swhid = SWHID::try_from(line.as_str())
            .with_context(|| format!("Could not parse SWHID {line}"))?;
        let node_id = permutation
            .get(
                mph.hash_swhid(&swhid)
                    .ok_or(anyhow!("Unknown SWHID {}", swhid))?, // rejected by MPH
            )
            .ok_or(anyhow!("Unknown SWHID {}", swhid))?; // MPH result too large
        if let Some(node2swhid) = &node2swhid {
            match node2swhid.get(node_id) {
                Err(OutOfBoundError { .. }) => bail!("Unknown SWHID {}", swhid),
                Ok(swhid2) if swhid == swhid2 => (), // expected result
                _ => bail!("Unknown SWHID {}", swhid), // hash collision
            }
        }
        println!("{node_id}");
    }

    Ok(())
}

fn hash_pseudonymized_persons_cmph(mph: PathBuf) -> Result<()> {
    // FIXME: duplicate of hash_pseudonymized_persons?
    log::info!("Loading MPH function...");
    let mph =
        GOVMPH::load(&mph).with_context(|| format!("Could not load MPH from {}", mph.display()))?;

    log::info!("Hashing input...");

    for (i, line) in std::io::stdin().lines().enumerate() {
        let line = line.with_context(|| format!("Could not read input line {i}"))?;
        println!(
            "{}",
            mph.hash_str(&line)
                .ok_or(anyhow!("Unknown value {}", line))?
        );
    }

    Ok(())
}

fn hash_pseudonymized_persons<MPHF: LoadablePersonMphf>(mph: PathBuf) -> Result<()> {
    log::info!("Loading MPH function...");
    let mph =
        MPHF::load(&mph).with_context(|| format!("Could not load MPH from {}", mph.display()))?;
    let hasher = PersonHasher::new(&mph);

    log::info!("Hashing input...");

    for (i, line) in std::io::stdin().lines().enumerate() {
        let line = line.with_context(|| format!("Could not read input line {i}"))?;
        println!(
            "{}",
            hasher
                .hash_pseudonymized_person(&line)
                .with_context(|| format!("Could not hash line {i}"))?
        );
    }

    Ok(())
}

#[cfg(not(feature = "compression"))]
fn hash_pseudonymized_persons_pthash_2024_08_23(_: PathBuf) -> Result<()> {
    bail!(
        "--workaround-2024-08-23 is not supported. Recompile swh-graph with --features=compression"
    );
}

#[cfg(all(feature = "compression", feature = "pthash"))]
fn hash_pseudonymized_persons_pthash_2024_08_23(mph: PathBuf) -> Result<()> {
    use pthash::{DictionaryDictionary, Hashable, Minimal, MurmurHash2_128, PartitionedPhf, Phf};

    pub struct LabelName<T: AsRef<[u8]>>(pub T);

    impl<T: AsRef<[u8]>> Hashable for LabelName<T> {
        type Bytes<'a>
            = &'a [u8]
        where
            T: 'a;
        fn as_bytes(&self) -> Self::Bytes<'_> {
            self.0.as_ref()
        }
    }

    log::info!("Loading MPH function...");
    let mph = <PartitionedPhf<Minimal, MurmurHash2_128, DictionaryDictionary> as Phf>::load(&mph)
        .with_context(|| format!("Could not load MPH from {}", mph.display()))?;

    log::info!("Hashing input...");

    for (i, line) in std::io::stdin().lines().enumerate() {
        let line = line.with_context(|| format!("Could not read input line {i}"))?;
        println!("{}", mph.hash(LabelName(&line)));
    }

    Ok(())
}

fn hash_person_fullnames<MPHF: LoadablePersonMphf>(mph: PathBuf, base64: bool) -> Result<()> {
    log::info!("Loading MPH function...");
    let mph =
        MPHF::load(&mph).with_context(|| format!("Could not load MPH from {}", mph.display()))?;
    let hasher = PersonHasher::new(&mph);

    log::info!("Hashing input...");

    for (i, line) in std::io::stdin().lines().enumerate() {
        let mut line = line
            .with_context(|| format!("Could not read input line {i}"))?
            .into_bytes();
        if base64 {
            line = base64_simd::STANDARD
                .decode_to_vec(&line)
                .with_context(|| {
                    format!("Invalid base64 line: {}", String::from_utf8_lossy(&line))
                })?;
        }
        println!(
            "{}",
            hasher
                .hash_person_fullname(&line)
                .with_context(|| format!("Unknown value {}", String::from_utf8_lossy(&line)))?
        );
    }

    Ok(())
}