From a8677f6b808fb010eaa4f9f6c2dedc929fbbaf59 Mon Sep 17 00:00:00 2001
From: JPaja <jovan.pavlovic99@gmail.com>
Date: Sat, 16 Oct 2021 20:58:41 +0200
Subject: [PATCH] Fix: Implementation base type gathering

Fixes #198
---
 .../ImplementationAnalyzerUtilities.cs        | 30 ++++++++++++++++---
 .../DotNetRelations.cs                        |  4 +--
 .../AbstractInheritanceTest.cs                |  9 +++---
 3 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/src/AsmResolver.Workspaces.DotNet/Analyzers/Implementation/ImplementationAnalyzerUtilities.cs b/src/AsmResolver.Workspaces.DotNet/Analyzers/Implementation/ImplementationAnalyzerUtilities.cs
index 145248b7d..cfb1ee54e 100644
--- a/src/AsmResolver.Workspaces.DotNet/Analyzers/Implementation/ImplementationAnalyzerUtilities.cs
+++ b/src/AsmResolver.Workspaces.DotNet/Analyzers/Implementation/ImplementationAnalyzerUtilities.cs
@@ -10,16 +10,38 @@ internal static class ImplementationAnalyzerUtilities
     {
         private static readonly SignatureComparer _comparer = new ();
 
+        private static IEnumerable<ITypeDefOrRef> GetBaseTypes(WorkspaceIndexNode baseNode)
+        {
+            var visited = new HashSet<WorkspaceIndexNode>();
+            var agenda = new Queue<WorkspaceIndexNode>();
+            agenda.Enqueue(baseNode);
+            while (agenda.Count != 0)
+            {
+                var node = agenda.Dequeue();
+                if(!visited.Add(node))
+                    continue;
+                var baseTypeNodes = node.ForwardRelations.GetNodes(DotNetRelations.BaseType);
+                foreach (var baseTypeNode in baseTypeNodes)
+                {
+                    var baseType = (ITypeDefOrRef)baseTypeNode.Subject;
+                    agenda.Enqueue(baseTypeNode);
+                    var baseTypeDefinitions = baseTypeNode.BackwardRelations.GetNodes(DotNetRelations.ReferenceType);
+                    foreach (var baseTypeDefinition in baseTypeDefinitions)
+                        agenda.Enqueue(baseTypeDefinition);
+                    yield return baseType;
+                }
+            }
+        }
+
         internal static IEnumerable<MethodDefinition> FindBaseMethods(this MethodDefinition subject,
             WorkspaceIndex index)
         {
             if (subject.DeclaringType is not { } declaringType)
                 yield break;
 
-            var baseTypes = index
-                .GetOrCreateNode(declaringType) // Get indexed declaring type.
-                .ForwardRelations.GetObjects(DotNetRelations.BaseType) // Get types that this declaring type is implementing.
-                .ToArray();
+            var declaringTypeNode = index.GetOrCreateNode(declaringType);
+
+            var baseTypes = GetBaseTypes(declaringTypeNode).ToArray();
 
             foreach (var baseType in baseTypes)
             {
diff --git a/src/AsmResolver.Workspaces.DotNet/DotNetRelations.cs b/src/AsmResolver.Workspaces.DotNet/DotNetRelations.cs
index f77f52605..b81e578c2 100644
--- a/src/AsmResolver.Workspaces.DotNet/DotNetRelations.cs
+++ b/src/AsmResolver.Workspaces.DotNet/DotNetRelations.cs
@@ -48,9 +48,9 @@ public static class DotNetRelations
             new Guid("ce11d2f6-a423-429d-ad37-2f073fdf63be"));
 
         /// <summary>
-        /// Describes the relationship between a type definition and its reference.
+        /// Describes the relationship between a type definition and its reference or specification.
         /// </summary>
-        public static readonly ObjectRelation<TypeDefinition, TypeReference> ReferenceType = new(
+        public static readonly ObjectRelation<TypeDefinition, ITypeDefOrRef> ReferenceType = new(
             "ReferenceType",
             new Guid("3cc86779-338c-4165-a00c-da547a2e8549"));
 
diff --git a/test/AsmResolver.Workspaces.DotNet.Tests/AbstractInheritanceTest.cs b/test/AsmResolver.Workspaces.DotNet.Tests/AbstractInheritanceTest.cs
index 897ede423..3f8bb85ae 100644
--- a/test/AsmResolver.Workspaces.DotNet.Tests/AbstractInheritanceTest.cs
+++ b/test/AsmResolver.Workspaces.DotNet.Tests/AbstractInheritanceTest.cs
@@ -23,13 +23,13 @@ public void InterfaceImplementation()
             var abstractMethod = abstractType.Methods.First(m=>m.Name == nameof(MyAboveAbstractClass.TestAboveAbstract));
 
             var implType1 = (TypeDefinition)module.LookupMember(typeof(MyDerivedAboveClass).MetadataToken);
-            var implMethod1 = abstractType.Methods.First(m => m.Name == nameof(MyAboveAbstractClass.TestAboveAbstract));
+            var implMethod1 = implType1.Methods.First(m => m.Name == nameof(MyAboveAbstractClass.TestAboveAbstract));
 
             var implType2 = (TypeDefinition)module.LookupMember(typeof(MyDerivedInbetweenClass).MetadataToken);
-            var implMethod2 = abstractType.Methods.First(m => m.Name == nameof(MyAboveAbstractClass.TestAboveAbstract));
+            var implMethod2 = implType2.Methods.First(m => m.Name == nameof(MyAboveAbstractClass.TestAboveAbstract));
 
             var implType3 = (TypeDefinition)module.LookupMember(typeof(MyDerivedClassGeneric).MetadataToken);
-            var implMethod3 = abstractType.Methods.First(m => m.Name == nameof(MyAboveAbstractClass.TestAboveAbstract));
+            var implMethod3 = implType3.Methods.First(m => m.Name == nameof(MyAboveAbstractClass.TestAboveAbstract));
 
 
             var workspace = new DotNetWorkspace();
@@ -37,8 +37,9 @@ public void InterfaceImplementation()
             workspace.Analyze();
 
             var node = workspace.Index.GetOrCreateNode(abstractMethod);
-            var implMethods = node.BackwardRelations.GetObjects(DotNetRelations.ImplementationMethod);
+            var implMethods = node.BackwardRelations.GetObjects(DotNetRelations.ImplementationMethod).ToArray();
 
+            Assert.Equal(3, implMethods.Length);
             Assert.Contains(implMethod1, implMethods);
             Assert.Contains(implMethod2, implMethods);
             Assert.Contains(implMethod3, implMethods);