scallion/

store.rs

use std::cmp::Ordering;
use std::collections::HashMap;
use std::ffi::OsStr;
use std::fmt::{self, Display};
use std::fs;
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::str::FromStr;

use log::{debug, info};
use rmp_serde::Serializer;
use serde::{Deserialize, Serialize};

use crate::data::{LicenseType, MatchData};
use crate::error::Error;

pub(crate) const CACHE_VERSION: &[u8] = b"scallion-00";
const HEADER_LENGTH: usize = CACHE_VERSION.len() + 5;

/// Entry for a specific license in the license [`Store`].
#[derive(Debug, Deserialize, Serialize)]
pub struct LicenseEntry {
    pub(crate) original: MatchData,
    pub(crate) aliases: Vec<String>,
    pub(crate) headers: Vec<MatchData>,
    pub(crate) alternates: Vec<MatchData>,
}

impl LicenseEntry {
    #[must_use]
    pub(crate) const fn new(original: MatchData) -> LicenseEntry {
        LicenseEntry {
            original,
            aliases: Vec::new(),
            alternates: Vec::new(),
            headers: Vec::new(),
        }
    }

    /// Retrieve original text of this license.
    #[must_use]
    pub const fn original(&self) -> &MatchData {
        &self.original
    }

    /// Retrieve list of aliases for this license.
    #[must_use]
    pub const fn aliases(&self) -> &[String] {
        self.aliases.as_slice()
    }

    /// Retrieve alternate variants of this license.
    #[must_use]
    pub const fn variants(&self) -> &[MatchData] {
        self.alternates.as_slice()
    }

    /// Retrieve header-only variants of this license.
    #[must_use]
    pub const fn headers(&self) -> &[MatchData] {
        self.headers.as_slice()
    }

    /// Add alias of this license's canonical name.
    pub fn add_alias(&mut self, name: String) {
        self.aliases.push(name);
    }

    /// Add alternate format of this license.
    pub fn add_variant(&mut self, data: MatchData) {
        self.alternates.push(data);
    }

    /// Add header-only variant of this license.
    pub fn add_header(&mut self, data: MatchData) {
        self.headers.push(data);
    }
}

/// A representation of a collection of known licenses.
///
/// This struct is generally what you want to start with if you're looking to
/// match text against a database of licenses. Load a cache from disk using
/// `from_cache`, then use the `analyze` function to determine what a text most
/// closely matches.
///
/// # Examples
///
/// ```rust,no_run
/// # use std::fs::File;
/// # use std::error::Error;
/// use scallion::{MatchData, Store};
///
/// # fn main() -> Result<(), Box<dyn Error>> {
/// let store = Store::from_cache(File::open("cache.bin")?)?;
/// let result = store.analyze(&MatchData::from("what's this"));
/// # Ok(())
/// # }
/// ```
#[derive(Debug, Default, Deserialize, Serialize)]
pub struct Store {
    licenses: HashMap<String, LicenseEntry>,
}

impl Store {
    /// Create a new `Store`.
    ///
    /// More often, you probably want to use `from_cache` instead of creating
    /// an empty store.
    #[must_use]
    pub fn new() -> Self {
        Store {
            licenses: HashMap::new(),
        }
    }

    /// Get the number of licenses in the store.
    ///
    /// This only counts licenses by name -- headers, aliases, and alternates
    /// aren't included in the count.
    #[must_use]
    pub fn len(&self) -> usize {
        self.licenses.len()
    }

