From b25b41e927f5bcfc35448e262248d9ecc50201af Mon Sep 17 00:00:00 2001
From: focustense <focustense@protonmail.com>
Date: Tue, 8 Nov 2022 15:56:53 -0500
Subject: [PATCH] Avoid infinite recursion in ReferenceFollower.

Tracks previously-visited keys in a set (a single set is reused per walk, in order to conserve memory) and stops walking the current node if it's a repeated form key.

Fixes #166.
---
 .../Analysis/ReferenceCheckerTests.cs         | 37 +++++++++++++++++++
 .../Analysis/ReferenceFollower.cs             | 20 ++++++----
 2 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/Focus.Providers.Mutagen.Tests/Analysis/ReferenceCheckerTests.cs b/Focus.Providers.Mutagen.Tests/Analysis/ReferenceCheckerTests.cs
index d5f29a9..3e40acb 100644
--- a/Focus.Providers.Mutagen.Tests/Analysis/ReferenceCheckerTests.cs
+++ b/Focus.Providers.Mutagen.Tests/Analysis/ReferenceCheckerTests.cs
@@ -292,6 +292,43 @@ public void WhenPluginsExcluded_SkipsExcludedRecords()
                 x => AssertPath(x, PathFrom(npc).Id("Face").Key("123456:badheadpart.esp", RecordType.HeadPart)));
         }
 
+        [Fact]
+        public void IgnoresCircularReferences()
+        {
+            var headPartKeys = groups
+                .AddRecords<HeadPart>(
+                    "main.esp",
+                    x =>
+                    {
+                        x.EditorID = "Hair";
+                        x.Color.SetTo(AddRecord<ColorRecord>("HairColor"));
+                        x.ExtraParts.AddRange(groups.AddRecords<HeadPart>("main.esp", hp =>
+                        {
+                            hp.EditorID = "Hairline";
+                            hp.Color.SetTo(AddRecord<ColorRecord>("HairlineColor"));
+                        }).ToFormKeys());
+                        x.ExtraParts.Add(x.FormKey.AsLinkGetter<IHeadPartGetter>());
+                        x.ExtraParts.Add(FormKey.Factory("123456:badheadpart.esp").AsLink<IHeadPartGetter>());
+                        x.TextureSet.SetTo(AddRecord<TextureSet>("HairTextureSet"));
+                    },
+                    x =>
+                    {
+                        x.EditorID = "Face";
+                        x.Color.SetTo(AddRecord<ColorRecord>("FaceColor"));
+                    })
+                .ToFormKeys();
+            var npc = new Npc(FormKey.Factory("123456:main.esp"), SkyrimRelease.SkyrimSE)
+            {
+                EditorID = "TestNpc",
+                HeadParts = new(headPartKeys.Select(x => x.AsLink<IHeadPartGetter>())),
+            };
+            var invalidPaths = checker.GetInvalidPaths(npc);
+
+            Assert.Collection(
+                invalidPaths,
+                x => AssertPath(x, PathFrom(npc).Id("Hair").Key("123456:badheadpart.esp", RecordType.HeadPart)));
+        }
+
         private FormKey AddRecord<T>(string editorId)
             where T : class, ISkyrimMajorRecord
         {
diff --git a/Focus.Providers.Mutagen/Analysis/ReferenceFollower.cs b/Focus.Providers.Mutagen/Analysis/ReferenceFollower.cs
index 8441d0e..ee1deac 100644
--- a/Focus.Providers.Mutagen/Analysis/ReferenceFollower.cs
+++ b/Focus.Providers.Mutagen/Analysis/ReferenceFollower.cs
@@ -52,7 +52,7 @@ public abstract class ReferenceFollower<T, TAccumulate, TResult> : IReferenceFol
         protected IGroupCache GroupCache { get; private init; }
 
         private readonly HashSet<string> excludedPluginNames = new(StringComparer.CurrentCultureIgnoreCase);
-        private readonly List<Func<T, TResult?, IEnumerable<TResult>>> routes = new();
+        private readonly List<Func<T, TResult?, ISet<FormKey>, IEnumerable<TResult>>> routes = new();
 
         public ReferenceFollower(IGroupCache groupCache)
         {
@@ -114,19 +114,19 @@ public IReferenceFollower<T> Follow<TNext>(
             subrecordFollower.AccumulatorCache = AccumulatorCache;
             subrecordFollower.WithPluginExclusions(excludedPluginNames);
             configure?.Invoke(subrecordFollower);
-            routes.Add((record, previous) => Walk(record, previous, linksSelector, subrecordFollower));
+            routes.Add((record, previous, visited) => Walk(record, previous, linksSelector, subrecordFollower, visited));
             return this;
         }
 
         public IReferenceFollower<T> FollowSelf(Func<T, IFormLinkGetter<T>?> linkSelector)
         {
-            routes.Add((record, previous) => Walk(record, previous, r => new[] { linkSelector(r) }, this));
+            routes.Add((record, previous, visited) => Walk(record, previous, r => new[] { linkSelector(r) }, this, visited));
             return this;
         }
 
         public IReferenceFollower<T> FollowSelf(Func<T, IEnumerable<IFormLinkGetter<T>?>?> linksSelector)
         {
-            routes.Add((record, previous) => Walk(record, previous, linksSelector, this));
+            routes.Add((record, previous, visited) => Walk(record, previous, linksSelector, this, visited));
             return this;
         }
 
@@ -139,16 +139,21 @@ public IReferenceFollower<T> WithPluginExclusions(IEnumerable<string> pluginName
 
         protected IEnumerable<IEnumerable<TResult>> WalkAll(T record)
         {
-            return routes.Select(walk => walk(record, default));
+            var visited = new HashSet<FormKey>();
+            return routes.Select(walk => walk(record, default, visited));
         }
 
         private IEnumerable<TResult> Walk<U>(
             T record,
             TResult? previous,
             Func<T, IEnumerable<IFormLinkGetter<U>?>?> linksSelector,
-            ReferenceFollower<U, TAccumulate, TResult> subrecordFollower)
+            ReferenceFollower<U, TAccumulate, TResult> subrecordFollower,
+            ISet<FormKey> visited)
             where U : class, ISkyrimMajorRecordGetter
         {
+            if (visited.Contains(record.FormKey))
+                yield break;
+            visited.Add(record.FormKey);
             var originKey = record.FormKey.ToRecordKey();
             var originType = typeof(T).GetRecordType();
             var current = AccumulatorCache.GetOrAdd(record.FormKey, _ => Visit(record));
@@ -174,11 +179,12 @@ private IEnumerable<TResult> Walk<U>(
                     continue;
                 var nextRecord = nextRecordWithSource.Value;
                 var subrecordResults = subrecordFollower.routes
-                    .Select(ps => ps(nextRecord, result))
+                    .Select(ps => ps(nextRecord, result, visited))
                     .SelectMany(results => results);
                 foreach (var subrecordPath in subrecordResults)
                     yield return subrecordPath;
             }
+            visited.Remove(record.FormKey);
         }
 
         protected abstract ReferenceFollower<TNext, TAccumulate, TResult> CreateChild<TNext>()