use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::hash::Hash;
use std::os::unix::fs::FileExt;
use std::path::Path;
use std::sync::Arc;
use anyhow::{anyhow, Result};
use arrow::array::{Array, StringArray};
use arrow::record_batch::RecordBatch;
use crate::codec;
use crate::index::{read_znippy_index_filtered, IndexFilter};
use crate::plugin::file_name;
pub const RUST_PKG_TYPE: i8 = 1;
pub const PYTHON_PKG_TYPE: i8 = 2;
pub const MAVEN_PKG_TYPE: i8 = 3;
pub const NPM_PKG_TYPE: i8 = 6;
pub const GEM_PKG_TYPE: i8 = 11;
pub const RPM_PKG_TYPE: i8 = 8;
pub const DEB_PKG_TYPE: i8 = 9;
pub const CONDA_PKG_TYPE: i8 = 14;
#[derive(Debug, Clone)]
pub struct FileLoc {
pub relative_path: String,
pub chunks: Vec<ChunkRef>,
pub uncompressed_size: u64,
}
#[derive(Debug, Clone, Copy)]
pub struct ChunkRef {
pub blob_offset: u64,
pub blob_size: u64,
pub fdata_offset: u64,
pub compressed: bool,
}
impl FileLoc {
fn read_bytes(&self, archive: &File) -> Result<Vec<u8>> {
let mut result = Vec::with_capacity(self.uncompressed_size as usize);
let mut blob = Vec::new();
let mut decomp = Vec::new();
for chunk in &self.chunks {
blob.resize(chunk.blob_size as usize, 0);
archive.read_exact_at(&mut blob, chunk.blob_offset)?;
if chunk.compressed {
codec::decompress_into(&blob, &mut decomp)?;
result.extend_from_slice(&decomp);
} else {
result.extend_from_slice(&blob);
}
}
Ok(result)
}
}
fn group_rows_by_file(batch: &RecordBatch) -> Result<HashMap<String, FileLoc>> {
use arrow::array::{BooleanArray, UInt64Array};
let col = |n: &str| {
batch
.column_by_name(n)
.ok_or_else(|| anyhow!("index missing column {n}"))
};
let paths = col("relative_path")?
.as_any()
.downcast_ref::<StringArray>()
.ok_or_else(|| anyhow!("relative_path not StringArray"))?;
let compressed = col("compressed")?
.as_any()
.downcast_ref::<BooleanArray>()
.ok_or_else(|| anyhow!("compressed not BooleanArray"))?;
let sizes = col("uncompressed_size")?
.as_any()
.downcast_ref::<UInt64Array>()
.ok_or_else(|| anyhow!("uncompressed_size not UInt64Array"))?;
let blob_offset = col("blob_offset")?
.as_any()
.downcast_ref::<UInt64Array>()
.ok_or_else(|| anyhow!("blob_offset not UInt64Array"))?;
let blob_size = col("blob_size")?
.as_any()
.downcast_ref::<UInt64Array>()
.ok_or_else(|| anyhow!("blob_size not UInt64Array"))?;
let fdata = col("fdata_offset")?
.as_any()
.downcast_ref::<UInt64Array>()
.ok_or_else(|| anyhow!("fdata_offset not UInt64Array"))?;
let mut by_path: HashMap<String, FileLoc> = HashMap::new();
for i in 0..batch.num_rows() {
let path = paths.value(i);
let entry = by_path.entry(path.to_string()).or_insert_with(|| FileLoc {
relative_path: path.to_string(),
chunks: Vec::new(),
uncompressed_size: 0,
});
entry.uncompressed_size += sizes.value(i);
entry.chunks.push(ChunkRef {
blob_offset: blob_offset.value(i),
blob_size: blob_size.value(i),
fdata_offset: fdata.value(i),
compressed: compressed.value(i),
});
}
for f in by_path.values_mut() {
f.chunks.sort_by_key(|c| c.fdata_offset);
}
Ok(by_path)
}
mod sealed {
pub trait Sealed {}
}
pub trait CoordSpec: sealed::Sealed + Send + Sync + 'static {
const PKG_TYPE: i8;
const SLOT: usize;
const COLS: &'static [&'static str];
const REQUIRED: usize;
type Key: Eq + Hash + Clone + Send + Sync;
type Extra: Clone + Send + Sync;
fn row(c: &Cells<'_>, path: &str) -> (Self::Key, Self::Extra);
fn replaces(_new: &Self::Extra) -> bool {
true
}
}
pub struct Cells<'a> {
cols: &'a [Option<&'a StringArray>],
i: usize,
}
impl Cells<'_> {
pub fn opt(&self, j: usize) -> Option<String> {
self.cols[j].filter(|a| !a.is_null(self.i)).map(|a| a.value(self.i).to_string())
}
pub fn get(&self, j: usize) -> String {
self.opt(j).unwrap_or_default()
}
pub fn non_empty(&self, j: usize) -> Option<String> {
self.opt(j).filter(|s| !s.is_empty())
}
}
pub struct CoordView<S: CoordSpec> {
archive: Arc<File>,
coords: HashMap<S::Key, (FileLoc, S::Extra)>,
}
pub struct Package<S: CoordSpec> {
archive: Arc<File>,
loc: FileLoc,
key: S::Key,
extra: S::Extra,
}
impl<S: CoordSpec> CoordView<S> {
pub(crate) fn build(path: &Path, archive: Arc<File>) -> Result<Self> {
let (_schema, batches) =
read_znippy_index_filtered(path, &IndexFilter { pkg_type: Some(S::PKG_TYPE), repo: None })?;
let mut coords: HashMap<S::Key, (FileLoc, S::Extra)> = HashMap::new();
for batch in &batches {
let cols: Vec<Option<&StringArray>> = S::COLS
.iter()
.map(|n| batch.column_by_name(n).and_then(|c| c.as_any().downcast_ref::<StringArray>()))
.collect();
if cols[..S::REQUIRED].iter().any(Option::is_none) {
continue;
}
let paths = batch
.column_by_name("relative_path")
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
.ok_or_else(|| anyhow!("missing relative_path"))?;
let locs = group_rows_by_file(batch)?;
let mut seen = HashSet::new();
for i in 0..batch.num_rows() {
let p = paths.value(i);
if !seen.insert(p) || cols[..S::REQUIRED].iter().any(|c| c.is_some_and(|a| a.is_null(i))) {
continue;
}
let Some(loc) = locs.get(p) else { continue };
let (key, extra) = S::row(&Cells { cols: &cols, i }, p);
if S::replaces(&extra) || !coords.contains_key(&key) {
coords.insert(key, (loc.clone(), extra));
}
}
}
Ok(Self { archive, coords })
}
fn package(&self, key: &S::Key) -> Option<Package<S>> {
let (loc, extra) = self.coords.get(key)?;
Some(Package { archive: Arc::clone(&self.archive), loc: loc.clone(), key: key.clone(), extra: extra.clone() })
}
fn entries(&self) -> impl Iterator<Item = (&S::Key, &S::Extra)> {
self.coords.iter().map(|(k, (_, e))| (k, e))
}
pub fn len(&self) -> usize {
self.coords.len()
}
pub fn is_empty(&self) -> bool {
self.coords.is_empty()
}
}
impl<S: CoordSpec> Package<S> {
pub fn relative_path(&self) -> &str {
&self.loc.relative_path
}
pub fn size(&self) -> u64 {
self.loc.uncompressed_size
}
pub fn bytes(&self) -> Result<Vec<u8>> {
self.loc.read_bytes(&self.archive)
}
pub fn into_bytes(self) -> Result<Vec<u8>> {
self.loc.read_bytes(&self.archive)
}
}
macro_rules! spec {
($t:ident) => {
#[doc = concat!("The coord spec behind [`", stringify!($t), "`]'s view.")]
pub struct $t;
impl sealed::Sealed for $t {}
};
}
fn s2(k: &(String, String)) -> (String, String) {
k.clone()
}
spec!(RustSpec);
impl CoordSpec for RustSpec {
const PKG_TYPE: i8 = RUST_PKG_TYPE;
const SLOT: usize = 0;
const COLS: &'static [&'static str] = &["crate_name", "version"];
const REQUIRED: usize = 2;
type Key = (String, String);
type Extra = ();
fn row(c: &Cells<'_>, _: &str) -> (Self::Key, ()) {
((c.get(0), c.get(1)), ())
}
}
pub type RustView = CoordView<RustSpec>;
pub type RustPackage = Package<RustSpec>;
impl RustView {
pub fn get(&self, name: &str, version: &str) -> Option<RustPackage> {
self.package(&(name.to_string(), version.to_string()))
}
pub fn list(&self) -> Vec<(String, String)> {
self.entries().map(|(k, _)| s2(k)).collect()
}
}
impl RustPackage {
pub fn name(&self) -> &str {
&self.key.0
}
pub fn version(&self) -> &str {
&self.key.1
}
}
spec!(MavenSpec);
impl CoordSpec for MavenSpec {
const PKG_TYPE: i8 = MAVEN_PKG_TYPE;
const SLOT: usize = 1;
const COLS: &'static [&'static str] = &["group_id", "artifact_id", "version", "classifier"];
const REQUIRED: usize = 3;
type Key = (String, String, String, Option<String>);
type Extra = ();
fn row(c: &Cells<'_>, path: &str) -> (Self::Key, ()) {
let (g, a, v) = (c.get(0), c.get(1), c.get(2));
let classifier = c.non_empty(3).or_else(|| derive_classifier(file_name(path), &a, &v));
((g, a, v, classifier), ())
}
}
pub type MavenView = CoordView<MavenSpec>;
pub type MavenPackage = Package<MavenSpec>;
impl MavenView {
pub fn get(&self, group: &str, artifact: &str, version: &str) -> Option<MavenPackage> {
self.get_classified(group, artifact, version, None)
}
pub fn get_classified(
&self,
group: &str,
artifact: &str,
version: &str,
classifier: Option<&str>,
) -> Option<MavenPackage> {
self.package(&(group.into(), artifact.into(), version.into(), classifier.map(str::to_string)))
}
pub fn list(&self) -> Vec<(String, String, String, Option<String>)> {
self.entries().map(|(k, _)| k.clone()).collect()
}
}
fn derive_classifier(filename: &str, artifact: &str, version: &str) -> Option<String> {
let stem = filename.rsplit_once('.').map(|(s, _)| s).unwrap_or(filename);
let rest = stem.strip_prefix(&format!("{artifact}-{version}"))?.strip_prefix('-')?;
(!rest.is_empty()).then(|| rest.to_string())
}
impl MavenPackage {
pub fn group(&self) -> &str {
&self.key.0
}
pub fn artifact(&self) -> &str {
&self.key.1
}
pub fn version(&self) -> &str {
&self.key.2
}
pub fn classifier(&self) -> Option<&str> {
self.key.3.as_deref()
}
pub fn coords(&self) -> (&str, &str, &str, Option<&str>) {
(&self.key.0, &self.key.1, &self.key.2, self.key.3.as_deref())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PythonKind {
Wheel,
Sdist,
}
spec!(PythonSpec);
impl CoordSpec for PythonSpec {
const PKG_TYPE: i8 = PYTHON_PKG_TYPE;
const SLOT: usize = 2;
const COLS: &'static [&'static str] = &["name", "version"];
const REQUIRED: usize = 2;
type Key = (String, String);
type Extra = PythonKind;
fn row(c: &Cells<'_>, path: &str) -> (Self::Key, PythonKind) {
let kind = if file_name(path).ends_with(".whl") { PythonKind::Wheel } else { PythonKind::Sdist };
((c.get(0), c.get(1)), kind)
}
fn replaces(new: &PythonKind) -> bool {
*new == PythonKind::Wheel
}
}
pub type PythonView = CoordView<PythonSpec>;
pub type PythonPackage = Package<PythonSpec>;
impl PythonView {
pub fn get(&self, name: &str, version: &str) -> Option<PythonPackage> {
self.package(&(name.to_string(), version.to_string()))
}
pub fn list(&self) -> Vec<(String, String)> {
self.entries().map(|(k, _)| s2(k)).collect()
}
}
impl PythonPackage {
pub fn name(&self) -> &str {
&self.key.0
}
pub fn version(&self) -> &str {
&self.key.1
}
pub fn kind(&self) -> PythonKind {
self.extra
}
}
spec!(NpmSpec);
impl CoordSpec for NpmSpec {
const PKG_TYPE: i8 = NPM_PKG_TYPE;
const SLOT: usize = 3;
const COLS: &'static [&'static str] = &["name", "version"];
const REQUIRED: usize = 2;
type Key = (String, String);
type Extra = ();
fn row(c: &Cells<'_>, _: &str) -> (Self::Key, ()) {
((c.get(0), c.get(1)), ())
}
}
pub type NpmView = CoordView<NpmSpec>;
pub type NpmPackage = Package<NpmSpec>;
impl NpmView {
pub fn get(&self, name: &str, version: &str) -> Option<NpmPackage> {
self.package(&(name.to_string(), version.to_string()))
}
pub fn list(&self) -> Vec<(String, String)> {
self.entries().map(|(k, _)| s2(k)).collect()
}
}
impl NpmPackage {
pub fn name(&self) -> &str {
&self.key.0
}
pub fn version(&self) -> &str {
&self.key.1
}
}
spec!(GemSpec);
impl CoordSpec for GemSpec {
const PKG_TYPE: i8 = GEM_PKG_TYPE;
const SLOT: usize = 4;
const COLS: &'static [&'static str] = &["name", "version", "platform"];
const REQUIRED: usize = 2;
type Key = (String, String, String);
type Extra = ();
fn row(c: &Cells<'_>, _: &str) -> (Self::Key, ()) {
((c.get(0), c.get(1), c.non_empty(2).unwrap_or_else(|| "ruby".into())), ())
}
}
pub type GemView = CoordView<GemSpec>;
pub type GemPackage = Package<GemSpec>;
impl GemView {
pub fn get(&self, name: &str, version: &str) -> Option<GemPackage> {
self.get_platform(name, version, "ruby")
}
pub fn get_platform(&self, name: &str, version: &str, platform: &str) -> Option<GemPackage> {
self.package(&(name.into(), version.into(), platform.into()))
}
pub fn list(&self) -> Vec<(String, String, String)> {
self.entries().map(|(k, _)| k.clone()).collect()
}
}
impl GemPackage {
pub fn name(&self) -> &str {
&self.key.0
}
pub fn version(&self) -> &str {
&self.key.1
}
pub fn platform(&self) -> &str {
&self.key.2
}
}
spec!(CondaSpec);
impl CoordSpec for CondaSpec {
const PKG_TYPE: i8 = CONDA_PKG_TYPE;
const SLOT: usize = 5;
const COLS: &'static [&'static str] = &["name", "version", "build", "subdir"];
const REQUIRED: usize = 2;
type Key = (String, String, String, String);
type Extra = ();
fn row(c: &Cells<'_>, _: &str) -> (Self::Key, ()) {
((c.get(0), c.get(1), c.get(2), c.get(3)), ())
}
}
pub type CondaView = CoordView<CondaSpec>;
pub type CondaPackage = Package<CondaSpec>;
impl CondaView {
pub fn get(&self, name: &str, version: &str) -> Option<CondaPackage> {
let key = self.coords.keys().find(|k| k.0 == name && k.1 == version)?.clone();
self.package(&key)
}
pub fn get_exact(&self, name: &str, version: &str, build: &str, subdir: &str) -> Option<CondaPackage> {
self.package(&(name.into(), version.into(), build.into(), subdir.into()))
}
pub fn list(&self) -> Vec<(String, String, String, String)> {
self.entries().map(|(k, _)| k.clone()).collect()
}
}
impl CondaPackage {
pub fn name(&self) -> &str {
&self.key.0
}
pub fn version(&self) -> &str {
&self.key.1
}
pub fn build(&self) -> &str {
&self.key.2
}
pub fn subdir(&self) -> &str {
&self.key.3
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RpmMetaRow {
pub name: String,
pub version: String,
pub release: String,
pub arch: String,
pub epoch: Option<String>,
pub summary: Option<String>,
pub license: Option<String>,
pub url: Option<String>,
pub vendor: Option<String>,
pub sourcerpm: Option<String>,
pub provides: Vec<String>,
pub requires: Vec<String>,
}
spec!(RpmSpec);
impl CoordSpec for RpmSpec {
const PKG_TYPE: i8 = RPM_PKG_TYPE;
const SLOT: usize = 6;
const COLS: &'static [&'static str] = &[
"name", "version", "release", "arch", "epoch", "summary", "license", "url", "vendor",
"sourcerpm", "provides", "requires",
];
const REQUIRED: usize = 2;
type Key = (String, String, String, String);
type Extra = RpmMetaRow;
fn row(c: &Cells<'_>, _: &str) -> (Self::Key, RpmMetaRow) {
let lines = |j| split_lines(c.opt(j));
let m = RpmMetaRow {
name: c.get(0),
version: c.get(1),
release: c.get(2),
arch: c.get(3),
epoch: c.opt(4),
summary: c.opt(5),
license: c.opt(6),
url: c.opt(7),
vendor: c.opt(8),
sourcerpm: c.opt(9),
provides: lines(10),
requires: lines(11),
};
((m.name.clone(), m.version.clone(), m.release.clone(), m.arch.clone()), m)
}
}
pub type RpmView = CoordView<RpmSpec>;
pub type RpmPackage = Package<RpmSpec>;
impl RpmView {
pub fn get(&self, name: &str, version: &str, release: &str, arch: &str) -> Option<RpmPackage> {
self.package(&(name.into(), version.into(), release.into(), arch.into()))
}
pub fn list(&self) -> Vec<(String, String, String, String, Option<String>)> {
self.entries().map(|(k, e)| (k.0.clone(), k.1.clone(), k.2.clone(), k.3.clone(), e.epoch.clone())).collect()
}
pub fn list_meta(&self) -> Vec<RpmMetaRow> {
self.entries().map(|(_, e)| e.clone()).collect()
}
}
impl RpmPackage {
pub fn name(&self) -> &str {
&self.key.0
}
pub fn version(&self) -> &str {
&self.key.1
}
pub fn release(&self) -> &str {
&self.key.2
}
pub fn arch(&self) -> &str {
&self.key.3
}
pub fn epoch(&self) -> Option<&str> {
self.extra.epoch.as_deref()
}
}
fn split_lines(v: Option<String>) -> Vec<String> {
v.map(|s| s.lines().filter(|l| !l.is_empty()).map(str::to_string).collect()).unwrap_or_default()
}
spec!(DebSpec);
impl CoordSpec for DebSpec {
const PKG_TYPE: i8 = DEB_PKG_TYPE;
const SLOT: usize = 7;
const COLS: &'static [&'static str] = &["name", "version", "arch", "control"];
const REQUIRED: usize = 1;
type Key = (String, String, String);
type Extra = Option<String>;
fn row(c: &Cells<'_>, _: &str) -> (Self::Key, Option<String>) {
((c.get(0), c.get(1), c.get(2)), c.opt(3))
}
}
pub type DebView = CoordView<DebSpec>;
pub type DebPackage = Package<DebSpec>;
impl DebView {
pub fn get(&self, name: &str, version: &str, arch: &str) -> Option<DebPackage> {
self.package(&(name.into(), version.into(), arch.into()))
}
pub fn list(&self) -> Vec<(String, String, String, Option<String>)> {
self.entries().map(|(k, e)| (k.0.clone(), k.1.clone(), k.2.clone(), e.clone())).collect()
}
}
impl DebPackage {
pub fn name(&self) -> &str {
&self.key.0
}
pub fn version(&self) -> &str {
&self.key.1
}
pub fn arch(&self) -> &str {
&self.key.2
}
pub fn control(&self) -> Option<&str> {
self.extra.as_deref()
}
}
pub(crate) const VIEW_SLOTS: usize = 8;