    /// Check if the store is empty.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.licenses.is_empty()
    }

    /// Add a single license to the store.
    ///
    /// If the license with the given name already existed, it and all of its
    /// variants will be replaced in the store, and returned.
    pub fn add_license(&mut self, name: String, data: MatchData) -> Option<LicenseEntry> {
        let entry = LicenseEntry::new(data);
        self.licenses.insert(name, entry)
    }

    /// Retrieve a single license entry from the store.
    ///
    /// Returns `None` If no license with the given name exists.
    #[must_use]
    pub fn get_license(&self, name: &str) -> Option<&LicenseEntry> {
        self.licenses.get(name)
    }

    /// Retrieve a mutable reference to a license entry from the store.
    ///
    /// Returns `None` If no license with the given name exists.
    #[must_use]
    pub fn get_license_mut(&mut self, name: &str) -> Option<&mut LicenseEntry> {
        self.licenses.get_mut(name)
    }

    /// Insert new unique license data or update an existing entry with a new alias.
    ///
    /// Returns the name of the existing license if only an alias was added,
    /// or `None` if a new entry was added.
    fn insert_or_add_alias(&mut self, name: &str, data: MatchData, header: Option<MatchData>) {
        // check if an identical license is already present
        let mut already_existed = None;
        self.licenses.iter_mut().for_each(|(key, ref mut value)| {
            if value.original.eq_data(&data) {
                value.aliases.push(name.to_string());
                already_existed = Some(key.as_str());
            }
        });
        if let Some(prev) = already_existed {
            info!("{name} already stored; added as an alias for {prev}");
            return;
        }

        let license = self
            .licenses
            .entry(name.to_string())
            .or_insert_with(|| LicenseEntry::new(data));

        if let Some(header_text) = header {
            license.headers = vec![header_text];
        }
    }

    /// Compare the given `MatchData` against all licenses in the `Store`.
    ///
    /// This parallelizes the search as much as it can to find the best match.
    /// Once a match is obtained, it can be optimized further; see methods on
    /// `MatchData` for more information.
    pub fn analyze<'a>(&'a self, text: &MatchData) -> Match<'a> {
        let mut res: Vec<PartialMatch<'a>>;

        let analyze_fold = |mut acc: Vec<PartialMatch<'a>>, (name, data): (&'a String, &'a LicenseEntry)| {
            acc.push(PartialMatch {
                score: data.original.match_score(text),
                name,
                license_type: LicenseType::Original,
                data: &data.original,
            });
            data.alternates.iter().for_each(|alt| {
                acc.push(PartialMatch {
                    score: alt.match_score(text),
                    name,
                    license_type: LicenseType::Alternate,
                    data: alt,
                });
            });
            data.headers.iter().for_each(|head| {
                acc.push(PartialMatch {
                    score: head.match_score(text),
                    name,
                    license_type: LicenseType::Header,
                    data: head,
                });
            });
            acc
        };

        // parallel analysis
        #[cfg(not(target_arch = "wasm32"))]
        {
            use rayon::prelude::*;
            res = self.licenses.par_iter().fold(Vec::new, analyze_fold).reduce(
                Vec::new,
                |mut a: Vec<PartialMatch<'a>>, b: Vec<PartialMatch<'a>>| {
                    a.extend(b);
                    a
                },
            );
            res.par_sort_unstable_by(|a, b| b.partial_cmp(a).unwrap());
        }

        // single-threaded analysis
        #[cfg(target_arch = "wasm32")]
        {
            res = self
                .licenses
                .iter()
                // len of licenses isn't strictly correct, but it'll do
                .fold(Vec::with_capacity(self.licenses.len()), analyze_fold);
            res.sort_unstable_by(|a, b| b.partial_cmp(a).unwrap());
        }

        let m = &res[0];

        Match {
            score: m.score,
            name: m.name,
            license_type: m.license_type,
            data: m.data,
        }
    }

    /// Create a store from a cache file.
    ///
    /// This method is highly useful for quickly loading a cache, as creating
    /// one from text data is rather slow. This method can typically load
    /// the full SPDX set from disk in 200-300 ms. The cache will be
    /// sanity-checked to ensure it was generated with a similar version of
    /// scallion.
    pub fn from_cache<R>(mut readable: R) -> Result<Store, Error>
    where
        R: Read + Sized,
    {
        let mut header = [0u8; HEADER_LENGTH];
        readable.read_exact(&mut header).map_err(|e| Error::io(e, None))?;

        let cf = match &header {
            b"scallion-00-zstd" => CF::Zstd,
            b"scallion-00-gzip" => CF::Gzip,
            b"scallion-00-none" => CF::None,
            _ => return Err(Error::cache_version(header.to_vec())),
        };

        match cf {
            CF::Zstd => {
                #[cfg(feature = "zstd")]
                {
                    let dec = zstd::Decoder::new(readable).map_err(|e| Error::io(e, None))?;
                    let store = rmp_serde::decode::from_read(dec)?;
                    Ok(store)
                }
                #[cfg(not(feature = "zstd"))]
                {
                    Err(Error::cache_format(cf))
                }
            },
            CF::Gzip => {
                #[cfg(feature = "gzip")]
                {
                    let dec = flate2::read::GzDecoder::new(readable);
                    let store = rmp_serde::decode::from_read(dec)?;
                    Ok(store)
                }
                #[cfg(not(feature = "gzip"))]
                {
                    Err(Error::cache_format(cf))
                }
            },
            CF::None => {
                let store = rmp_serde::decode::from_read(readable)?;
                Ok(store)
            },
        }
    }

    /// Serialize the current store.
    ///
    /// The output will be a `MessagePack`'d gzip'd or zstd'd binary stream that should be
    /// written to disk.
    pub fn to_cache<W>(&self, mut writable: W, format: CF) -> Result<(), Error>
    where
        W: Write + Sized,
    {
        let serialize = || -> Result<Vec<u8>, Error> {
            // This currently sits around 3.7MiB, so go up to 4 to fit comfortably
            let mut buf = Vec::with_capacity(4 * 1024 * 1024);
            let mut serializer = Serializer::new(&mut buf);
            self.serialize(&mut serializer)?;
            Ok(buf)
        };

        match format {
            CF::Zstd => {
                #[cfg(feature = "zstd")]
                {
                    writable.write_all(CACHE_VERSION).map_err(|e| Error::io(e, None))?;
                    writable.write_all(b"-zstd").map_err(|e| Error::io(e, None))?;

                    let serialized = serialize()?;
                    let mut enc = zstd::Encoder::new(writable, 21).map_err(|e| Error::io(e, None))?;

                    io::copy(&mut serialized.as_slice(), &mut enc).map_err(|e| Error::io(e, None))?;
                    enc.finish().map_err(|e| Error::io(e, None))?;
                    Ok(())
                }
                #[cfg(not(feature = "zstd"))]
                {
                    Err(Error::cache_format(format))
                }
            },
            CF::Gzip => {
                #[cfg(feature = "gzip")]
                {
                    writable.write_all(CACHE_VERSION).map_err(|e| Error::io(e, None))?;
                    writable.write_all(b"-gzip").map_err(|e| Error::io(e, None))?;

                    let serialized = serialize()?;
                    let mut enc = flate2::write::GzEncoder::new(writable, flate2::Compression::default());

                    io::copy(&mut serialized.as_slice(), &mut enc).map_err(|e| Error::io(e, None))?;
                    enc.finish().map_err(|e| Error::io(e, None))?;
                    Ok(())
                }
                #[cfg(not(feature = "gzip"))]
                {
                    Err(Error::cache_format(format))
                }
            },
            CF::None => {
                writable.write_all(CACHE_VERSION).map_err(|e| Error::io(e, None))?;
                writable.write_all(b"-none").map_err(|e| Error::io(e, None))?;

                let serialized = serialize()?;

                io::copy(&mut serialized.as_slice(), &mut writable).map_err(|e| Error::io(e, None))?;
                Ok(())
            },
        }
    }

    /// Fill the store with SPDX JSON data.
    ///
    /// This function is very specific to the format of SPDX's
    /// `license-list-data` repository. It reads all JSON files in the
    /// `json/details` directory and creates entries inside the store for
    /// matching.
    pub fn load_spdx<P: AsRef<Path>>(&mut self, dir: P) -> Result<(), Error> {
        let paths = locate_json_files(dir)?;

        for path in paths {
            let parsed = parse_json_file(&path)?;

            if parsed.deprecated {
                debug!("Skipping {} (deprecated)", parsed.name);
                continue;
            }
            info!("Processing {}", parsed.name);

            let data = MatchData::new(&parsed.text);
            let header = parsed.header.as_deref().map(MatchData::new);

            self.insert_or_add_alias(&parsed.name, data, header);
        }

        Ok(())
    }
}

