From 059ec8cde1b29f9a4aadc47c3b25331b04c3f632 Mon Sep 17 00:00:00 2001
From: LurkingNinja <dev@lurking-ninja.com>
Date: Mon, 20 Nov 2023 20:28:49 -0800
Subject: [PATCH] Removing limitations on "using" directives.

---
 CHANGELOG.md          |  4 ++++
 Common.cs             | 10 +++++-----
 DomainReloadSG.csproj |  4 ++--
 README.md             |  9 +++++----
 SourceGenerator.cs    | 11 +++++++++--
 5 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5bbb26e..8d20071 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 All notable changes to this project will be documented in this file.
 
 [READ ME](./README.md) | [MIT LICENSE](./LICENSE)
+## [0.0.3] - 2022-11-20
+### Changed
+- Limitation of missing using directives was removed.
+
 ## [0.0.2] - 2022-11-06
 ### Added
 - This CHANGELOG
diff --git a/Common.cs b/Common.cs
index e0f1816..fb98eea 100644
--- a/Common.cs
+++ b/Common.cs
@@ -18,10 +18,9 @@ internal static class Common
          * {0} name space if exists
          * {1} closing bracket for namespace if needed
          * {2} class definition
+         * {3} using directives
          */
-        private const string NS_TEMPLATE = @"using System;
-using UnityEngine;
-
+        private const string NS_TEMPLATE = @"{3}
 {0}
     {2}
 {1}";
@@ -50,13 +49,14 @@ static void ApplyStaticFieldsAndEventHandlers()
         }}
     }}";
 
-        internal static string NamespaceTemplateResolve(string nameSpace, string source)
+        internal static string NamespaceTemplateResolve(string usingDirectives, string nameSpace, string source)
         {
             var ns = GetNamespaceTemplate(nameSpace);
             return string.Format(NS_TEMPLATE,
                 /*{0}*/ns.Item1,
                 /*{1}*/ns.Item2,
-                /*{2}*/source);
+                /*{2}*/source,
+                /*{3}*/usingDirectives);
         }
 
         internal static bool IsPartial(ClassDeclarationSyntax cds) =>
diff --git a/DomainReloadSG.csproj b/DomainReloadSG.csproj
index bf92cbf..bcbcc53 100644
--- a/DomainReloadSG.csproj
+++ b/DomainReloadSG.csproj
@@ -4,8 +4,8 @@
         <TargetFramework>netstandard2.0</TargetFramework>
         <RootNamespace>LurkingNinja.DomainReloadSG</RootNamespace>
         <RunPostBuildEvent>OnOutputUpdated</RunPostBuildEvent>
-        <AssemblyVersion>0.2</AssemblyVersion>
-        <FileVersion>0.2</FileVersion>
+        <AssemblyVersion>0.3</AssemblyVersion>
+        <FileVersion>0.3</FileVersion>
     </PropertyGroup>
 
     <ItemGroup>
diff --git a/README.md b/README.md
index e8beed8..cb778c0 100644
--- a/README.md
+++ b/README.md
@@ -25,13 +25,13 @@ You have two ways to use this repository:
 ### Limitations
 - The class you're augmenting _has to be_ decorated with the ```partial``` keyword.
 - Currently there is no way of excluding static fields or methods from this service.
-- If you're using extra namespaces currently you need to explicitly write it in the code as opposed of in a ```using``` statement.
 
 ### Exclusion
 - Decorate your partial class with ```[NoDomainReloadSupport]``` attribute if you want to exclude it completely.
 
 ### Example
 ```csharp
+using LurkingNinja.FirstNameSpace.SecondNameSpace.ForTest;
 using UnityEngine;
 
 public partial class TestStaticOnNoDomainReload : MonoBehaviour
@@ -44,6 +44,7 @@ public partial class TestStaticOnNoDomainReload : MonoBehaviour
         Debug.Log($"Started with {_number}");
         _number += 10;
         Application.quitting += OnQuit;
+        OtherTestEvent.OnChangeSomethingStatic += OnQuit;
         Debug.Log($"Ended with {_number}");
     }
 
@@ -59,17 +60,17 @@ Obviously the ```Edit > Project Settings > Editor > Enter Play Mode``` is set an
 The generated source code:
 ```TestStaticOnNoDomainReload_codegen.cs```
 ```csharp
-using System;
+using LurkingNinja.FirstNameSpace.SecondNameSpace.ForTest;
 using UnityEngine;
 
     public partial class TestStaticOnNoDomainReload 
     {
-        [UnityEngine.RuntimeInitializeOnLoadMethod(UnityEngine.RuntimeInitializeLoadType.SubsystemRegistration)]
+        [RuntimeInitializeOnLoadMethod(RuntimeInitializeLoadType.SubsystemRegistration)]
         static void ApplyStaticFieldsAndEventHandlers()
         {
 			_number = default;
 			Application.quitting -= OnQuit;
-
+			OtherTestEvent.OnChangeSomethingStatic -= OnQuit;
         }
     }
 ```
\ No newline at end of file
diff --git a/SourceGenerator.cs b/SourceGenerator.cs
index aee17d2..c6d2691 100644
--- a/SourceGenerator.cs
+++ b/SourceGenerator.cs
@@ -24,8 +24,15 @@ public void Execute(GeneratorExecutionContext context)
 
             var nameSpace = Common.GetNamespace(dsr.ClassToAugment);
 
-            var lines = new StringBuilder();
+            var usingDirectives = new StringBuilder();
+            foreach (var child in dsr.ClassToAugment.SyntaxTree.GetRoot().ChildNodes())
+                if (child.IsKind(SyntaxKind.UsingDirective))
+                    usingDirectives.AppendLine(child.ToString());
+            usingDirectives.AppendLine("using System;");
+            usingDirectives.AppendLine("using UnityEngine;");
 
+            var lines = new StringBuilder();
+            
             foreach (var field in dsr.fields)
             {
                 if (field.Declaration.Variables.Count < 1) continue;
@@ -46,7 +53,7 @@ public void Execute(GeneratorExecutionContext context)
                 /*{1}*/lines);
             
             Common.AddSource(context, dsr.ClassToAugment.Identifier.ToFullString().Trim(), 
-                Common.NamespaceTemplateResolve(nameSpace, source));
+                Common.NamespaceTemplateResolve(usingDirectives.ToString(), nameSpace, source));
         }
     }