use std::{
collections::{BTreeMap, BTreeSet},
fs,
num::NonZeroUsize,
path::{Component, Path, PathBuf},
sync::Arc,
};
use anyhow::{bail, ensure, Context, Result};
use futures::{stream, StreamExt};
use neoengram_core::{chunk_file, FileNode, Index};
use walkdir::WalkDir;
use crate::repository::{is_neoengram_dir_name, Repository};
pub(crate) async fn execute(path: PathBuf) -> Result<()> {
execute_with_options(path, false).await
}
pub(crate) async fn execute_with_options(path: PathBuf, all: bool) -> Result<()> {
if contains_repository_component(&path) {
bail!("不能把 NeoEngram 内部目录加入对象库: {}", path.display());
}
let current_dir = std::env::current_dir().context("无法确定当前工作目录")?;
let repository = Repository::discover(¤t_dir)?;
repository.ensure_no_unfinished_transactions()?;
let traversal_path = path.clone();
let traversal_current_dir = current_dir;
let repository_root = repository.root().to_path_buf();
let collection = tokio::task::spawn_blocking(move || {
collect_files(
&traversal_path,
&traversal_current_dir,
&repository_root,
all,
)
})
.await
.with_context(|| format!("目录遍历任务异常终止: {}", path.display()))??;
for input in &collection.files {
repository.validate_logical_path(&input.repository_path)?;
}
let object_store = repository.object_store();
let staging_dir = repository.staging_dir();
let concurrency = std::thread::available_parallelism().map_or(1, NonZeroUsize::get);
let tasks = collection.files.into_iter().map(|input| {
let task_object_store = Arc::clone(&object_store);
let task_staging_dir = staging_dir.clone();
async move {
let disk_path = input.disk_path;
let mut file_node = chunk_file(disk_path.clone(), task_object_store, task_staging_dir)
.await
.with_context(|| format!("无法处理文件: {}", disk_path.display()))?;
file_node.path = input.repository_path;
print_progress(&file_node);
Ok::<FileNode, anyhow::Error>(file_node)
}
});
let results: Vec<Result<FileNode>> = stream::iter(tasks)
.buffer_unordered(concurrency)
.collect()
.await;
let mut staged_files = Vec::with_capacity(results.len());
for result in results {
staged_files.push(result?);
}
if staged_files.is_empty() && !all {
return Ok(());
}
let index_repository = repository.clone();
let deletion_scope = all.then_some(collection.repository_scope);
tokio::task::spawn_blocking(move || {
update_index(&index_repository, staged_files, deletion_scope.as_deref())
})
.await
.context("index 更新任务异常终止")??;
Ok(())
}
struct InputFile {
disk_path: PathBuf,
repository_path: String,
}
struct FileCollection {
files: Vec<InputFile>,
repository_scope: String,
}
fn collect_files(
path: &Path,
current_dir: &Path,
repository_root: &Path,
allow_missing: bool,
) -> Result<FileCollection> {
let requested_path = if path.is_absolute() {
path.to_path_buf()
} else {
current_dir.join(path)
};
let metadata = match fs::symlink_metadata(&requested_path) {
Ok(metadata) => metadata,
Err(error) if error.kind() == std::io::ErrorKind::NotFound && allow_missing => {
let normalized = normalize_absolute_path(&requested_path)?;
ensure_missing_path_is_inside_repository(&normalized, repository_root)?;
return Ok(FileCollection {
files: Vec::new(),
repository_scope: repository_path(&normalized, repository_root, true)?,
});
}
Err(error) => {
return Err(error)
.with_context(|| format!("无法访问路径: {}", requested_path.display()));
}
};
ensure!(
!metadata.file_type().is_symlink(),
"add 不跟随符号链接: {}",
requested_path.display()
);
ensure!(
metadata.is_file() || metadata.is_dir(),
"路径不是普通文件或目录: {}",
requested_path.display()
);
let canonical_path = fs::canonicalize(&requested_path)
.with_context(|| format!("无法解析路径: {}", requested_path.display()))?;
ensure!(
canonical_path.starts_with(repository_root),
"不能添加仓库之外的路径: {}",
requested_path.display()
);
let repository_scope = repository_path(&canonical_path, repository_root, true)?;
let mut files = Vec::new();
if metadata.is_file() {
files.push(input_file(canonical_path, repository_root)?);
} else {
let entries = WalkDir::new(&canonical_path)
.follow_links(false)
.into_iter()
.filter_entry(|entry| !is_neoengram_dir_name(entry.file_name()));
for entry in entries {
let entry =
entry.with_context(|| format!("遍历目录失败: {}", canonical_path.display()))?;
if entry.file_type().is_file() {
files.push(input_file(entry.into_path(), repository_root)?);
}
}
}
files.sort_unstable_by(|left, right| left.repository_path.cmp(&right.repository_path));
Ok(FileCollection {
files,
repository_scope,
})
}
fn input_file(disk_path: PathBuf, repository_root: &Path) -> Result<InputFile> {
let repository_path = repository_path(&disk_path, repository_root, false)?;
Ok(InputFile {
disk_path,
repository_path,
})
}
fn repository_path(path: &Path, repository_root: &Path, allow_root: bool) -> Result<String> {
let relative = path.strip_prefix(repository_root).with_context(|| {
format!(
"文件不在仓库根目录内: {}(仓库根目录: {})",
path.display(),
repository_root.display()
)
})?;
let mut components = Vec::new();
for component in relative.components() {
let Component::Normal(name) = component else {
bail!("文件路径未规范化: {}", path.display());
};
ensure!(
!is_neoengram_dir_name(name),
"不能添加 NeoEngram 内部文件: {}",
path.display()
);
let name = name
.to_str()
.with_context(|| format!("文件路径不是有效 UTF-8: {}", path.display()))?;
components.push(name);
}
ensure!(
allow_root || !components.is_empty(),
"不能把仓库根目录当作文件添加"
);
Ok(components.join("/"))
}
fn normalize_absolute_path(path: &Path) -> Result<PathBuf> {
ensure!(path.is_absolute(), "内部错误:待规范化路径不是绝对路径");
let mut normalized = PathBuf::new();
for component in path.components() {
match component {
Component::Prefix(prefix) => normalized.push(prefix.as_os_str()),
Component::RootDir => normalized.push(component.as_os_str()),
Component::CurDir => {}
Component::ParentDir => {
ensure!(
normalized.pop(),
"路径越过了文件系统根目录: {}",
path.display()
);
}
Component::Normal(name) => normalized.push(name),
}
}
Ok(normalized)
}
fn ensure_missing_path_is_inside_repository(path: &Path, repository_root: &Path) -> Result<()> {
ensure!(
path.starts_with(repository_root),
"不能添加仓库之外的路径: {}",
path.display()
);
let relative = path
.strip_prefix(repository_root)
.context("无法计算待删除路径的仓库相对位置")?;
let mut current = repository_root.to_path_buf();
let components: Vec<_> = relative.components().collect();
for (index, component) in components.iter().enumerate() {
let Component::Normal(name) = component else {
bail!("文件路径未规范化: {}", path.display());
};
ensure!(
!is_neoengram_dir_name(name),
"不能添加 NeoEngram 内部文件: {}",
path.display()
);
current.push(name);
match fs::symlink_metadata(¤t) {
Ok(metadata) => {
ensure!(
!metadata.file_type().is_symlink(),
"路径祖先不能是符号链接: {}",
current.display()
);
if index + 1 < components.len() {
ensure!(metadata.is_dir(), "路径祖先不是目录: {}", current.display());
}
}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => break,
Err(error) => {
return Err(error).with_context(|| format!("无法检查路径: {}", current.display()));
}
}
}
Ok(())
}
fn update_index(
repository: &Repository,
staged_files: Vec<FileNode>,
deletion_scope: Option<&str>,
) -> Result<()> {
let _lock = repository.acquire_write_lock()?;
let index = repository.read_index()?;
let mut files: BTreeMap<String, FileNode> = index
.files
.into_iter()
.map(|file| (file.path.clone(), file))
.collect();
let staged_paths: BTreeSet<&str> = staged_files.iter().map(|file| file.path.as_str()).collect();
if let Some(scope) = deletion_scope {
files.retain(|path, _| {
!path_is_in_scope(path, scope) || staged_paths.contains(path.as_str())
});
}
for file in staged_files {
files.insert(file.path.clone(), file);
}
repository.write_index(&Index {
format_version: index.format_version,
files: files.into_values().collect(),
})
}
fn path_is_in_scope(path: &str, scope: &str) -> bool {
scope.is_empty()
|| path == scope
|| path
.strip_prefix(scope)
.is_some_and(|suffix| suffix.starts_with('/'))
}
fn contains_repository_component(path: &Path) -> bool {
path.components().any(
|component| matches!(component, Component::Normal(name) if is_neoengram_dir_name(name)),
)
}
fn print_progress(file_node: &FileNode) {
println!(
"Processed: {}, {} chunks created",
file_node.path,
file_node.chunks.len()
);
}