/// License data cache compression format.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum CompressionFormat {
    /// zstd compression (using libzstd)
    Zstd,
    /// gzip compression (using `zlib-rs`, or `miniz_oxide` on WASM)
    Gzip,
    /// no compression
    None,
}

use CompressionFormat as CF;

impl Display for CompressionFormat {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            CF::Zstd => write!(f, "zstd"),
            CF::Gzip => write!(f, "gzip"),
            CF::None => write!(f, "none"),
        }
    }
}

impl FromStr for CompressionFormat {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "zstd" => Ok(CF::Zstd),
            "gzip" => Ok(CF::Gzip),
            "none" => Ok(CF::None),
            _ => Err(format!("Invalid compression format: '{s}'")),
        }
    }
}

/// Information about text that was compared against licenses in the store.
///
/// This only contains information about the overall match; to uncover more
/// data you can run methods like `optimize_bounds` on `MatchData`.
///
/// Its lifetime is tied to the lifetime of the `Store` it was generated from.
#[derive(Clone)]
pub struct Match<'a> {
    /// Confidence score of the match, ranging from 0 to 1.
    pub score: f32,
    /// The name of the closest matching license in the `Store`. This will
    /// always be something that exists in the store, regardless of the score.
    pub name: &'a str,
    /// The type of the license that matched. Useful to know if the match was
    /// the complete text, a header, or something else.
    pub license_type: LicenseType,
    /// A reference to the license data that matched inside the `Store`. May be
    /// useful for diagnostic purposes or to further optimize the result.
    pub data: &'a MatchData,
}

