Tracker analyzer, yield return orig analyzer

.gitignore

@@ -2,4 +2,5 @@ bin/
 obj/
 /packages/
 riderModule.iml
-/_ReSharper.Caches/
+/_ReSharper.Caches/
+.vs

CelesteAnalyzer/CelesteAnalyzer.Sample/Examples.cs

@@ -2,8 +2,10 @@
 // ReSharper disable UnusedMember.Global
 
 using System;
+using System.Collections;
 using System.Numerics;
 using System.Runtime.CompilerServices;
+using System.Text;
 using Celeste.Mod.Helpers;
 using Mono.Cecil.Cil;
 using MonoMod.Cil;
@@ -50,6 +52,8 @@ public void Bad()
             {
                 return orig(arg, arg2);
             };
+            On.Celeste.Player.OnSomeRoutine += RoutineOnHook;
+
             IL.Celeste.Player.NormalUpdate += static ctx =>
             {
                 var cursor = new ILCursor(ctx);
@@ -64,6 +68,18 @@ public void Bad()
             HookTarget(4);
         }
 
+        private static IEnumerator RoutineOnHook(Func<IEnumerator> orig)
+        {
+            yield return new SwapImmediately(orig());
+
+            var origRoutine = orig();
+            while (origRoutine.MoveNext())
+                yield return origRoutine.Current;
+
+
+            yield return orig();
+        }
+
         private int Cb(int arg)
         {
             return arg * 2;
@@ -92,6 +108,7 @@ private void NotAHook(Action<int> orig, int arg)
     }
 
     [CustomEntity, Other]
+    [Tracked]
     class MyEntity : Entity
     {
         public MyEntity(Vector2 blah)
@@ -101,6 +118,15 @@ public MyEntity(Vector2 blah)
             cb = () =>
             {
                 Console.WriteLine(Scene);
+                // we don't have access to this class and it's not tracked, no warning here
+                new EntityList().FindAll<StringBuilder>();
+                new EntityList().FindAll<MyEntity>();
+
+
+
+                new Tracker().GetEntities<MyEntity>();
+                new Tracker().GetEntities<OtherEntity>();
+                new Tracker().GetEntities<MyTrigger>();
             };
         }
 
@@ -112,7 +138,14 @@ public object DoStuff()
         private Action cb;
     }
 
+    [TrackedAs(typeof(MyEntity))]
+    class OtherEntity : Entity
+    {
+
+    }
+
     [CustomEntity($"MyTrigger = {nameof(Generator)}")]
+    //[TrackedAs(typeof(Trigger))]
     class MyTrigger : Trigger
     {
         void Generator(int a)
@@ -133,6 +166,8 @@ namespace On.Celeste
     public class Player
     {
         public static event Func<Func<int, string, int>, int, string, int> OnNormalUpdate;
+
+        public static event Func<Func<IEnumerator>, IEnumerator> OnSomeRoutine;
     }
 }
 
@@ -146,13 +181,27 @@ public class Player
 
 namespace Celeste.Mod.Helpers
 {
+    public record SwapImmediately(IEnumerator Enumerable);
+
     public class CustomEntity : Attribute
     {
         public CustomEntity(params string[] ids)
         {
 
         }
     }
+
+    public class Tracked : Attribute
+    {
+    }
+
+    public class TrackedAs : Attribute
+    {
+        public TrackedAs(Type trackedAs)
+        {
+
+        }
+    }
 
     public class Other : Attribute
     {
@@ -168,6 +217,22 @@ public class Entity
     public class Trigger : Entity
     {
     }
+
+    public class EntityList
+    {
+        public System.Collections.Generic.List<T> FindAll<T>()
+        {
+            return new();
+        } 
+    }
+
+    public class Tracker
+    {
+        public System.Collections.Generic.List<Entity> GetEntities<T>() where T : Entity
+        {
+            return new();
+        }
+    }
 }
 
 

CelesteAnalyzer/CelesteAnalyzer/CelesteAnalyzer.csproj

@@ -1,8 +1,7 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
 
     <PropertyGroup>
         <TargetFramework>netstandard2.0</TargetFramework>
-        <PackageTargetFallback>net452</PackageTargetFallback>
         <Nullable>enable</Nullable>
         <LangVersion>latest</LangVersion>
 
@@ -39,8 +38,8 @@
             <PrivateAssets>all</PrivateAssets>
             <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
         </PackageReference>
-        <PackageReference Include="Microsoft.CodeAnalysis.CSharp" Version="4.6.0" PrivateAssets="all" />
-        <PackageReference Include="Microsoft.CodeAnalysis.CSharp.Workspaces" Version="4.6.0" PrivateAssets="all" />
+        <PackageReference Include="Microsoft.CodeAnalysis.CSharp" Version="4.4.0" PrivateAssets="all" />
+        <PackageReference Include="Microsoft.CodeAnalysis.CSharp.Workspaces" Version="4.4.0" PrivateAssets="all" />
     </ItemGroup>
 
     <ItemGroup>

CelesteAnalyzer/CelesteAnalyzer/DiagnosticIds.cs

@@ -15,4 +15,8 @@ public static class DiagnosticIds
     internal const string CustomEntityGeneratorInvalid = "CL0011";
     internal const string CustomEntityNoIDs = "CL0012";
     internal const string UsingSceneInWrongPlace = "CL0013";
+    internal const string DontUseFindAll = "CL0014";
+    internal const string TrackerUsedOnUntrackedType = "CL0015";
+    internal const string InvalidTrackedAs = "CL0016";
+    internal const string DontYieldReturnOrig = "CL0017";
 }

