use std::borrow::Cow;
use std::path::{Path, PathBuf};
use serde::Deserialize;
use crate::extract::SCHEMA_VER;
use crate::extract::{FileMapL1, FileMapL2};
use crate::hashing::{self, Hash};
use crate::store::{Store, StoreError, check_schema};
#[cfg(feature = "code-search")]
use crate::store_blob_codec::read_u16;
use crate::store_blob_codec::{
BlobKind, CompressedPayload, checked_u32_len, compress_payload, corrupt_blob, decompress_payload, encode_prefix,
envelope_prefix, read_u32,
};
const SINGLE_HEADER_LEN: usize = 16;
const FILEMAP_HEADER_LEN: usize = 24;
#[cfg(feature = "code-search")]
const CHUNK_HEADER_LEN: usize = 28;
#[cfg(feature = "code-search")]
const MAX_PEEK_ITEMS: usize = 10_000_000;
struct FileMapEnvelope<'a> {
l1: CompressedPayload<'a>,
l2: CompressedPayload<'a>,
}
#[cfg(feature = "code-search")]
struct ChunkEnvelopePeek {
schema_ver: u16,
embedding_dim: u16,
embedding_model: String,
config_digest: String,
chunk_count: usize,
embedding_count: usize,
}
#[derive(Deserialize)]
struct BlobSchemaPeek {
schema_ver: u16,
}
pub(crate) fn read_if_exists(path: &Path) -> Result<Option<Vec<u8>>, StoreError> {
match std::fs::read(path) {
Ok(bytes) => Ok(Some(bytes)),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(source) => Err(StoreError::Io {
path: path.to_path_buf(),
source,
}),
}
}
pub fn read_l1_blob_in(blobs_dir: &Path, hash_hex: &str) -> Result<Option<FileMapL1>, StoreError> {
let path = blobs_dir.join(format!("{hash_hex}.fm.msgpack"));
let Some(bytes) = read_if_exists(&path)? else {
return Ok(None);
};
let map = parse_filemap_l1(&path, &bytes)?;
check_schema(map.schema_ver)?;
Ok(Some(map))
}
pub fn read_l2_blob_in(blobs_dir: &Path, hash_hex: &str) -> Result<Option<FileMapL2>, StoreError> {
let path = blobs_dir.join(format!("{hash_hex}.fm.msgpack"));
let Some(bytes) = read_if_exists(&path)? else {
return Ok(None);
};
match parse_filemap_l2(&path, &bytes)? {
Some(map) => {
check_schema(map.schema_ver)?;
Ok(Some(map))
}
None => Ok(None),
}
}
fn frame_slices(bytes: &[u8]) -> Option<(&[u8], &[u8])> {
let header: [u8; 4] = bytes.get(0..4)?.try_into().ok()?;
let l1_len = u32::from_le_bytes(header) as usize;
let rest = bytes.get(4..)?;
let l1 = rest.get(..l1_len)?;
let l2 = &rest[l1_len..];
Some((l1, l2))
}
fn filemap_envelope<'a>(path: &Path, bytes: &'a [u8]) -> Result<FileMapEnvelope<'a>, StoreError> {
let prefix = envelope_prefix(path, bytes)?.ok_or_else(|| corrupt_blob(path))?;
if prefix.kind != BlobKind::FileMap {
return Err(corrupt_blob(path));
}
let l1_raw_len = read_u32(path, bytes, 8)?;
let l1_stored_len = read_u32(path, bytes, 12)?;
let l2_raw_len = read_u32(path, bytes, 16)?;
let l2_stored_len = read_u32(path, bytes, 20)?;
let l1_end = FILEMAP_HEADER_LEN
.checked_add(l1_stored_len)
.ok_or_else(|| corrupt_blob(path))?;
let l2_end = l1_end.checked_add(l2_stored_len).ok_or_else(|| corrupt_blob(path))?;
if l2_end != bytes.len() {
return Err(corrupt_blob(path));
}
Ok(FileMapEnvelope {
l1: CompressedPayload {
uncompressed_len: l1_raw_len,
bytes: &bytes[FILEMAP_HEADER_LEN..l1_end],
},
l2: CompressedPayload {
uncompressed_len: l2_raw_len,
bytes: &bytes[l1_end..l2_end],
},
})
}
pub(crate) fn frame_filemap(l1: &FileMapL1, l2: Option<&FileMapL2>) -> Result<Vec<u8>, StoreError> {
let l1_bytes = rmp_serde::to_vec_named(l1)?;
let l2_bytes = match l2 {
Some(map) => rmp_serde::to_vec_named(map)?,
None => Vec::new(),
};
let l1_compressed = compress_payload(&l1_bytes)?;
let l2_compressed = if l2_bytes.is_empty() {
Vec::new()
} else {
compress_payload(&l2_bytes)?
};
let mut out = Vec::with_capacity(FILEMAP_HEADER_LEN + l1_compressed.len() + l2_compressed.len());
encode_prefix(BlobKind::FileMap, l1.schema_ver, &mut out);
out.extend_from_slice(&checked_u32_len(l1_bytes.len())?.to_le_bytes());
out.extend_from_slice(&checked_u32_len(l1_compressed.len())?.to_le_bytes());
out.extend_from_slice(&checked_u32_len(l2_bytes.len())?.to_le_bytes());
out.extend_from_slice(&checked_u32_len(l2_compressed.len())?.to_le_bytes());
out.extend_from_slice(&l1_compressed);
out.extend_from_slice(&l2_compressed);
Ok(out)
}
pub(crate) fn parse_filemap_l1(path: &Path, bytes: &[u8]) -> Result<FileMapL1, StoreError> {
let l1 = match envelope_prefix(path, bytes)? {
Some(_) => Cow::Owned(decompress_payload(path, filemap_envelope(path, bytes)?.l1)?),
None => Cow::Borrowed(frame_slices(bytes).ok_or_else(|| corrupt_blob(path))?.0),
};
Ok(rmp_serde::from_slice(&l1)?)
}
pub(crate) fn parse_filemap_l2(path: &Path, bytes: &[u8]) -> Result<Option<FileMapL2>, StoreError> {
let l2 = match envelope_prefix(path, bytes)? {
Some(_) => {
let envelope = filemap_envelope(path, bytes)?;
if envelope.l2.uncompressed_len == 0 && envelope.l2.bytes.is_empty() {
return Ok(None);
}
Cow::Owned(decompress_payload(path, envelope.l2)?)
}
None => Cow::Borrowed(frame_slices(bytes).ok_or_else(|| corrupt_blob(path))?.1),
};
if l2.is_empty() {
Ok(None)
} else {
Ok(Some(rmp_serde::from_slice(&l2)?))
}
}
fn encode_single_blob(bytes: &[u8], schema_ver: u16) -> Result<Vec<u8>, StoreError> {
let compressed = compress_payload(bytes)?;
let mut out = Vec::with_capacity(SINGLE_HEADER_LEN + compressed.len());
encode_prefix(BlobKind::Single, schema_ver, &mut out);
out.extend_from_slice(&checked_u32_len(bytes.len())?.to_le_bytes());
out.extend_from_slice(&checked_u32_len(compressed.len())?.to_le_bytes());
out.extend_from_slice(&compressed);
Ok(out)
}
fn single_payload<'a>(path: &Path, bytes: &'a [u8], expected: BlobKind) -> Result<CompressedPayload<'a>, StoreError> {
let prefix = envelope_prefix(path, bytes)?.ok_or_else(|| corrupt_blob(path))?;
if prefix.kind != expected {
return Err(corrupt_blob(path));
}
let uncompressed_len = read_u32(path, bytes, 8)?;
let stored_len = read_u32(path, bytes, 12)?;
let payload_end = SINGLE_HEADER_LEN
.checked_add(stored_len)
.ok_or_else(|| corrupt_blob(path))?;
if payload_end != bytes.len() {
return Err(corrupt_blob(path));
}
Ok(CompressedPayload {
uncompressed_len,
bytes: &bytes[SINGLE_HEADER_LEN..payload_end],
})
}
fn decode_single_or_legacy<'a>(path: &Path, bytes: &'a [u8], expected: BlobKind) -> Result<Cow<'a, [u8]>, StoreError> {
match envelope_prefix(path, bytes)? {
Some(_) => Ok(Cow::Owned(decompress_payload(
path,
single_payload(path, bytes, expected)?,
)?)),
None => Ok(Cow::Borrowed(bytes)),
}
}
fn is_current_single_blob(path: &Path) -> bool {
let Ok(bytes) = std::fs::read(path) else {
return false;
};
let Ok(Some(prefix)) = envelope_prefix(path, &bytes) else {
return false;
};
if prefix.kind != BlobKind::Single || prefix.schema_ver != SCHEMA_VER {
return false;
}
let Ok(decoded) = decode_single_or_legacy(path, &bytes, BlobKind::Single) else {
return false;
};
rmp_serde::from_slice::<BlobSchemaPeek>(&decoded).is_ok_and(|peek| peek.schema_ver == SCHEMA_VER)
}
fn is_current_filemap_blob(path: &Path) -> bool {
let Ok(bytes) = std::fs::read(path) else {
return false;
};
let Ok(Some(prefix)) = envelope_prefix(path, &bytes) else {
return false;
};
if prefix.kind != BlobKind::FileMap || prefix.schema_ver != SCHEMA_VER {
return false;
}
let Ok(l1) = parse_filemap_l1(path, &bytes) else {
return false;
};
if l1.schema_ver != SCHEMA_VER || l1.extract_epoch < crate::extract::EXTRACT_EPOCH {
return false;
}
match parse_filemap_l2(path, &bytes) {
Ok(Some(l2)) => l2.schema_ver == SCHEMA_VER,
Ok(None) => true,
Err(_) => false,
}
}
#[cfg(feature = "code-search")]
fn encode_chunk_blob(blob: &crate::chunk::CodeChunkBlob) -> Result<Vec<u8>, StoreError> {
let bytes = rmp_serde::to_vec_named(blob)?;
let compressed = compress_payload(&bytes)?;
let model = blob.embedding_model.as_bytes();
let model_len = u16::try_from(model.len()).map_err(|_| StoreError::BlobTooLarge)?;
let digest = chunk_digest_region(&blob.config_digest);
let mut out = Vec::with_capacity(CHUNK_HEADER_LEN + model.len() + digest.len() + compressed.len());
encode_prefix(BlobKind::Chunk, blob.schema_ver, &mut out);
out.extend_from_slice(&checked_u32_len(bytes.len())?.to_le_bytes());
out.extend_from_slice(&checked_u32_len(compressed.len())?.to_le_bytes());
out.extend_from_slice(&blob.embedding_dim.to_le_bytes());
out.extend_from_slice(&model_len.to_le_bytes());
out.extend_from_slice(&checked_u32_len(blob.chunks.len())?.to_le_bytes());
out.extend_from_slice(&checked_u32_len(blob.embeddings.len())?.to_le_bytes());
out.extend_from_slice(model);
out.extend_from_slice(digest);
out.extend_from_slice(&compressed);
Ok(out)
}
#[cfg(feature = "code-search")]
fn chunk_digest_region(digest: &str) -> &[u8] {
if digest.len() == crate::config::rules::DIGEST_LEN {
digest.as_bytes()
} else {
&[]
}
}
#[cfg(feature = "code-search")]
struct ChunkHeader {
peek: ChunkEnvelopePeek,
uncompressed_len: usize,
payload_start: usize,
payload_end: usize,
}
#[cfg(feature = "code-search")]
fn chunk_header(path: &Path, bytes: &[u8], total_len: usize) -> Result<ChunkHeader, StoreError> {
let prefix = envelope_prefix(path, bytes)?.ok_or_else(|| corrupt_blob(path))?;
if prefix.kind != BlobKind::Chunk {
return Err(corrupt_blob(path));
}
let uncompressed_len = read_u32(path, bytes, 8)?;
let stored_len = read_u32(path, bytes, 12)?;
let embedding_dim = read_u16(path, bytes, 16)?;
let model_len = read_u16(path, bytes, 18)? as usize;
let chunk_count = read_u32(path, bytes, 20)?;
let embedding_count = read_u32(path, bytes, 24)?;
if chunk_count > MAX_PEEK_ITEMS || embedding_count > MAX_PEEK_ITEMS {
return Err(corrupt_blob(path));
}
let model_end = CHUNK_HEADER_LEN
.checked_add(model_len)
.ok_or_else(|| corrupt_blob(path))?;
let legacy_end = model_end.checked_add(stored_len).ok_or_else(|| corrupt_blob(path))?;
let digest_len = if legacy_end == total_len {
0
} else {
crate::config::rules::DIGEST_LEN
};
let digest_end = model_end + digest_len;
let payload_end = digest_end.checked_add(stored_len).ok_or_else(|| corrupt_blob(path))?;
if payload_end != total_len {
return Err(corrupt_blob(path));
}
let embedding_model = std::str::from_utf8(
bytes
.get(CHUNK_HEADER_LEN..model_end)
.ok_or_else(|| corrupt_blob(path))?,
)
.map_err(|_| corrupt_blob(path))?
.to_string();
let config_digest = std::str::from_utf8(bytes.get(model_end..digest_end).ok_or_else(|| corrupt_blob(path))?)
.map_err(|_| corrupt_blob(path))?
.to_string();
Ok(ChunkHeader {
peek: ChunkEnvelopePeek {
schema_ver: prefix.schema_ver,
embedding_dim,
embedding_model,
config_digest,
chunk_count,
embedding_count,
},
uncompressed_len,
payload_start: digest_end,
payload_end,
})
}
#[cfg(feature = "code-search")]
fn chunk_payload<'a>(path: &Path, bytes: &'a [u8]) -> Result<CompressedPayload<'a>, StoreError> {
let header = chunk_header(path, bytes, bytes.len())?;
Ok(CompressedPayload {
uncompressed_len: header.uncompressed_len,
bytes: &bytes[header.payload_start..header.payload_end],
})
}
#[cfg(feature = "code-search")]
fn read_chunk_header_prefix(path: &Path) -> Result<Option<(Vec<u8>, usize, bool)>, StoreError> {
use std::io::Read;
let io_err = |source| StoreError::Io {
path: path.to_path_buf(),
source,
};
let mut file = match std::fs::File::open(path) {
Ok(f) => f,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(e) => return Err(io_err(e)),
};
let total_len = usize::try_from(file.metadata().map_err(io_err)?.len()).map_err(|_| corrupt_blob(path))?;
let mut head = vec![0u8; CHUNK_HEADER_LEN.min(total_len)];
file.read_exact(&mut head).map_err(io_err)?;
if envelope_prefix(path, &head)?.is_none() {
file.read_to_end(&mut head).map_err(io_err)?;
return Ok(Some((head, total_len, false)));
}
let model_len = read_u16(path, &head, 18)? as usize;
let want = (CHUNK_HEADER_LEN + model_len + crate::config::rules::DIGEST_LEN).min(total_len);
let have = head.len();
head.resize(want, 0);
file.read_exact(&mut head[have..]).map_err(io_err)?;
Ok(Some((head, total_len, true)))
}
#[cfg(feature = "code-search")]
fn decode_chunk_or_legacy<'a>(path: &Path, bytes: &'a [u8]) -> Result<Cow<'a, [u8]>, StoreError> {
match envelope_prefix(path, bytes)? {
Some(_) => Ok(Cow::Owned(decompress_payload(path, chunk_payload(path, bytes)?)?)),
None => Ok(Cow::Borrowed(bytes)),
}
}
#[cfg(feature = "code-search")]
fn public_chunk_peek(header: ChunkEnvelopePeek) -> crate::chunk::CodeChunkBlobPeek {
crate::chunk::CodeChunkBlobPeek {
schema_ver: header.schema_ver,
embedding_dim: header.embedding_dim,
embedding_model: header.embedding_model,
config_digest: header.config_digest,
chunks: (0..header.chunk_count).map(|_| serde::de::IgnoredAny).collect(),
embeddings: (0..header.embedding_count).map(|_| serde::de::IgnoredAny).collect(),
}
}
thread_local! {
static TMP_SUFFIX: String = format!(
"{}.{:?}.tmp",
std::process::id(),
std::thread::current().id()
);
}
pub(crate) fn write_bytes_atomic(path: PathBuf, bytes: &[u8]) -> Result<(), StoreError> {
use std::fs::OpenOptions;
use std::io::Write;
let tmp = TMP_SUFFIX.with(|suffix| path.with_extension(format!("msgpack.{suffix}")));
{
let mut f = OpenOptions::new()
.create(true)
.write(true)
.truncate(true)
.open(&tmp)
.map_err(|source| StoreError::Io {
path: tmp.clone(),
source,
})?;
f.write_all(bytes).map_err(|source| StoreError::Io {
path: tmp.clone(),
source,
})?;
}
if let Err(source) = std::fs::rename(&tmp, &path) {
let _ = std::fs::remove_file(&tmp);
return Err(StoreError::Io { path, source });
}
Ok(())
}
pub(crate) fn write_blob<T: serde::Serialize>(path: PathBuf, value: &T) -> Result<(), StoreError> {
if is_current_single_blob(&path) {
return Ok(());
}
let bytes = rmp_serde::to_vec_named(value)?;
let encoded = encode_single_blob(&bytes, SCHEMA_VER)?;
write_bytes_atomic(path, &encoded)
}
#[cfg(feature = "documents")]
pub(crate) fn write_blob_overwrite<T: serde::Serialize>(path: PathBuf, value: &T) -> Result<(), StoreError> {
let bytes = rmp_serde::to_vec_named(value)?;
let encoded = encode_single_blob(&bytes, SCHEMA_VER)?;
write_bytes_atomic(path, &encoded)
}
impl Store {
pub fn blob_path_fm(&self, hash: &Hash) -> PathBuf {
let buf = hashing::hex_buf(hash);
self.blob_path_fm_hex(hashing::hex_str(&buf))
}
pub fn blob_path_fm_hex(&self, hash_hex: &str) -> PathBuf {
self.blobs_dir.join(format!("{hash_hex}.fm.msgpack"))
}
#[cfg(feature = "documents")]
pub fn blob_path_doc(&self, hash: &Hash) -> PathBuf {
let buf = hashing::hex_buf(hash);
self.blob_path_doc_hex(hashing::hex_str(&buf))
}
#[cfg(feature = "documents")]
pub fn blob_path_doc_hex(&self, hash_hex: &str) -> PathBuf {
self.blobs_dir.join(format!("{hash_hex}.doc.msgpack"))
}
pub fn read_l1_by_hex(&self, hash_hex: &str) -> Result<Option<FileMapL1>, StoreError> {
read_l1_blob_in(&self.blobs_dir, hash_hex)
}
pub fn read_l2_by_hex(&self, hash_hex: &str) -> Result<Option<FileMapL2>, StoreError> {
read_l2_blob_in(&self.blobs_dir, hash_hex)
}
pub fn write_filemap_hex(&self, hash_hex: &str, l1: &FileMapL1, l2: Option<&FileMapL2>) -> Result<(), StoreError> {
if is_current_filemap_blob(&self.blob_path_fm_hex(hash_hex)) {
return Ok(());
}
self.overwrite_filemap_hex(hash_hex, l1, l2)
}
pub fn overwrite_filemap_hex(
&self,
hash_hex: &str,
l1: &FileMapL1,
l2: Option<&FileMapL2>,
) -> Result<(), StoreError> {
let bytes = frame_filemap(l1, l2)?;
write_bytes_atomic(self.blob_path_fm_hex(hash_hex), &bytes)
}
#[cfg(feature = "documents")]
pub fn write_doc(&self, hash: &Hash, map: &crate::extract::doc::FileMapDoc) -> Result<(), StoreError> {
write_blob_overwrite(self.blob_path_doc(hash), map)
}
#[cfg(feature = "documents")]
pub fn read_doc_by_hex(&self, hash_hex: &str) -> Result<Option<crate::extract::doc::FileMapDoc>, StoreError> {
let path = self.blob_path_doc_hex(hash_hex);
let Some(bytes) = read_if_exists(&path)? else {
return Ok(None);
};
let decoded = decode_single_or_legacy(&path, &bytes, BlobKind::Single)?;
let map: crate::extract::doc::FileMapDoc = rmp_serde::from_slice(&decoded)?;
check_schema(map.schema_ver)?;
Ok(Some(map))
}
#[cfg(feature = "documents")]
pub fn touch_doc_blob(&self, hash_hex: &str) {
let path = self.blob_path_doc_hex(hash_hex);
if let Err(error) = std::fs::File::options()
.write(true)
.open(&path)
.and_then(|file| file.set_modified(std::time::SystemTime::now()))
{
tracing::debug!(%error, path = %path.display(), "touch_doc_blob: refreshing mtime failed");
}
}
pub fn blob_path_rref_hex(&self, hash_hex: &str) -> PathBuf {
self.blobs_dir.join(format!("{hash_hex}.rref.msgpack"))
}
pub fn write_resolved_hex(
&self,
hash_hex: &str,
refs: &crate::intel::model::FileResolvedRefs,
) -> Result<(), StoreError> {
write_blob(self.blob_path_rref_hex(hash_hex), refs)
}
pub fn read_resolved_by_hex(
&self,
hash_hex: &str,
) -> Result<Option<crate::intel::model::FileResolvedRefs>, StoreError> {
let path = self.blob_path_rref_hex(hash_hex);
let Some(bytes) = read_if_exists(&path)? else {
return Ok(None);
};
let decoded = decode_single_or_legacy(&path, &bytes, BlobKind::Single)?;
let refs: crate::intel::model::FileResolvedRefs = rmp_serde::from_slice(&decoded)?;
check_schema(refs.schema_ver)?;
Ok(Some(refs))
}
#[cfg(feature = "code-search")]
pub fn blob_path_chunk_hex(&self, hash_hex: &str) -> PathBuf {
self.blobs_dir.join(format!("{hash_hex}.chunk.msgpack"))
}
#[cfg(feature = "code-search")]
pub fn write_chunks_hex(&self, hash_hex: &str, blob: &crate::chunk::CodeChunkBlob) -> Result<(), StoreError> {
let encoded = encode_chunk_blob(blob)?;
write_bytes_atomic(self.blob_path_chunk_hex(hash_hex), &encoded)
}
#[cfg(feature = "code-search")]
pub fn read_chunks_by_hex(&self, hash_hex: &str) -> Result<Option<crate::chunk::CodeChunkBlob>, StoreError> {
let path = self.blob_path_chunk_hex(hash_hex);
let Some(bytes) = read_if_exists(&path)? else {
return Ok(None);
};
let decoded = decode_chunk_or_legacy(&path, &bytes)?;
let blob: crate::chunk::CodeChunkBlob = rmp_serde::from_slice(&decoded)?;
check_schema(blob.schema_ver)?;
Ok(Some(blob))
}
#[cfg(feature = "code-search")]
pub fn peek_chunk_state(&self, hash_hex: &str) -> Result<Option<crate::chunk::CodeChunkBlobPeek>, StoreError> {
let path = self.blob_path_chunk_hex(hash_hex);
let Some((bytes, total_len, enveloped)) = read_chunk_header_prefix(&path)? else {
return Ok(None);
};
let peek = if enveloped {
public_chunk_peek(chunk_header(&path, &bytes, total_len)?.peek)
} else {
rmp_serde::from_slice(&bytes)?
};
check_schema(peek.schema_ver)?;
Ok(Some(peek))
}
}
#[cfg(test)]
#[path = "store_blob_tests.rs"]
mod tests;