/// A lighter version of Match to be used during analysis.
/// Reduces the need for cloning a bunch of fields.
struct PartialMatch<'a> {
    pub name: &'a str,
    pub score: f32,
    pub license_type: LicenseType,
    pub data: &'a MatchData,
}

impl PartialOrd for PartialMatch<'_> {
    fn partial_cmp(&self, other: &PartialMatch<'_>) -> Option<Ordering> {
        self.score.partial_cmp(&other.score)
    }
}

impl PartialEq for PartialMatch<'_> {
    fn eq(&self, other: &PartialMatch<'_>) -> bool {
        self.score.eq(&other.score) && self.name == other.name && self.license_type == other.license_type
    }
}

impl fmt::Debug for Match<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Match {{ score: {}, name: {}, license_type: {:?} }}",
            self.score, self.name, self.license_type
        )
    }
}

fn locate_json_files<P: AsRef<Path>>(dir: P) -> Result<Vec<PathBuf>, Error> {
    // locate all json files in the directory
    let mut paths: Vec<_> = fs::read_dir(&dir)
        .map_err(|e| Error::io(e, Some(dir.as_ref().to_path_buf())))?
        .filter_map(Result::ok)
        .map(|e| e.path())
        .filter(|p| p.is_file() && p.extension().unwrap_or_else(|| OsStr::new("")) == "json")
        .collect();

    // sort without extensions; otherwise dashes and dots muck it up
    paths.sort_by(|a, b| a.file_stem().unwrap().cmp(b.file_stem().unwrap()));

    Ok(paths)
}

#[derive(Deserialize)]
struct LicenseListData {
    #[serde(rename = "licenseId")]
    name: String,
    #[serde(rename = "isDeprecatedLicenseId")]
    deprecated: bool,
    #[serde(rename = "licenseText")]
    text: String,
    #[serde(rename = "standardLicenseHeader")]
    header: Option<String>,
    // incomplete
}

fn parse_json_file<P: AsRef<Path>>(path: P) -> Result<LicenseListData, Error> {
    let path = path.as_ref().to_path_buf();
    let data = fs::read_to_string(&path).map_err(|e| Error::io(e, Some(path.clone())))?;
    serde_json::from_str(&data).map_err(|e| Error::spdx(e, path))
}