CelesteAnalyzer/CelesteAnalyzer/EntityAnalyzer.cs

@@ -12,8 +12,6 @@ namespace CelesteAnalyzer;
 [DiagnosticAnalyzer(LanguageNames.CSharp)]
 public class EntityAnalyzer : DiagnosticAnalyzer
 {
-    private const string Category = "Usage";
-
     private static readonly DiagnosticDescriptor CustomEntityWithNoValidCtorRule 
         = Utils.CreateDiagnostic(DiagnosticIds.CustomEntityWithNoValidCtor);
 
@@ -98,9 +96,7 @@ private void AnalyzeClassDecl(SymbolAnalysisContext ctx)
         var ids = customEntityAttr.ConstructorArguments.First().Values;
         if (ids.Length == 0)
         {
-            var customEntityAttrSyntax = syntax.Value.AttributeLists
-                .SelectMany(a => a.Attributes)
-                .FirstOrDefault(a => a.Name.ToString() is "CustomEntity");
+            var customEntityAttrSyntax = Utils.GetAttributeSyntaxFromClassDef(customEntityAttr, syntax.Value);
             ctx.ReportDiagnostic(Diagnostic.Create(CustomEntityNoIDsRule, customEntityAttrSyntax?.GetLocation()));
         }
 

CelesteAnalyzer/CelesteAnalyzer/HookAnalyzer.cs

@@ -1,7 +1,9 @@
 using System;
+using System.Collections;
 using System.Collections.Immutable;
 using System.Linq;
 using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Diagnostics;
 using Microsoft.CodeAnalysis.Operations;
@@ -14,16 +16,17 @@ namespace CelesteAnalyzer;
 [DiagnosticAnalyzer(LanguageNames.CSharp)]
 public class HookAnalyzer : DiagnosticAnalyzer
 {
-    private const string Category = "Usage";
-
     private static readonly DiagnosticDescriptor CallOrigInHooksRule
         = Utils.CreateDiagnostic(DiagnosticIds.CallOrigInHooks);
 
     private static readonly DiagnosticDescriptor HooksShouldBeStaticRule
         = Utils.CreateDiagnostic(DiagnosticIds.HooksShouldBeStatic);
 
+    private static readonly DiagnosticDescriptor DontYieldReturnOrigRule
+        = Utils.CreateDiagnostic(DiagnosticIds.DontYieldReturnOrig);
+
     public override ImmutableArray<DiagnosticDescriptor> SupportedDiagnostics { get; } =
-        ImmutableArray.Create(CallOrigInHooksRule, HooksShouldBeStaticRule);
+        ImmutableArray.Create(CallOrigInHooksRule, HooksShouldBeStaticRule, DontYieldReturnOrigRule);
 
     public override void Initialize(AnalysisContext context)
     {
@@ -119,17 +122,44 @@ private static void AnalyzeHook(OperationAnalysisContext context, IMethodSymbol
         // hooks should call orig in at least one code path
         if (bodySyntax is not null)
         {
+            bool origCalled = false;
+            bool isEnumeratorMethod = methodSymbol.ReturnType.Name == nameof(IEnumerator);
+
             foreach (var st in bodySyntax.DescendantNodes())
             {
-                if (st is InvocationExpressionSyntax invocationExpressionSyntax &&
-                    invocationExpressionSyntax.Expression.ToString() == firstParam.Name)
+                if (IsOrig(st, firstParam))
+                {
+                    origCalled = true;
+
+                    // no point in checking the method any more if this isn't an enumerator
+                    if (!isEnumeratorMethod)
+                        return;
+                }
+
+                if (isEnumeratorMethod && st is YieldStatementSyntax yield)
                 {
-                    return;
+                    if (yield.ReturnOrBreakKeyword.IsKind(SyntaxKind.ReturnKeyword))
+                    {
+                        // avoid yield return orig();
+                        if (IsOrig(yield.Expression, firstParam))
+                        {
+                            context.ReportDiagnostic(Diagnostic.Create(DontYieldReturnOrigRule, yield.GetLocation(), firstParam.Type.Name));
+                        }
+                    }
+                }
+
+                static bool IsOrig(SyntaxNode? st, IParameterSymbol orig)
+                {
+                    return st is InvocationExpressionSyntax invocationExpressionSyntax &&
+                           invocationExpressionSyntax.Expression.ToString() == orig.Name;
                 }
             }
 
-            var diagnostic = Diagnostic.Create(CallOrigInHooksRule, loc, firstParam.Type.Name);
-            context.ReportDiagnostic(diagnostic);
+            if (!origCalled)
+            {
+                var diagnostic = Diagnostic.Create(CallOrigInHooksRule, loc, firstParam.Type.Name);
+                context.ReportDiagnostic(diagnostic);
+            }
         }
     }
 }

CelesteAnalyzer/CelesteAnalyzer/IlCursorAnalyzer.cs

@@ -13,8 +13,6 @@ namespace CelesteAnalyzer;
 [DiagnosticAnalyzer(LanguageNames.CSharp)]
 public class IlCursorAnalyzer : DiagnosticAnalyzer
 {
-    private const string Category = "Usage";
-
     private static readonly DiagnosticDescriptor DontUseLambdasRule
         = Utils.CreateDiagnostic(DiagnosticIds.DontUseLambdas);