Merge pull request #14 from Nexus-Mods/find-subpath-method

Find subpath method

src/NexusMods.Paths/FileTree/FileTreeNode.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Runtime.CompilerServices;
 using JetBrains.Annotations;
 
 namespace NexusMods.Paths.FileTree;
@@ -18,6 +19,7 @@ public class FileTreeNode<TPath, TValue> : IFileTree<FileTreeNode<TPath, TValue>
 
     private readonly Dictionary<RelativePath, FileTreeNode<TPath, TValue>> _children;
     private FileTreeNode<TPath, TValue>? _parent;
+    private ushort _depth;
 
     /// <summary>
     /// Constructs a new <see cref="FileTreeNode{TPath,TValue}"/> with the given path, name, parent and value.
@@ -34,6 +36,7 @@ public FileTreeNode(TPath path, RelativePath name, bool isFile, TValue? value)
         _isFile = isFile;
         Value = value;
         _children = new Dictionary<RelativePath, FileTreeNode<TPath, TValue>>();
+        _depth = 0;
     }
 
     /// <summary>
@@ -78,6 +81,12 @@ public FileTreeNode<TPath, TValue> Parent
         }
     }
 
+    /// <summary>
+    /// Returns the depth of the node in the tree, the top node in the tree will have
+    /// a depth of 0. So `bar` in `/foo/bar/baz` will have a depth of 2 due to the `/` having a depth of 0.
+    /// </summary>
+    public ushort Depth => _depth;
+
     /// <inheritdoc />
     public FileTreeNode<TPath, TValue> Root
     {
@@ -145,6 +154,7 @@ public void AddChildren(IEnumerable<FileTreeNode<TPath, TValue>> children)
         foreach (var child in children)
         {
             child._parent = this;
+            child._depth = (ushort)(child._parent._depth + 1);
             AddChild(child);
         }
     }
@@ -159,6 +169,7 @@ public void AddChildren(IEnumerable<FileTreeNode<TPath, TValue>> children)
     public void AddChild(FileTreeNode<TPath, TValue> child)
     {
         child._parent = this;
+        child._depth = (ushort)(child._parent._depth + 1);
         _children.Add(child.Name, child);
     }
 
@@ -201,6 +212,57 @@ public void Deconstruct(out TPath path, out TValue? value)
         value = Value;
     }
 
+
+    /// <summary>
+    /// For a given subpath, find all nodes that match the subpath. The returned nodes will be the top of the subpath.
+    /// This allows for searching `/foo/bar/baz/qux` by passing `bar/baz` to this function and this function will then
+    /// return the `bar` node.
+    /// </summary>
+    /// <param name="subpath"></param>
+    /// <returns></returns>
+    public IEnumerable<FileTreeNode<TPath, TValue>> FindSubPath(RelativePath subpath)
+    {
+        var parts = subpath.Parts.ToArray();
+
+        var children = GetAllNodes().Where(n => n.RelativePath.StartsWith(parts[0]));
+
+        foreach (var (_, childNode) in children)
+        {
+            var notFound = false;
+            var node = childNode;
+            foreach (var part in parts[1..])
+            {
+                if (!node.Children.TryGetValue(part, out var found))
+                {
+                    notFound = true;
+                    break;
+                }
+
+                node = found;
+            }
+
+            if (notFound) continue;
+            yield return childNode;
+        }
+    }
+
+    /// <summary>
+    /// Gets all nodes in the tree.
+    /// </summary>
+    /// <returns></returns>
+    public IEnumerable<(RelativePath RelativePath, FileTreeNode<TPath, TValue> Node)> GetAllNodes()
+    {
+        foreach (var (path, node) in Children)
+        {
+            yield return (path, node);
+            foreach (var tuple in node.GetAllNodes())
+            {
+                yield return tuple;
+            }
+        }
+
+    }
+
     /// <summary>
     /// Populates the tree with the given collection of file entries.
     /// </summary>

tests/NexusMods.Paths.Tests/FileTree/RelativeFileTreeTests.cs

@@ -159,6 +159,29 @@ public void Test_GetAllDescendentFilesDictionary(string path, string[] expectedD
             .BeEquivalentTo(expectedDescendentPaths.Select(x => (RelativePath)x));
     }
 
+    [Theory]
+    [InlineData("bar", new [] {"foo/bar"}, new[] {2})]
+    [InlineData("foo", new [] {"foo"}, new[] {1})]
+    [InlineData("bazer", new [] {"baz/bazer"}, new[] {2})]
+
+    public void Test_FindSubPath(string prefix, string[] paths, int[] depths)
+    {
+        var tree = MakeTestTree();
+
+        var nodes = tree.FindSubPath((RelativePath)prefix).ToArray();
+        nodes.Should().HaveCount(paths.Length);
+        nodes.Should().HaveCount(depths.Length);
+
+        nodes.Select(f => (string)f.Path).Should().Contain(paths);
+
+        foreach (var (path, depth) in paths.Zip(depths))
+        {
+            var node = nodes.First(n => n.Path.StartsWith(path));
+            node.Depth.Should().Be((ushort)depth);
+        }
+
+    }
+
     private static FileTreeNode<RelativePath, int> MakeTestTree()
     {
         Dictionary<RelativePath, int> fileEntries;