use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use anyhow::{Result, anyhow};
use arrow::array::{
Array, BinaryArray, BinaryBuilder, BooleanArray, BooleanBuilder, Float64Array,
Float64Builder, Int8Array, Int8Builder, Int64Array, Int64Builder, StringArray, StringBuilder,
};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use serde::{Deserialize, Serialize};
use crate::index::{
FORMAT_VERSION_KEY, META_MODULE, ZNIPPY_FORMAT_VERSION, check_format_version,
read_reserved_section_bytes,
};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum MetaValue {
Str(String),
I64(i64),
F64(f64),
Bool(bool),
Bytes(Vec<u8>),
}
impl MetaValue {
fn tag(&self) -> i8 {
match self {
MetaValue::Str(_) => 0,
MetaValue::I64(_) => 1,
MetaValue::F64(_) => 2,
MetaValue::Bool(_) => 3,
MetaValue::Bytes(_) => 4,
}
}
pub fn as_str(&self) -> Option<&str> {
match self {
MetaValue::Str(s) => Some(s),
_ => None,
}
}
pub fn as_i64(&self) -> Option<i64> {
match self {
MetaValue::I64(v) => Some(*v),
_ => None,
}
}
pub fn as_f64(&self) -> Option<f64> {
match self {
MetaValue::F64(v) => Some(*v),
_ => None,
}
}
pub fn as_bool(&self) -> Option<bool> {
match self {
MetaValue::Bool(v) => Some(*v),
_ => None,
}
}
pub fn as_bytes(&self) -> Option<&[u8]> {
match self {
MetaValue::Bytes(b) => Some(b),
_ => None,
}
}
}
impl From<&str> for MetaValue {
fn from(v: &str) -> Self {
MetaValue::Str(v.to_string())
}
}
impl From<String> for MetaValue {
fn from(v: String) -> Self {
MetaValue::Str(v)
}
}
impl From<i64> for MetaValue {
fn from(v: i64) -> Self {
MetaValue::I64(v)
}
}
impl From<f64> for MetaValue {
fn from(v: f64) -> Self {
MetaValue::F64(v)
}
}
impl From<bool> for MetaValue {
fn from(v: bool) -> Self {
MetaValue::Bool(v)
}
}
impl From<Vec<u8>> for MetaValue {
fn from(v: Vec<u8>) -> Self {
MetaValue::Bytes(v)
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MetaEntry {
pub relative_path: Option<String>,
pub key: String,
pub value: MetaValue,
}
impl MetaEntry {
pub fn entry(relative_path: impl Into<String>, key: impl Into<String>, value: impl Into<MetaValue>) -> Self {
Self {
relative_path: Some(relative_path.into()),
key: key.into(),
value: value.into(),
}
}
pub fn archive(key: impl Into<String>, value: impl Into<MetaValue>) -> Self {
Self { relative_path: None, key: key.into(), value: value.into() }
}
pub fn path(&self) -> Option<&str> {
self.relative_path.as_deref()
}
fn sort_key(&self) -> (&str, Option<&str>) {
(self.key.as_str(), self.relative_path.as_deref())
}
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct MetaTable {
rows: Vec<MetaEntry>,
}
impl MetaTable {
pub fn new() -> Self {
Self::default()
}
pub fn from_rows(rows: Vec<MetaEntry>) -> Self {
Self { rows }
}
pub fn insert(
&mut self,
relative_path: impl Into<String>,
key: impl Into<String>,
value: impl Into<MetaValue>,
) -> &mut Self {
self.rows.push(MetaEntry::entry(relative_path, key, value));
self
}
pub fn insert_archive(&mut self, key: impl Into<String>, value: impl Into<MetaValue>) -> &mut Self {
self.rows.push(MetaEntry::archive(key, value));
self
}
pub fn extend(&mut self, rows: impl IntoIterator<Item = MetaEntry>) -> &mut Self {
self.rows.extend(rows);
self
}
pub fn rows(&self) -> &[MetaEntry] {
&self.rows
}
pub fn len(&self) -> usize {
self.rows.len()
}
pub fn is_empty(&self) -> bool {
self.rows.is_empty()
}
fn sorted_rows(&self) -> Vec<&MetaEntry> {
let mut v: Vec<&MetaEntry> = self.rows.iter().collect();
v.sort_by(|a, b| a.sort_key().cmp(&b.sort_key()));
v
}
}
pub fn meta_schema() -> Arc<Schema> {
let fields = vec![
Field::new("relative_path", DataType::Utf8, true),
Field::new("key", DataType::Utf8, false),
Field::new("value_type", DataType::Int8, false),
Field::new("v_str", DataType::Utf8, true),
Field::new("v_i64", DataType::Int64, true),
Field::new("v_f64", DataType::Float64, true),
Field::new("v_bool", DataType::Boolean, true),
Field::new("v_bytes", DataType::Binary, true),
];
let mut md = HashMap::new();
md.insert(FORMAT_VERSION_KEY.to_string(), ZNIPPY_FORMAT_VERSION.to_string());
Arc::new(Schema::new_with_metadata(fields, md))
}
pub fn build_meta_batch(table: &MetaTable) -> Result<RecordBatch> {
let rows = table.sorted_rows();
let n = rows.len();
let mut path_b = StringBuilder::with_capacity(n, n * 24);
let mut key_b = StringBuilder::with_capacity(n, n * 16);
let mut tag_b = Int8Builder::with_capacity(n);
let mut s_b = StringBuilder::with_capacity(n, n * 16);
let mut i_b = Int64Builder::with_capacity(n);
let mut f_b = Float64Builder::with_capacity(n);
let mut bo_b = BooleanBuilder::with_capacity(n);
let mut by_b = BinaryBuilder::with_capacity(n, n * 16);
for r in rows {
match &r.relative_path {
Some(p) => path_b.append_value(p),
None => path_b.append_null(),
}
key_b.append_value(&r.key);
tag_b.append_value(r.value.tag());
match &r.value {
MetaValue::Str(s) => {
s_b.append_value(s);
i_b.append_null();
f_b.append_null();
bo_b.append_null();
by_b.append_null();
}
MetaValue::I64(v) => {
s_b.append_null();
i_b.append_value(*v);
f_b.append_null();
bo_b.append_null();
by_b.append_null();
}
MetaValue::F64(v) => {
s_b.append_null();
i_b.append_null();
f_b.append_value(*v);
bo_b.append_null();
by_b.append_null();
}
MetaValue::Bool(v) => {
s_b.append_null();
i_b.append_null();
f_b.append_null();
bo_b.append_value(*v);
by_b.append_null();
}
MetaValue::Bytes(b) => {
s_b.append_null();
i_b.append_null();
f_b.append_null();
bo_b.append_null();
by_b.append_value(b);
}
}
}
RecordBatch::try_new(meta_schema(), vec![
Arc::new(path_b.finish()),
Arc::new(key_b.finish()),
Arc::new(tag_b.finish()),
Arc::new(s_b.finish()),
Arc::new(i_b.finish()),
Arc::new(f_b.finish()),
Arc::new(bo_b.finish()),
Arc::new(by_b.finish()),
])
.map_err(|e| anyhow!("build meta sub-index batch: {e}"))
}
#[derive(Debug, Clone, PartialEq)]
pub struct MetaIndex {
rows: Vec<MetaEntry>,
}
impl MetaIndex {
pub fn from_rows(mut rows: Vec<MetaEntry>) -> Self {
rows.sort_by(|a, b| a.sort_key().cmp(&b.sort_key()));
Self { rows }
}
pub fn len(&self) -> usize {
self.rows.len()
}
pub fn is_empty(&self) -> bool {
self.rows.is_empty()
}
pub fn iter(&self) -> std::slice::Iter<'_, MetaEntry> {
self.rows.iter()
}
pub fn find_by_key(&self, key: &str) -> &[MetaEntry] {
let lo = self.rows.partition_point(|r| r.key.as_str() < key);
let hi = self.rows.partition_point(|r| r.key.as_str() <= key);
&self.rows[lo..hi]
}
pub fn find_by_prefix(&self, prefix: &str) -> &[MetaEntry] {
let lo = self.rows.partition_point(|r| r.key.as_str() < prefix);
let hi = self.rows.partition_point(|r| r.key.as_str() < prefix || r.key.starts_with(prefix));
&self.rows[lo..hi]
}
pub fn archive_value(&self, key: &str) -> Option<&MetaValue> {
self.find_by_key(key)
.iter()
.find(|r| r.relative_path.is_none())
.map(|r| &r.value)
}
pub fn keys(&self) -> Vec<&str> {
let mut out: Vec<&str> = Vec::new();
for r in &self.rows {
if out.last() != Some(&r.key.as_str()) {
out.push(r.key.as_str());
}
}
out
}
pub fn to_table(&self) -> MetaTable {
MetaTable::from_rows(self.rows.clone())
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum ArchiveMeta {
NoMetadata,
Index(MetaIndex),
}
impl ArchiveMeta {
pub fn index(&self) -> Option<&MetaIndex> {
match self {
ArchiveMeta::NoMetadata => None,
ArchiveMeta::Index(i) => Some(i),
}
}
pub fn is_searchable(&self) -> bool {
matches!(self, ArchiveMeta::Index(_))
}
pub fn find_by_key(&self, key: &str) -> MetaSearch<'_> {
match self {
ArchiveMeta::NoMetadata => MetaSearch::NoMetadata,
ArchiveMeta::Index(i) => MetaSearch::Hits(i.find_by_key(key)),
}
}
pub fn find_by_prefix(&self, prefix: &str) -> MetaSearch<'_> {
match self {
ArchiveMeta::NoMetadata => MetaSearch::NoMetadata,
ArchiveMeta::Index(i) => MetaSearch::Hits(i.find_by_prefix(prefix)),
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum MetaSearch<'a> {
NoMetadata,
Hits(&'a [MetaEntry]),
}
impl<'a> MetaSearch<'a> {
pub fn hits(&self) -> Option<&'a [MetaEntry]> {
match self {
MetaSearch::NoMetadata => None,
MetaSearch::Hits(h) => Some(h),
}
}
pub fn found_any(&self) -> bool {
matches!(self, MetaSearch::Hits(h) if !h.is_empty())
}
}
pub fn decode_meta_section(bytes: &[u8]) -> Result<MetaIndex> {
use arrow::ipc::reader::StreamReader;
let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None)
.map_err(|e| anyhow!("meta sub-index: not a readable Arrow stream: {e}"))?;
check_format_version(reader.schema().metadata())?;
let mut rows = Vec::new();
for batch in reader {
let batch = batch.map_err(|e| anyhow!("meta sub-index read: {e}"))?;
decode_meta_batch_into(&batch, &mut rows)?;
}
Ok(MetaIndex::from_rows(rows))
}
fn col<'a, T: 'static>(batch: &'a RecordBatch, name: &str) -> Result<&'a T> {
batch
.column_by_name(name)
.ok_or_else(|| anyhow!("meta sub-index missing column {name:?}"))?
.as_any()
.downcast_ref::<T>()
.ok_or_else(|| anyhow!("meta sub-index column {name:?} has an unexpected Arrow type"))
}
fn decode_meta_batch_into(batch: &RecordBatch, out: &mut Vec<MetaEntry>) -> Result<()> {
let paths = col::<StringArray>(batch, "relative_path")?;
let keys = col::<StringArray>(batch, "key")?;
let tags = col::<Int8Array>(batch, "value_type")?;
let v_str = col::<StringArray>(batch, "v_str")?;
let v_i64 = col::<Int64Array>(batch, "v_i64")?;
let v_f64 = col::<Float64Array>(batch, "v_f64")?;
let v_bool = col::<BooleanArray>(batch, "v_bool")?;
let v_bytes = col::<BinaryArray>(batch, "v_bytes")?;
out.reserve(batch.num_rows());
for i in 0..batch.num_rows() {
let want = |present: bool, what: &str| -> Result<()> {
anyhow::ensure!(present, "meta row {i} declares {what} but that column is null");
Ok(())
};
let value = match tags.value(i) {
0 => {
want(v_str.is_valid(i), "a string value")?;
MetaValue::Str(v_str.value(i).to_string())
}
1 => {
want(v_i64.is_valid(i), "an i64 value")?;
MetaValue::I64(v_i64.value(i))
}
2 => {
want(v_f64.is_valid(i), "an f64 value")?;
MetaValue::F64(v_f64.value(i))
}
3 => {
want(v_bool.is_valid(i), "a bool value")?;
MetaValue::Bool(v_bool.value(i))
}
4 => {
want(v_bytes.is_valid(i), "a bytes value")?;
MetaValue::Bytes(v_bytes.value(i).to_vec())
}
other => return Err(anyhow!("meta row {i} has unknown value_type {other}")),
};
out.push(MetaEntry {
relative_path: paths.is_valid(i).then(|| paths.value(i).to_string()),
key: keys.value(i).to_string(),
value,
});
}
Ok(())
}
pub fn read_archive_meta(path: &Path) -> Result<ArchiveMeta> {
match read_reserved_section_bytes(path, META_MODULE)? {
None => Ok(ArchiveMeta::NoMetadata),
Some(bytes) => Ok(ArchiveMeta::Index(decode_meta_section(&bytes)?)),
}
}
#[cfg(test)]
mod tests {
use super::*;
fn sample() -> MetaTable {
let mut t = MetaTable::new();
t.insert("app/main.wasm", "build-thing", MetaValue::Bytes(vec![0, 97, 115, 109, 1]))
.insert("app/main.wasm", "build-thing.abi", "wasi-p2")
.insert("app/main.wasm", "size", 5i64)
.insert("lib/util.rs", "build-thing.abi", "native")
.insert("lib/util.rs", "coverage", 0.87f64)
.insert("lib/util.rs", "vendored", false)
.insert_archive("producer", "znippy")
.insert_archive("build-thing", MetaValue::Bytes(vec![1, 2, 3]));
t
}
#[test]
fn every_value_type_and_both_scopes_round_trip() {
let t = sample();
let batch = build_meta_batch(&t).unwrap();
assert_eq!(batch.num_rows(), t.len());
let mut rows = Vec::new();
decode_meta_batch_into(&batch, &mut rows).unwrap();
let idx = MetaIndex::from_rows(rows);
assert_eq!(idx.len(), t.len());
let bt = idx.find_by_key("build-thing");
assert_eq!(bt.len(), 2, "one entry-scoped + one archive-scoped");
assert_eq!(bt[0].path(), None, "archive-scoped sorts first (NULL path)");
assert_eq!(bt[0].value.as_bytes(), Some(&[1u8, 2, 3][..]));
assert_eq!(bt[1].path(), Some("app/main.wasm"));
assert_eq!(bt[1].value.as_bytes(), Some(&[0u8, 97, 115, 109, 1][..]));
assert_eq!(idx.find_by_key("size")[0].value.as_i64(), Some(5));
assert_eq!(idx.find_by_key("coverage")[0].value.as_f64(), Some(0.87));
assert_eq!(idx.find_by_key("vendored")[0].value.as_bool(), Some(false));
assert_eq!(idx.archive_value("producer").and_then(MetaValue::as_str), Some("znippy"));
let ordered: Vec<_> = idx.iter().map(|r| (r.key.as_str(), r.path())).collect();
let mut want = ordered.clone();
want.sort();
assert_eq!(ordered, want, "rows must be stored key-major and sorted");
let json = serde_json::to_string(t.rows()).unwrap();
let back: Vec<MetaEntry> = serde_json::from_str(&json).unwrap();
assert_eq!(back, t.rows());
}
#[test]
fn key_and_prefix_search_return_exactly_the_matching_rows() {
let idx = MetaIndex::from_rows(sample().rows().to_vec());
let exact = idx.find_by_key("build-thing");
assert_eq!(exact.len(), 2, "exact key must NOT sweep in `build-thing.abi`");
assert!(exact.iter().all(|r| r.key == "build-thing"));
let pre = idx.find_by_prefix("build-thing");
assert_eq!(pre.len(), 4, "prefix picks up build-thing + build-thing.abi ×2");
assert!(pre.iter().all(|r| r.key.starts_with("build-thing")));
let paths: Vec<_> = idx.find_by_key("build-thing.abi").iter().filter_map(|r| r.path()).collect();
assert_eq!(paths, vec!["app/main.wasm", "lib/util.rs"]);
assert!(idx.find_by_key("absent").is_empty());
assert!(idx.find_by_prefix("nope").is_empty());
assert_eq!(idx.find_by_prefix("").len(), idx.len(), "empty prefix matches all");
assert_eq!(idx.keys(), vec!["build-thing", "build-thing.abi", "coverage", "producer", "size", "vendored"]);
}
#[test]
fn no_metadata_is_not_an_empty_index() {
let absent = ArchiveMeta::NoMetadata;
let empty = ArchiveMeta::Index(MetaIndex::from_rows(Vec::new()));
assert_ne!(absent, empty, "the two states must not compare equal");
assert!(!absent.is_searchable(), "an archive with no section was not searched");
assert!(empty.is_searchable(), "a present-but-empty section WAS searched");
assert!(absent.index().is_none());
assert!(empty.index().is_some_and(MetaIndex::is_empty));
let a = absent.find_by_key("build-thing");
let e = empty.find_by_key("build-thing");
assert_eq!(a, MetaSearch::NoMetadata);
assert_eq!(e, MetaSearch::Hits(&[]));
assert!(a.hits().is_none(), "absent must not present itself as zero hits");
assert_eq!(e.hits(), Some(&[][..]), "empty IS zero hits, honestly");
assert!(!a.found_any() && !e.found_any());
}
#[test]
fn a_malformed_row_errors_rather_than_defaulting() {
use arrow::array::{BinaryArray, BooleanArray, Float64Array, Int8Array, Int64Array, StringArray};
let mk = |tag: i8, with_value: bool| {
RecordBatch::try_new(meta_schema(), vec![
Arc::new(StringArray::from(vec![Some("a")])),
Arc::new(StringArray::from(vec![Some("k")])),
Arc::new(Int8Array::from(vec![tag])),
Arc::new(StringArray::from(vec![with_value.then_some("v")])),
Arc::new(Int64Array::from(vec![None::<i64>])),
Arc::new(Float64Array::from(vec![None::<f64>])),
Arc::new(BooleanArray::from(vec![None::<bool>])),
Arc::new(BinaryArray::from(vec![None::<&[u8]>])),
])
.unwrap()
};
let mut rows = Vec::new();
assert!(
decode_meta_batch_into(&mk(0, false), &mut rows).is_err(),
"a row declaring a string with a NULL string column must error"
);
assert!(
decode_meta_batch_into(&mk(9, true), &mut rows).is_err(),
"an unknown value_type must error, not be skipped or defaulted"
);
assert!(decode_meta_batch_into(&mk(0, true), &mut rows).is_ok(), "the well-formed control decodes");
assert_eq!(rows.len(), 1, "only the well-formed row was produced");
}
}