Merge pull request #137 from JunNishimura/#136

fix restore bug

cmd/restore.go

@@ -19,27 +19,39 @@ import (
 func restoreIndex(rootGoitPath, path string, index *store.Index, tree *object.Tree) error {
 	// get entry
 	_, _, isEntryFound := index.GetEntry([]byte(path))
-	if !isEntryFound {
-		return fmt.Errorf("error: pathspec '%s' did not match any file(s) known to goit", path)
-	}
 
 	// get node
 	node, isNodeFound := object.GetNode(tree.Children, path)
 
-	// restore index
-	if isNodeFound { // if node is in the last commit
-		// change hash
-		isUpdated, err := client.Idx.Update(rootGoitPath, node.Hash, []byte(path))
-		if err != nil {
-			return fmt.Errorf("fail to update index: %w", err)
-		}
-		if !isUpdated {
-			return errors.New("fail to restore index")
+	// if the path is registered in the Index
+	if isEntryFound {
+		// restore index
+		if isNodeFound { // if the file is updated
+			// change hash
+			isUpdated, err := index.Update(rootGoitPath, node.Hash, []byte(path))
+			if err != nil {
+				return fmt.Errorf("fail to update index: %w", err)
+			}
+			if !isUpdated {
+				return errors.New("fail to restore index")
+			}
+		} else { // if the file is newly added
+			// delete entry
+			if err := index.DeleteEntry(rootGoitPath, []byte(path)); err != nil {
+				return fmt.Errorf("fail to delete entry: %w", err)
+			}
 		}
-	} else { // if node is not in the last commit
-		// delete entry
-		if err := index.DeleteEntry(rootGoitPath, []byte(path)); err != nil {
-			return fmt.Errorf("fail to delete entry: %w", err)
+	} else { // if the path is not registered in the index,
+		if isNodeFound { // if the file is deleted
+			isUpdated, err := index.Update(rootGoitPath, node.Hash, []byte(path))
+			if err != nil {
+				return fmt.Errorf("fail to update index: %w", err)
+			}
+			if !isUpdated {
+				return errors.New("fail to restore index")
+			}
+		} else {
+			return fmt.Errorf("error: pathspec '%s' did not match any file(s) known to goit", path)
 		}
 	}
 
@@ -57,17 +69,28 @@ func restoreWorkingDirectory(rootGoitPath, path string, index *store.Index) erro
 		return fmt.Errorf("fail to get object '%s': %w", path, err)
 	}
 
-	// restore file
+	// get abs path
 	absPath, err := filepath.Abs(path)
 	if err != nil {
 		return fmt.Errorf("fail to get abs path '%s': %w", path, err)
 	}
+
+	// make sure the parent path exists
+	parentPath := filepath.Dir(absPath)
+	if _, err := os.Stat(parentPath); os.IsNotExist(err) {
+		if err := os.MkdirAll(parentPath, 0777); err != nil {
+			return fmt.Errorf("fail to make directory %s: %w", parentPath, err)
+		}
+	}
+
+	// restore file
 	f, err := os.Create(absPath)
 	if err != nil {
 		return fmt.Errorf("%w: %s", ErrIOHandling, absPath)
 	}
 	defer f.Close()
 
+	// write contents to the file
 	if _, err := f.WriteString(string(obj.Data)); err != nil {
 		return fmt.Errorf("fail to write to file '%s': %w", absPath, err)
 	}
@@ -92,17 +115,6 @@ var restoreCmd = &cobra.Command{
 			return errors.New("fatal: you must specify path(s) to restore")
 		}
 
-		// check if arg exists
-		for _, arg := range args {
-			argPath, err := filepath.Abs(arg)
-			if err != nil {
-				return fmt.Errorf("fail to get absolute path of %s: %w", arg, err)
-			}
-			if _, err := os.Stat(argPath); os.IsNotExist(err) {
-				return fmt.Errorf("error: pathspec '%s' did not match any file(s) known to goit", arg)
-			}
-		}
-
 		// get staged option
 		isStaged, err := cmd.Flags().GetBool("staged")
 		if err != nil {
@@ -134,8 +146,31 @@ var restoreCmd = &cobra.Command{
 					return fmt.Errorf("fail to get arg abs path: %w", err)
 				}
 				f, err := os.Stat(argAbsPath)
-				if err != nil {
-					return fmt.Errorf("%w: %s", ErrIOHandling, argAbsPath)
+				if os.IsNotExist(err) { // even if the file is not found, the file might be the deleted file
+					// get node
+					cleanedArg := filepath.Clean(arg)
+					cleanedArg = strings.ReplaceAll(cleanedArg, `\`, "/")
+					node, isNodeFound := object.GetNode(tree.Children, cleanedArg)
+					if !isNodeFound {
+						return fmt.Errorf("error: pathspec '%s' did not match any file(s) known to goit", arg)
+					}
+
+					// check if the arg is dir or not
+					if len(node.Children) > 0 { // node is directory
+						paths := node.GetPaths()
+
+						for _, path := range paths {
+							if err := restoreIndex(client.RootGoitPath, path, client.Idx, tree); err != nil {
+								return err
+							}
+						}
+					} else { // node is a file
+						if err := restoreIndex(client.RootGoitPath, cleanedArg, client.Idx, tree); err != nil {
+							return err
+						}
+					}
+
+					continue
 				}
 
 				if f.IsDir() { // directory
@@ -177,8 +212,31 @@ var restoreCmd = &cobra.Command{
 					return fmt.Errorf("fail to get arg abs path: %w", err)
 				}
 				f, err := os.Stat(argAbsPath)
-				if err != nil {
-					return fmt.Errorf("%w: %s", ErrIOHandling, argAbsPath)
+				if os.IsNotExist(err) {
+					// check if the arg is registered in the index
+					cleanedArg := filepath.Clean(arg)
+					cleanedArg = strings.ReplaceAll(cleanedArg, `\`, "/")
+					_, _, isRegistered := client.Idx.GetEntry([]byte(cleanedArg))
+					isRegisteredAsDir := client.Idx.IsRegisteredAsDirectory(cleanedArg)
+
+					if !(isRegistered || isRegisteredAsDir) {
+						return fmt.Errorf("error: pathspec '%s' did not match any file(s) known to goit", arg)
+					}
+
+					if isRegisteredAsDir {
+						entries := client.Idx.GetEntriesByDirectory(cleanedArg)
+						for _, entry := range entries {
+							if err := restoreWorkingDirectory(client.RootGoitPath, string(entry.Path), client.Idx); err != nil {
+								return err
+							}
+						}
+					} else {
+						if err := restoreWorkingDirectory(client.RootGoitPath, cleanedArg, client.Idx); err != nil {
+							return err
+						}
+					}
+
+					continue
 				}
 
 				if f.IsDir() { // directory

internal/object/tree.go

@@ -17,6 +17,34 @@ type Node struct {
 	Children []*Node
 }
 
+func (n *Node) GetPaths() []string {
+	if len(n.Children) == 0 {
+		return []string{n.Name}
+	}
+
+	var paths []string
+	for _, child := range n.Children {
+		ps := child.getPaths(n.Name)
+		paths = append(paths, ps...)
+	}
+
+	return paths
+}
+
+func (n *Node) getPaths(parentDir string) []string {
+	if len(n.Children) == 0 {
+		return []string{fmt.Sprintf("%s/%s", parentDir, n.Name)}
+	}
+
+	var paths []string
+	for _, child := range n.Children {
+		ps := child.getPaths(fmt.Sprintf("%s/%s", parentDir, n.Name))
+		paths = append(paths, ps...)
+	}
+
+	return paths
+}
+
 type Tree struct {
 	object   *Object
 	Children []*Node
@@ -173,7 +201,11 @@ func GetNode(children []*Node, path string) (*Node, bool) {
 			if len(node.Children) == 0 {
 				return node, true
 			}
-			return GetNode(node.Children, pathSplit[1])
+			if len(pathSplit) > 1 {
+				return GetNode(node.Children, pathSplit[1])
+			} else {
+				return node, true
+			}
 		} else if node.Name < searchName {
 			left = middle + 1
 		} else {

internal/object/tree_test.go

@@ -7,6 +7,89 @@ import (
 	"testing"
 )
 
+func TestNewGetPaths(t *testing.T) {
+	type fields struct {
+		node Node
+	}
+	type test struct {
+		name   string
+		fields fields
+		want   []string
+	}
+	tests := []*test{
+		func() *test {
+			hash, _ := hex.DecodeString("87f3c49bccf2597484ece08746d3ee5defaba335")
+
+			return &test{
+				name: "success: empty",
+				fields: fields{
+					node: Node{
+						Hash:     hash,
+						Name:     "test.txt",
+						Children: nil,
+					},
+				},
+				want: []string{"test.txt"},
+			}
+		}(),
+		func() *test {
+			hash, _ := hex.DecodeString("87f3c49bccf2597484ece08746d3ee5defaba335")
+
+			return &test{
+				name: "success: multiple nodes",
+				fields: fields{
+					node: Node{
+						Hash: hash,
+						Name: "dir",
+						Children: []*Node{
+							{
+								Hash: hash,
+								Name: "dir2",
+								Children: []*Node{
+									{
+										Hash:     hash,
+										Name:     "test.txt",
+										Children: nil,
+									},
+									{
+										Hash:     hash,
+										Name:     "test2.txt",
+										Children: nil,
+									},
+								},
+							},
+							{
+								Hash:     hash,
+								Name:     "test.txt",
+								Children: nil,
+							},
+							{
+								Hash:     hash,
+								Name:     "test2.txt",
+								Children: nil,
+							},
+						},
+					},
+				},
+				want: []string{
+					"dir/dir2/test.txt",
+					"dir/dir2/test2.txt",
+					"dir/test.txt",
+					"dir/test2.txt",
+				},
+			}
+		}(),
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := tt.fields.node.GetPaths()
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("got = %v, want = %v", got, tt.want)
+			}
+		})
+	}
+}
+
 func TestNewTree(t *testing.T) {
 	type args struct {
 		object *Object

internal/store/index.go

@@ -90,6 +90,19 @@ func (idx *Index) GetEntry(path []byte) (int, *Entry, bool) {
 	return newEntryFlag, nil, false
 }
 
+func (idx *Index) GetEntriesByDirectory(dirName string) []*Entry {
+	var entries []*Entry
+
+	dirRegexp := regexp.MustCompile(fmt.Sprintf(`%s\/.+`, dirName))
+	for _, entry := range idx.Entries {
+		if dirRegexp.Match(entry.Path) {
+			entries = append(entries, entry)
+		}
+	}
+
+	return entries
+}
+
 func (idx *Index) IsRegisteredAsDirectory(dirName string) bool {
 	if idx.EntryNum == 0 {
 		return false