1use std::collections::HashMap;
37use std::path::Path;
38use std::sync::Arc;
39
40use anyhow::{Result, anyhow};
41use arrow::array::{
42 Array, BinaryArray, BinaryBuilder, BooleanArray, BooleanBuilder, Float64Array,
43 Float64Builder, Int8Array, Int8Builder, Int64Array, Int64Builder, StringArray, StringBuilder,
44};
45use arrow::datatypes::{DataType, Field, Schema};
46use arrow::record_batch::RecordBatch;
47use serde::{Deserialize, Serialize};
48
49use crate::index::{
50 FORMAT_VERSION_KEY, META_MODULE, ZNIPPY_FORMAT_VERSION, check_format_version,
51 read_reserved_section_bytes,
52};
53
54#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
63pub enum MetaValue {
64 Str(String),
65 I64(i64),
66 F64(f64),
67 Bool(bool),
68 Bytes(Vec<u8>),
70}
71
72impl MetaValue {
73 fn tag(&self) -> i8 {
75 match self {
76 MetaValue::Str(_) => 0,
77 MetaValue::I64(_) => 1,
78 MetaValue::F64(_) => 2,
79 MetaValue::Bool(_) => 3,
80 MetaValue::Bytes(_) => 4,
81 }
82 }
83
84 pub fn as_str(&self) -> Option<&str> {
85 match self {
86 MetaValue::Str(s) => Some(s),
87 _ => None,
88 }
89 }
90 pub fn as_i64(&self) -> Option<i64> {
91 match self {
92 MetaValue::I64(v) => Some(*v),
93 _ => None,
94 }
95 }
96 pub fn as_f64(&self) -> Option<f64> {
97 match self {
98 MetaValue::F64(v) => Some(*v),
99 _ => None,
100 }
101 }
102 pub fn as_bool(&self) -> Option<bool> {
103 match self {
104 MetaValue::Bool(v) => Some(*v),
105 _ => None,
106 }
107 }
108 pub fn as_bytes(&self) -> Option<&[u8]> {
109 match self {
110 MetaValue::Bytes(b) => Some(b),
111 _ => None,
112 }
113 }
114}
115
116impl From<&str> for MetaValue {
117 fn from(v: &str) -> Self {
118 MetaValue::Str(v.to_string())
119 }
120}
121impl From<String> for MetaValue {
122 fn from(v: String) -> Self {
123 MetaValue::Str(v)
124 }
125}
126impl From<i64> for MetaValue {
127 fn from(v: i64) -> Self {
128 MetaValue::I64(v)
129 }
130}
131impl From<f64> for MetaValue {
132 fn from(v: f64) -> Self {
133 MetaValue::F64(v)
134 }
135}
136impl From<bool> for MetaValue {
137 fn from(v: bool) -> Self {
138 MetaValue::Bool(v)
139 }
140}
141impl From<Vec<u8>> for MetaValue {
142 fn from(v: Vec<u8>) -> Self {
143 MetaValue::Bytes(v)
144 }
145}
146
147#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
150pub struct MetaEntry {
151 pub relative_path: Option<String>,
153 pub key: String,
154 pub value: MetaValue,
155}
156
157impl MetaEntry {
158 pub fn entry(relative_path: impl Into<String>, key: impl Into<String>, value: impl Into<MetaValue>) -> Self {
159 Self {
160 relative_path: Some(relative_path.into()),
161 key: key.into(),
162 value: value.into(),
163 }
164 }
165
166 pub fn archive(key: impl Into<String>, value: impl Into<MetaValue>) -> Self {
167 Self { relative_path: None, key: key.into(), value: value.into() }
168 }
169
170 pub fn path(&self) -> Option<&str> {
176 self.relative_path.as_deref()
177 }
178
179 fn sort_key(&self) -> (&str, Option<&str>) {
182 (self.key.as_str(), self.relative_path.as_deref())
183 }
184}
185
186#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
197pub struct MetaTable {
198 rows: Vec<MetaEntry>,
199}
200
201impl MetaTable {
202 pub fn new() -> Self {
203 Self::default()
204 }
205
206 pub fn from_rows(rows: Vec<MetaEntry>) -> Self {
207 Self { rows }
208 }
209
210 pub fn insert(
212 &mut self,
213 relative_path: impl Into<String>,
214 key: impl Into<String>,
215 value: impl Into<MetaValue>,
216 ) -> &mut Self {
217 self.rows.push(MetaEntry::entry(relative_path, key, value));
218 self
219 }
220
221 pub fn insert_archive(&mut self, key: impl Into<String>, value: impl Into<MetaValue>) -> &mut Self {
223 self.rows.push(MetaEntry::archive(key, value));
224 self
225 }
226
227 pub fn extend(&mut self, rows: impl IntoIterator<Item = MetaEntry>) -> &mut Self {
228 self.rows.extend(rows);
229 self
230 }
231
232 pub fn rows(&self) -> &[MetaEntry] {
233 &self.rows
234 }
235
236 pub fn len(&self) -> usize {
237 self.rows.len()
238 }
239
240 pub fn is_empty(&self) -> bool {
241 self.rows.is_empty()
242 }
243
244 fn sorted_rows(&self) -> Vec<&MetaEntry> {
247 let mut v: Vec<&MetaEntry> = self.rows.iter().collect();
248 v.sort_by(|a, b| a.sort_key().cmp(&b.sort_key()));
249 v
250 }
251}
252
253pub fn meta_schema() -> Arc<Schema> {
262 let fields = vec![
263 Field::new("relative_path", DataType::Utf8, true),
265 Field::new("key", DataType::Utf8, false),
266 Field::new("value_type", DataType::Int8, false),
267 Field::new("v_str", DataType::Utf8, true),
268 Field::new("v_i64", DataType::Int64, true),
269 Field::new("v_f64", DataType::Float64, true),
270 Field::new("v_bool", DataType::Boolean, true),
271 Field::new("v_bytes", DataType::Binary, true),
272 ];
273 let mut md = HashMap::new();
274 md.insert(FORMAT_VERSION_KEY.to_string(), ZNIPPY_FORMAT_VERSION.to_string());
275 Arc::new(Schema::new_with_metadata(fields, md))
276}
277
278pub fn build_meta_batch(table: &MetaTable) -> Result<RecordBatch> {
280 let rows = table.sorted_rows();
281 let n = rows.len();
282
283 let mut path_b = StringBuilder::with_capacity(n, n * 24);
284 let mut key_b = StringBuilder::with_capacity(n, n * 16);
285 let mut tag_b = Int8Builder::with_capacity(n);
286 let mut s_b = StringBuilder::with_capacity(n, n * 16);
287 let mut i_b = Int64Builder::with_capacity(n);
288 let mut f_b = Float64Builder::with_capacity(n);
289 let mut bo_b = BooleanBuilder::with_capacity(n);
290 let mut by_b = BinaryBuilder::with_capacity(n, n * 16);
291
292 for r in rows {
293 match &r.relative_path {
294 Some(p) => path_b.append_value(p),
295 None => path_b.append_null(),
296 }
297 key_b.append_value(&r.key);
298 tag_b.append_value(r.value.tag());
299 match &r.value {
301 MetaValue::Str(s) => {
302 s_b.append_value(s);
303 i_b.append_null();
304 f_b.append_null();
305 bo_b.append_null();
306 by_b.append_null();
307 }
308 MetaValue::I64(v) => {
309 s_b.append_null();
310 i_b.append_value(*v);
311 f_b.append_null();
312 bo_b.append_null();
313 by_b.append_null();
314 }
315 MetaValue::F64(v) => {
316 s_b.append_null();
317 i_b.append_null();
318 f_b.append_value(*v);
319 bo_b.append_null();
320 by_b.append_null();
321 }
322 MetaValue::Bool(v) => {
323 s_b.append_null();
324 i_b.append_null();
325 f_b.append_null();
326 bo_b.append_value(*v);
327 by_b.append_null();
328 }
329 MetaValue::Bytes(b) => {
330 s_b.append_null();
331 i_b.append_null();
332 f_b.append_null();
333 bo_b.append_null();
334 by_b.append_value(b);
335 }
336 }
337 }
338
339 RecordBatch::try_new(meta_schema(), vec![
340 Arc::new(path_b.finish()),
341 Arc::new(key_b.finish()),
342 Arc::new(tag_b.finish()),
343 Arc::new(s_b.finish()),
344 Arc::new(i_b.finish()),
345 Arc::new(f_b.finish()),
346 Arc::new(bo_b.finish()),
347 Arc::new(by_b.finish()),
348 ])
349 .map_err(|e| anyhow!("build meta sub-index batch: {e}"))
350}
351
352#[derive(Debug, Clone, PartialEq)]
362pub struct MetaIndex {
363 rows: Vec<MetaEntry>,
364}
365
366impl MetaIndex {
367 pub fn from_rows(mut rows: Vec<MetaEntry>) -> Self {
370 rows.sort_by(|a, b| a.sort_key().cmp(&b.sort_key()));
371 Self { rows }
372 }
373
374 pub fn len(&self) -> usize {
375 self.rows.len()
376 }
377
378 pub fn is_empty(&self) -> bool {
381 self.rows.is_empty()
382 }
383
384 pub fn iter(&self) -> std::slice::Iter<'_, MetaEntry> {
386 self.rows.iter()
387 }
388
389 pub fn find_by_key(&self, key: &str) -> &[MetaEntry] {
394 let lo = self.rows.partition_point(|r| r.key.as_str() < key);
395 let hi = self.rows.partition_point(|r| r.key.as_str() <= key);
396 &self.rows[lo..hi]
397 }
398
399 pub fn find_by_prefix(&self, prefix: &str) -> &[MetaEntry] {
402 let lo = self.rows.partition_point(|r| r.key.as_str() < prefix);
403 let hi = self.rows.partition_point(|r| r.key.as_str() < prefix || r.key.starts_with(prefix));
404 &self.rows[lo..hi]
405 }
406
407 pub fn archive_value(&self, key: &str) -> Option<&MetaValue> {
409 self.find_by_key(key)
410 .iter()
411 .find(|r| r.relative_path.is_none())
412 .map(|r| &r.value)
413 }
414
415 pub fn keys(&self) -> Vec<&str> {
417 let mut out: Vec<&str> = Vec::new();
418 for r in &self.rows {
419 if out.last() != Some(&r.key.as_str()) {
420 out.push(r.key.as_str());
421 }
422 }
423 out
424 }
425
426 pub fn to_table(&self) -> MetaTable {
429 MetaTable::from_rows(self.rows.clone())
430 }
431}
432
433#[derive(Debug, Clone, PartialEq)]
440pub enum ArchiveMeta {
441 NoMetadata,
445 Index(MetaIndex),
448}
449
450impl ArchiveMeta {
451 pub fn index(&self) -> Option<&MetaIndex> {
454 match self {
455 ArchiveMeta::NoMetadata => None,
456 ArchiveMeta::Index(i) => Some(i),
457 }
458 }
459
460 pub fn is_searchable(&self) -> bool {
462 matches!(self, ArchiveMeta::Index(_))
463 }
464
465 pub fn find_by_key(&self, key: &str) -> MetaSearch<'_> {
467 match self {
468 ArchiveMeta::NoMetadata => MetaSearch::NoMetadata,
469 ArchiveMeta::Index(i) => MetaSearch::Hits(i.find_by_key(key)),
470 }
471 }
472
473 pub fn find_by_prefix(&self, prefix: &str) -> MetaSearch<'_> {
475 match self {
476 ArchiveMeta::NoMetadata => MetaSearch::NoMetadata,
477 ArchiveMeta::Index(i) => MetaSearch::Hits(i.find_by_prefix(prefix)),
478 }
479 }
480}
481
482#[derive(Debug, Clone, PartialEq)]
488pub enum MetaSearch<'a> {
489 NoMetadata,
490 Hits(&'a [MetaEntry]),
491}
492
493impl<'a> MetaSearch<'a> {
494 pub fn hits(&self) -> Option<&'a [MetaEntry]> {
497 match self {
498 MetaSearch::NoMetadata => None,
499 MetaSearch::Hits(h) => Some(h),
500 }
501 }
502
503 pub fn found_any(&self) -> bool {
505 matches!(self, MetaSearch::Hits(h) if !h.is_empty())
506 }
507}
508
509pub fn decode_meta_section(bytes: &[u8]) -> Result<MetaIndex> {
511 use arrow::ipc::reader::StreamReader;
512
513 let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None)
514 .map_err(|e| anyhow!("meta sub-index: not a readable Arrow stream: {e}"))?;
515 check_format_version(reader.schema().metadata())?;
518
519 let mut rows = Vec::new();
520 for batch in reader {
521 let batch = batch.map_err(|e| anyhow!("meta sub-index read: {e}"))?;
522 decode_meta_batch_into(&batch, &mut rows)?;
523 }
524 Ok(MetaIndex::from_rows(rows))
525}
526
527fn col<'a, T: 'static>(batch: &'a RecordBatch, name: &str) -> Result<&'a T> {
528 batch
529 .column_by_name(name)
530 .ok_or_else(|| anyhow!("meta sub-index missing column {name:?}"))?
531 .as_any()
532 .downcast_ref::<T>()
533 .ok_or_else(|| anyhow!("meta sub-index column {name:?} has an unexpected Arrow type"))
534}
535
536fn decode_meta_batch_into(batch: &RecordBatch, out: &mut Vec<MetaEntry>) -> Result<()> {
537 let paths = col::<StringArray>(batch, "relative_path")?;
538 let keys = col::<StringArray>(batch, "key")?;
539 let tags = col::<Int8Array>(batch, "value_type")?;
540 let v_str = col::<StringArray>(batch, "v_str")?;
541 let v_i64 = col::<Int64Array>(batch, "v_i64")?;
542 let v_f64 = col::<Float64Array>(batch, "v_f64")?;
543 let v_bool = col::<BooleanArray>(batch, "v_bool")?;
544 let v_bytes = col::<BinaryArray>(batch, "v_bytes")?;
545
546 out.reserve(batch.num_rows());
547 for i in 0..batch.num_rows() {
548 let want = |present: bool, what: &str| -> Result<()> {
551 anyhow::ensure!(present, "meta row {i} declares {what} but that column is null");
552 Ok(())
553 };
554 let value = match tags.value(i) {
555 0 => {
556 want(v_str.is_valid(i), "a string value")?;
557 MetaValue::Str(v_str.value(i).to_string())
558 }
559 1 => {
560 want(v_i64.is_valid(i), "an i64 value")?;
561 MetaValue::I64(v_i64.value(i))
562 }
563 2 => {
564 want(v_f64.is_valid(i), "an f64 value")?;
565 MetaValue::F64(v_f64.value(i))
566 }
567 3 => {
568 want(v_bool.is_valid(i), "a bool value")?;
569 MetaValue::Bool(v_bool.value(i))
570 }
571 4 => {
572 want(v_bytes.is_valid(i), "a bytes value")?;
573 MetaValue::Bytes(v_bytes.value(i).to_vec())
574 }
575 other => return Err(anyhow!("meta row {i} has unknown value_type {other}")),
576 };
577 out.push(MetaEntry {
578 relative_path: paths.is_valid(i).then(|| paths.value(i).to_string()),
579 key: keys.value(i).to_string(),
580 value,
581 });
582 }
583 Ok(())
584}
585
586pub fn read_archive_meta(path: &Path) -> Result<ArchiveMeta> {
596 match read_reserved_section_bytes(path, META_MODULE)? {
597 None => Ok(ArchiveMeta::NoMetadata),
598 Some(bytes) => Ok(ArchiveMeta::Index(decode_meta_section(&bytes)?)),
599 }
600}
601
602#[cfg(test)]
603mod tests {
604 use super::*;
605
606 fn sample() -> MetaTable {
607 let mut t = MetaTable::new();
608 t.insert("app/main.wasm", "build-thing", MetaValue::Bytes(vec![0, 97, 115, 109, 1]))
609 .insert("app/main.wasm", "build-thing.abi", "wasi-p2")
610 .insert("app/main.wasm", "size", 5i64)
611 .insert("lib/util.rs", "build-thing.abi", "native")
612 .insert("lib/util.rs", "coverage", 0.87f64)
613 .insert("lib/util.rs", "vendored", false)
614 .insert_archive("producer", "znippy")
615 .insert_archive("build-thing", MetaValue::Bytes(vec![1, 2, 3]));
616 t
617 }
618
619 #[test]
623 fn every_value_type_and_both_scopes_round_trip() {
624 let t = sample();
625 let batch = build_meta_batch(&t).unwrap();
626 assert_eq!(batch.num_rows(), t.len());
627
628 let mut rows = Vec::new();
629 decode_meta_batch_into(&batch, &mut rows).unwrap();
630 let idx = MetaIndex::from_rows(rows);
631 assert_eq!(idx.len(), t.len());
632
633 let bt = idx.find_by_key("build-thing");
635 assert_eq!(bt.len(), 2, "one entry-scoped + one archive-scoped");
636 assert_eq!(bt[0].path(), None, "archive-scoped sorts first (NULL path)");
637 assert_eq!(bt[0].value.as_bytes(), Some(&[1u8, 2, 3][..]));
638 assert_eq!(bt[1].path(), Some("app/main.wasm"));
639 assert_eq!(bt[1].value.as_bytes(), Some(&[0u8, 97, 115, 109, 1][..]));
640
641 assert_eq!(idx.find_by_key("size")[0].value.as_i64(), Some(5));
642 assert_eq!(idx.find_by_key("coverage")[0].value.as_f64(), Some(0.87));
643 assert_eq!(idx.find_by_key("vendored")[0].value.as_bool(), Some(false));
644 assert_eq!(idx.archive_value("producer").and_then(MetaValue::as_str), Some("znippy"));
645
646 let ordered: Vec<_> = idx.iter().map(|r| (r.key.as_str(), r.path())).collect();
648 let mut want = ordered.clone();
649 want.sort();
650 assert_eq!(ordered, want, "rows must be stored key-major and sorted");
651
652 let json = serde_json::to_string(t.rows()).unwrap();
654 let back: Vec<MetaEntry> = serde_json::from_str(&json).unwrap();
655 assert_eq!(back, t.rows());
656 }
657
658 #[test]
662 fn key_and_prefix_search_return_exactly_the_matching_rows() {
663 let idx = MetaIndex::from_rows(sample().rows().to_vec());
664
665 let exact = idx.find_by_key("build-thing");
666 assert_eq!(exact.len(), 2, "exact key must NOT sweep in `build-thing.abi`");
667 assert!(exact.iter().all(|r| r.key == "build-thing"));
668
669 let pre = idx.find_by_prefix("build-thing");
670 assert_eq!(pre.len(), 4, "prefix picks up build-thing + build-thing.abi ×2");
671 assert!(pre.iter().all(|r| r.key.starts_with("build-thing")));
672
673 let paths: Vec<_> = idx.find_by_key("build-thing.abi").iter().filter_map(|r| r.path()).collect();
675 assert_eq!(paths, vec!["app/main.wasm", "lib/util.rs"]);
676
677 assert!(idx.find_by_key("absent").is_empty());
678 assert!(idx.find_by_prefix("nope").is_empty());
679 assert_eq!(idx.find_by_prefix("").len(), idx.len(), "empty prefix matches all");
680 assert_eq!(idx.keys(), vec!["build-thing", "build-thing.abi", "coverage", "producer", "size", "vendored"]);
681 }
682
683 #[test]
687 fn no_metadata_is_not_an_empty_index() {
688 let absent = ArchiveMeta::NoMetadata;
689 let empty = ArchiveMeta::Index(MetaIndex::from_rows(Vec::new()));
690
691 assert_ne!(absent, empty, "the two states must not compare equal");
692 assert!(!absent.is_searchable(), "an archive with no section was not searched");
693 assert!(empty.is_searchable(), "a present-but-empty section WAS searched");
694 assert!(absent.index().is_none());
695 assert!(empty.index().is_some_and(MetaIndex::is_empty));
696
697 let a = absent.find_by_key("build-thing");
699 let e = empty.find_by_key("build-thing");
700 assert_eq!(a, MetaSearch::NoMetadata);
701 assert_eq!(e, MetaSearch::Hits(&[]));
702 assert!(a.hits().is_none(), "absent must not present itself as zero hits");
703 assert_eq!(e.hits(), Some(&[][..]), "empty IS zero hits, honestly");
704 assert!(!a.found_any() && !e.found_any());
705 }
706
707 #[test]
710 fn a_malformed_row_errors_rather_than_defaulting() {
711 use arrow::array::{BinaryArray, BooleanArray, Float64Array, Int8Array, Int64Array, StringArray};
712
713 let mk = |tag: i8, with_value: bool| {
714 RecordBatch::try_new(meta_schema(), vec![
715 Arc::new(StringArray::from(vec![Some("a")])),
716 Arc::new(StringArray::from(vec![Some("k")])),
717 Arc::new(Int8Array::from(vec![tag])),
718 Arc::new(StringArray::from(vec![with_value.then_some("v")])),
719 Arc::new(Int64Array::from(vec![None::<i64>])),
720 Arc::new(Float64Array::from(vec![None::<f64>])),
721 Arc::new(BooleanArray::from(vec![None::<bool>])),
722 Arc::new(BinaryArray::from(vec![None::<&[u8]>])),
723 ])
724 .unwrap()
725 };
726
727 let mut rows = Vec::new();
728 assert!(
729 decode_meta_batch_into(&mk(0, false), &mut rows).is_err(),
730 "a row declaring a string with a NULL string column must error"
731 );
732 assert!(
733 decode_meta_batch_into(&mk(9, true), &mut rows).is_err(),
734 "an unknown value_type must error, not be skipped or defaulted"
735 );
736 assert!(decode_meta_batch_into(&mk(0, true), &mut rows).is_ok(), "the well-formed control decodes");
737 assert_eq!(rows.len(), 1, "only the well-formed row was produced");
738 }
739}