From 41a17634b76e4c3b9438bb6a8dd5bb2bf8354c2d Mon Sep 17 00:00:00 2001
From: mrbeanc <mrbeanc@qq.com>
Date: Sat, 23 Dec 2023 12:52:21 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E9=87=8D=E6=9E=84restore=5Fworktree()?=
 =?UTF-8?q?=E9=80=BB=E8=BE=91=EF=BC=8C=E6=AD=A3=E7=A1=AE=E6=81=A2=E5=A4=8D?=
 =?UTF-8?q?dir=E4=B8=AD[=E5=B7=B2=E5=88=A0=E9=99=A4=E6=96=87=E4=BB=B6]?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/commands/restore.rs | 75 ++++++++++++++++++++++++++---------------
 1 file changed, 48 insertions(+), 27 deletions(-)

diff --git a/src/commands/restore.rs b/src/commands/restore.rs
index 1762911..1b914ee 100644
--- a/src/commands/restore.rs
+++ b/src/commands/restore.rs
@@ -1,5 +1,7 @@
+use std::collections::{HashMap, HashSet};
 use std::{fs, path::PathBuf};
 
+use crate::utils::util::{get_absolute_path, list_files};
 use crate::{
     head,
     models::{commit::Commit, index::Index, object::Hash},
@@ -9,46 +11,65 @@ use crate::{
 
 /** 根据filter restore workdir */
 pub fn restore_worktree(filter: Option<&Vec<PathBuf>>, target_blobs: &Vec<(PathBuf, Hash)>) {
-    let all = filter.is_none(); //是否恢复所有文件
     let paths: Vec<PathBuf> = if let Some(filter) = filter {
         filter.clone()
     } else {
         vec![get_working_dir().unwrap()] //None == all(workdir), '.' == cur_dir
     };
-    let dot = paths.contains(&PathBuf::from(".")); //是否包含当前目录
-    let paths = util::integrate_paths(&paths); // file paths
 
     let target_blobs = target_blobs // 转为绝对路径 //TODO tree改变路径表示方式后，这里需要修改
         .iter()
         .map(|(path, hash)| (util::to_workdir_absolute_path(path), hash.clone()))
-        .collect::<Vec<(PathBuf, Hash)>>();
+        .collect::<HashMap<PathBuf, Hash>>();
+
+    let dirs: Vec<PathBuf> = paths.iter().filter(|path| path.is_dir()).cloned().collect();
+    let del_files = target_blobs //统计所有目录中(包括None & '.')，删除的文件
+        .iter()
+        .filter(|(path, _)| {
+            if !path.exists() {
+                for dir in &dirs {
+                    if util::is_parent_dir(path, dir) {
+                        //需要包含在指定dir内
+                        return true;
+                    }
+                }
+            }
+            false
+        })
+        .map(|(path, _)| path.clone())
+        .collect::<HashSet<PathBuf>>(); //HashSet自动去重
+    let mut paths = util::integrate_paths(&paths); //存在的文件路径
+    paths.extend(del_files); //不存在的文件路径
 
-    //TODO @mrbeanc all & dot比较特殊，需要包含被删除的文件，逻辑和add类似 我明天写 @mrbeanc 传递一个目录也需要包含被删除的文件
     let index = Index::new();
     let store = Store::new();
-    for (path, hash) in &target_blobs {
-        if !paths.contains(path) {
-            continue; //不在指定路径内
-        }
-        if path.exists() {
-            let file_hash = util::calc_file_hash(&path); //TODO tree没有存修改时间，所以这里只能用hash判断
-            if file_hash == *hash {
-                continue; //文件未修改 不需要还原
-            }
-        }
-        //文件不存在或已修改
-        store.restore_to_file(hash, &path);
-    }
 
-    //处理工作区的新文件
-    for path in paths {
-        if target_blobs.iter().any(|(target_path, _)| target_path == &path) {
-            //TODO 最好返回HashMap 方便优化
-            continue; //已处理
-        }
-        //未找到，则对于target_commit来说是新文件；若已跟踪，则删除；若未跟踪，则保留
-        if index.tracked(&path) {
-            fs::remove_file(&path).unwrap();
+    for path in &paths {
+        assert!(path.is_absolute() && !path.is_dir()); // 绝对路径且不是目录
+        if !path.exists() {
+            //文件不存在于workdir
+            if target_blobs.contains_key(path) {
+                //文件存在于target_commit
+                store.restore_to_file(&target_blobs[path], &path);
+            } else {
+                //在target_commit和workdir中都不存在(非法路径)
+                println!("fatal: pathspec '{}' did not match any files", path.display());
+            }
+        } else {
+            //文件存在，有两种情况：1.修改 2.新文件
+            if target_blobs.contains_key(path) {
+                //文件已修改(modified)
+                let file_hash = util::calc_file_hash(&path); //TODO tree没有存修改时间，所以这里只能用hash判断
+                if file_hash != target_blobs[path] {
+                    store.restore_to_file(&target_blobs[path], &path);
+                }
+            } else {
+                //新文件，也分两种情况：1.已跟踪，需要删除 2.未跟踪，保留
+                if index.tracked(path) {
+                    //文件已跟踪
+                    fs::remove_file(&path).unwrap();
+                }
+            }
         }
     }
 }