From c975d02ef0bf503785b15751bf3c226fd957538c Mon Sep 17 00:00:00 2001
From: mrbeanc <mrbeanc@qq.com>
Date: Sat, 23 Dec 2023 17:47:57 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E6=AD=A3restore=5Findex()?=
 =?UTF-8?q?=E8=A1=8C=E4=B8=BA=EF=BC=8C=E4=BB=8Eindex=E7=AD=9B=E9=80=89?=
 =?UTF-8?q?=E6=96=87=E4=BB=B6=E8=80=8C=E9=9D=9E=E5=B7=A5=E4=BD=9C=E5=8C=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/commands/restore.rs | 29 +++++++++++++++++++++++------
 src/commands/switch.rs  |  2 +-
 src/utils/util.rs       |  6 +++++-
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/src/commands/restore.rs b/src/commands/restore.rs
index 14254dd..434553c 100644
--- a/src/commands/restore.rs
+++ b/src/commands/restore.rs
@@ -16,7 +16,10 @@ use crate::{
 };
 
 /// 统计[工作区]中的dirs文件夹中，相对于target_blobs已删除的文件
-fn get_worktree_deleted_files_in_dirs(dirs: &Vec<PathBuf>, target_blobs: &HashMap<PathBuf, Hash>) -> HashSet<PathBuf> {
+fn get_worktree_deleted_files_in_dirs(
+    dirs: &HashSet<PathBuf>,
+    target_blobs: &HashMap<PathBuf, Hash>,
+) -> HashSet<PathBuf> {
     target_blobs //统计所有目录中(包括None & '.')，删除的文件
         .iter()
         .filter(|(path, _)| {
@@ -38,7 +41,7 @@ fn get_worktree_deleted_files_in_dirs(dirs: &Vec<PathBuf>, target_blobs: &HashMa
 /// 统计[暂存区index]中相对于target_blobs已删除的文件，且包含在指定dirs内
 fn get_index_deleted_files_in_dirs(
     index: &Index,
-    dirs: &Vec<PathBuf>,
+    dirs: &HashSet<PathBuf>,
     target_blobs: &HashMap<PathBuf, Hash>,
 ) -> HashSet<PathBuf> {
     target_blobs //统计index中相对target已删除的文件，且包含在指定dir内
@@ -79,13 +82,13 @@ fn preprocess_blobs(blobs: &Vec<(PathBuf, Hash)>) -> HashMap<PathBuf, Hash> {
 
 /** 根据filter restore workdir */
 pub fn restore_worktree(filter: Option<&Vec<PathBuf>>, target_blobs: &Vec<(PathBuf, Hash)>) {
-    let paths = preprocess_filters(filter); //预处理filter 将None转化为workdir
+    let input_paths = preprocess_filters(filter); //预处理filter 将None转化为workdir
     let target_blobs = preprocess_blobs(target_blobs); //预处理target_blobs 转化为绝对路径HashMap
 
-    let dirs = util::filter_dirs(&paths); //统计所有目录
+    let dirs = util::filter_dirs(&input_paths); //统计所有目录
     let deleted_files = get_worktree_deleted_files_in_dirs(&dirs, &target_blobs); //统计所有目录中已删除的文件
 
-    let mut file_paths = util::integrate_paths(&paths); //整合存在的文件（绝对路径）
+    let mut file_paths = util::integrate_paths(&input_paths); //整合存在的文件（绝对路径）
     file_paths.extend(deleted_files); //已删除的文件
 
     let index = Index::new();
@@ -131,7 +134,21 @@ pub fn restore_index(filter: Option<&Vec<PathBuf>>, target_blobs: &Vec<(PathBuf,
     let dirs = util::filter_dirs(&input_paths); //统计所有目录
     let deleted_files_index = get_index_deleted_files_in_dirs(&index, &dirs, &target_blobs); //统计所有目录中已删除的文件
 
-    let mut file_paths = util::integrate_paths(&input_paths); //整合存在的文件（绝对路径）
+    // 1. 获取输入中的[文件路径]
+    let mut file_paths: HashSet<PathBuf> = util::filter_files(&input_paths)
+        .iter()
+        .map(|path| util::get_absolute_path(path))
+        .collect();
+    // 2.获取index中包含于dirs的文件（使用dirs进行筛选）
+    for index_file in index.get_tracked_files() {
+        for dir in &dirs {
+            if util::is_parent_dir(&index_file, dir) {
+                //需要包含在指定dir内
+                file_paths.insert(index_file.clone());
+            }
+        }
+    }
+    // 3.补充index中已删除的文件（相较于target_blobs）
     file_paths.extend(deleted_files_index); //已删除的文件
 
     for path in &file_paths {
diff --git a/src/commands/switch.rs b/src/commands/switch.rs
index 5ba31f8..eab2578 100644
--- a/src/commands/switch.rs
+++ b/src/commands/switch.rs
@@ -91,7 +91,7 @@ mod test {
 
     use super::*;
     #[test]
-    #[ignore] // TODO 等待restore实现后再测试
+    // #[ignore] // TODO 等待restore实现后再测试
     fn test_switch() {
         util::setup_test_with_clean_mit();
         util::list_workdir_files().iter().for_each(|f| {
diff --git a/src/utils/util.rs b/src/utils/util.rs
index 695c5cc..49386e7 100644
--- a/src/utils/util.rs
+++ b/src/utils/util.rs
@@ -171,10 +171,14 @@ pub fn format_time(time: &std::time::SystemTime) -> String {
 }
 
 /// 过滤出路径数组中的目录
-pub fn filter_dirs(paths: &Vec<PathBuf>) -> Vec<PathBuf> {
+pub fn filter_dirs(paths: &Vec<PathBuf>) -> HashSet<PathBuf> {
     paths.iter().filter(|path| path.is_dir()).cloned().collect()
 }
 
+pub fn filter_files(paths: &Vec<PathBuf>) -> HashSet<PathBuf> {
+    paths.iter().filter(|path| path.is_file()).cloned().collect()
+}
+
 /// 将路径中的分隔符统一为当前系统的分隔符
 fn unify_path_separator(path: &Path) -> PathBuf {
     #[cfg(windows)]