From ef4be0b959b7f64080a04e5e123c47a3b2e8adfb Mon Sep 17 00:00:00 2001
From: Dimitar Dobrev <dpldobrev@yahoo.com>
Date: Tue, 8 Jul 2014 22:28:08 +0300
Subject: [PATCH] Fixed the generation of code in a single file.

Signed-off-by: Dimitar Dobrev <dpldobrev@yahoo.com>
---
 src/Generator/Generator.cs                    | 74 ++++++++++---------
 src/Generator/Generators/CLI/CLIGenerator.cs  |  6 +-
 .../Generators/CLI/CLIHeadersTemplate.cs      |  6 +-
 .../Generators/CLI/CLISourcesTemplate.cs      |  6 +-
 .../Generators/CLI/CLITextTemplate.cs         |  6 +-
 .../Generators/CSharp/CSharpGenerator.cs      |  4 +-
 .../Generators/CSharp/CSharpTextTemplate.cs   | 68 ++++++++---------
 src/Generator/Generators/Template.cs          | 19 ++---
 tests/CSharpTemp/AnotherUnit.cpp              |  5 ++
 tests/CSharpTemp/AnotherUnit.h                |  3 +
 tests/CSharpTemp/CSharpTemp.Tests.cs          |  6 ++
 tests/CSharpTemp/CSharpTemp.cs                |  1 +
 12 files changed, 111 insertions(+), 93 deletions(-)
 create mode 100644 tests/CSharpTemp/AnotherUnit.cpp
 create mode 100644 tests/CSharpTemp/AnotherUnit.h

diff --git a/src/Generator/Generator.cs b/src/Generator/Generator.cs
index 8f3a6093..69d0dac4 100644
--- a/src/Generator/Generator.cs
+++ b/src/Generator/Generator.cs
@@ -67,49 +67,57 @@ namespace CppSharp.Generators
         {
             var outputs = new List<GeneratorOutput>();
 
-            foreach (var unit in Driver.ASTContext.TranslationUnits)
-            {
-                if (!unit.IsGenerated || !unit.HasDeclarations)
-                    continue;
+            var units = Driver.ASTContext.TranslationUnits.Where(
+                u => u.IsGenerated && u.HasDeclarations && !u.IsSystemHeader && u.IsValid);
+            if (Driver.Options.IsCSharpGenerator && Driver.Options.GenerateSingleCSharpFile)
+                GenerateSingleTemplate(units, outputs);
+            else
+                foreach (var unit in units)
+                    GenerateTemplate(unit, outputs);
+            return outputs;
+        }
 
-                if (unit.IsSystemHeader || !unit.IsValid)
-                    continue;
+        private void GenerateSingleTemplate(IEnumerable<TranslationUnit> units, ICollection<GeneratorOutput> outputs)
+        {
+            var output = new GeneratorOutput
+            {
+                TranslationUnit = new TranslationUnit
+                {
+                    FilePath = string.Format("{0}.cs", Driver.Options.OutputNamespace ?? Driver.Options.LibraryName)
+                },
+                Templates = Generate(units)
+            };
+            output.Templates[0].Process();
+            outputs.Add(output);
+
+            OnUnitGenerated(output);
+        }
 
-                var templates = Generate(unit);
-                if (templates.Count == 0)
-                    continue;
+        private void GenerateTemplate(TranslationUnit unit, ICollection<GeneratorOutput> outputs)
+        {
+            var templates = Generate(new[] { unit });
+            if (templates.Count == 0)
+                return;
 
-                if (templates.Count == 1)
-                {
-                    templates[0].Process(Template.Order.First | Template.Order.Last);
-                }
-                else
-                {
-                    templates.First().Process(Template.Order.First);
-                    for (var i = 1; i < templates.Count - 1; i++)
-                    {
-                        templates[i].Process(Template.Order.InBetween);
-                    }
-                    templates.Last().Process(Template.Order.Last);
-                }
-
-                var output = new GeneratorOutput
-                    {
-                        TranslationUnit = unit,
-                        Templates = templates
-                    };
-                outputs.Add(output);
-
-                OnUnitGenerated(output);
+            foreach (var template in templates)
+            {
+                template.Process();
             }
 
-            return outputs;
+            var output = new GeneratorOutput
+            {
+                TranslationUnit = unit,
+                Templates = templates
+            };
+            outputs.Add(output);
+
+            OnUnitGenerated(output);
         }
 
         /// <summary>
         /// Generates the outputs for a given translation unit.
         /// </summary>
-        public abstract List<Template> Generate(TranslationUnit unit);
+        public abstract List<Template> Generate(IEnumerable<TranslationUnit> units);
 
         public static string GeneratedIdentifier(string id)
         {
diff --git a/src/Generator/Generators/CLI/CLIGenerator.cs b/src/Generator/Generators/CLI/CLIGenerator.cs
index 5241b3ec..4afbf37f 100644
--- a/src/Generator/Generators/CLI/CLIGenerator.cs
+++ b/src/Generator/Generators/CLI/CLIGenerator.cs
@@ -17,14 +17,14 @@ namespace CppSharp.Generators.CLI
             Type.TypePrinterDelegate += type => type.Visit(typePrinter);
         }
 
-        public override List<Template> Generate(TranslationUnit unit)
+        public override List<Template> Generate(IEnumerable<TranslationUnit> units)
         {
             var outputs = new List<Template>();
 
-            var header = new CLIHeadersTemplate(Driver, unit);
+            var header = new CLIHeadersTemplate(Driver, units);
             outputs.Add(header);
 
-            var source = new CLISourcesTemplate(Driver, unit);
+            var source = new CLISourcesTemplate(Driver, units);
             outputs.Add(source);
 
             return outputs;
diff --git a/src/Generator/Generators/CLI/CLIHeadersTemplate.cs b/src/Generator/Generators/CLI/CLIHeadersTemplate.cs
index 7f583d7c..03f36ce6 100644
--- a/src/Generator/Generators/CLI/CLIHeadersTemplate.cs
+++ b/src/Generator/Generators/CLI/CLIHeadersTemplate.cs
@@ -15,12 +15,12 @@ namespace CppSharp.Generators.CLI
     {
         public override string FileExtension { get { return "h"; } }
 
-        public CLIHeadersTemplate(Driver driver, TranslationUnit unit)
-            : base(driver, unit)
+        public CLIHeadersTemplate(Driver driver, IEnumerable<TranslationUnit> units)
+            : base(driver, units)
         {
         }
 
-        public override void Process(Order order)
+        public override void Process()
         {
             PushBlock(BlockKind.Header);
             PopBlock();
diff --git a/src/Generator/Generators/CLI/CLISourcesTemplate.cs b/src/Generator/Generators/CLI/CLISourcesTemplate.cs
index 8d0e8684..b9568359 100644
--- a/src/Generator/Generators/CLI/CLISourcesTemplate.cs
+++ b/src/Generator/Generators/CLI/CLISourcesTemplate.cs
@@ -16,13 +16,13 @@ namespace CppSharp.Generators.CLI
     /// </summary>
     public class CLISourcesTemplate : CLITextTemplate
     {
-        public CLISourcesTemplate(Driver driver, TranslationUnit unit)
-            : base(driver, unit)
+        public CLISourcesTemplate(Driver driver, IEnumerable<TranslationUnit> units)
+            : base(driver, units)
         {
             
         }
 
-        public override void Process(Order order)
+        public override void Process()
         {
             PushBlock(BlockKind.Header);
             PopBlock();
diff --git a/src/Generator/Generators/CLI/CLITextTemplate.cs b/src/Generator/Generators/CLI/CLITextTemplate.cs
index 63aa0798..02d2735d 100644
--- a/src/Generator/Generators/CLI/CLITextTemplate.cs
+++ b/src/Generator/Generators/CLI/CLITextTemplate.cs
@@ -61,8 +61,8 @@ namespace CppSharp.Generators.CLI
 
         public ISet<Include> Includes;
 
-        protected CLITextTemplate(Driver driver, TranslationUnit unit)
-            : base(driver, unit)
+        protected CLITextTemplate(Driver driver, IEnumerable<TranslationUnit> units)
+            : base(driver, units)
         {
             TypePrinter = new CLITypePrinter(driver);
             Includes = new HashSet<Include>();
@@ -70,7 +70,7 @@ namespace CppSharp.Generators.CLI
 
         public abstract override string FileExtension { get; }
 
-        public abstract override void Process(Order order);
+        public abstract override void Process();
 
         #region Helpers
 
diff --git a/src/Generator/Generators/CSharp/CSharpGenerator.cs b/src/Generator/Generators/CSharp/CSharpGenerator.cs
index f65b8151..7b5d9350 100644
--- a/src/Generator/Generators/CSharp/CSharpGenerator.cs
+++ b/src/Generator/Generators/CSharp/CSharpGenerator.cs
@@ -16,11 +16,11 @@ namespace CppSharp.Generators.CSharp
             CppSharp.AST.Type.TypePrinterDelegate += type => type.Visit(typePrinter).Type;
         }
 
-        public override List<Template> Generate(TranslationUnit unit)
+        public override List<Template> Generate(IEnumerable<TranslationUnit> units)
         {
             var outputs = new List<Template>();
 
-            var template = new CSharpTextTemplate(Driver, unit, typePrinter, expressionPrinter);
+            var template = new CSharpTextTemplate(Driver, units, typePrinter, expressionPrinter);
             outputs.Add(template);
 
             return outputs;
diff --git a/src/Generator/Generators/CSharp/CSharpTextTemplate.cs b/src/Generator/Generators/CSharp/CSharpTextTemplate.cs
index 40b6727d..ff71ab7f 100644
--- a/src/Generator/Generators/CSharp/CSharpTextTemplate.cs
+++ b/src/Generator/Generators/CSharp/CSharpTextTemplate.cs
@@ -91,8 +91,8 @@ namespace CppSharp.Generators.CSharp
             get { return "cs"; }
         }
 
-        public CSharpTextTemplate(Driver driver, TranslationUnit unit, CSharpTypePrinter typePrinter, CSharpExpressionPrinter expressionPrinter)
-            : base(driver, unit)
+        public CSharpTextTemplate(Driver driver, IEnumerable<TranslationUnit> units, CSharpTypePrinter typePrinter, CSharpExpressionPrinter expressionPrinter)
+            : base(driver, units)
         {
             TypePrinter = typePrinter;
             ExpressionPrinter = expressionPrinter;
@@ -126,39 +126,36 @@ namespace CppSharp.Generators.CSharp
 
         #endregion
 
-        public override void Process(Order order)
+        public override void Process()
         {
             GenerateHeader();
 
-            if (!Options.GenerateSingleCSharpFile || order.HasFlag(Order.First))
+            PushBlock(CSharpBlockKind.Usings);
+            WriteLine("using System;");
+            WriteLine("using System.Runtime.InteropServices;");
+            WriteLine("using System.Security;");
+            foreach (var customUsingStatement in Options.DependentNameSpaces)
             {
-                PushBlock(CSharpBlockKind.Usings);
-                WriteLine("using System;");
-                WriteLine("using System.Runtime.InteropServices;");
-                WriteLine("using System.Security;");
-                foreach (var customUsingStatement in Options.DependentNameSpaces)
-                {
-                    WriteLine(string.Format("using {0};", customUsingStatement));
-                }
-                PopBlock(NewLineKind.BeforeNextBlock);
+                WriteLine(string.Format("using {0};", customUsingStatement));
+            }
+            PopBlock(NewLineKind.BeforeNextBlock);
 
-                if (Options.GenerateLibraryNamespace)
-                {
-                    PushBlock(CSharpBlockKind.Namespace);
-                    WriteLine("namespace {0}", Driver.Options.OutputNamespace);
-                    WriteStartBraceIndent();
-                }
+            if (Options.GenerateLibraryNamespace)
+            {
+                PushBlock(CSharpBlockKind.Namespace);
+                WriteLine("namespace {0}", Driver.Options.OutputNamespace);
+                WriteStartBraceIndent();
             }
 
-            GenerateDeclContext(TranslationUnit);
+            foreach (var unit in TranslationUnits)
+            {
+                GenerateDeclContext(unit);
+            }
 
-            if (!Options.GenerateSingleCSharpFile || order.HasFlag(Order.Last))
+            if (Options.GenerateLibraryNamespace)
             {
-                if (Options.GenerateLibraryNamespace)
-                {
-                    WriteCloseBraceIndent();
-                    PopBlock(NewLineKind.BeforeNextBlock);
-                }
+                WriteCloseBraceIndent();
+                PopBlock(NewLineKind.BeforeNextBlock);
             }
         }
 
@@ -217,7 +214,7 @@ namespace CppSharp.Generators.CSharp
             {
                 PushBlock(CSharpBlockKind.Functions);
                 WriteLine("public unsafe partial class {0}",
-                    TranslationUnit.FileNameWithoutExtension);
+                    TranslationUnits[0].FileNameWithoutExtension);
                 WriteStartBraceIndent();
 
                 PushBlock(CSharpBlockKind.InternalsClass);
@@ -459,14 +456,17 @@ namespace CppSharp.Generators.CSharp
             typePrinter.PushContext(CSharpTypePrinterContextKind.Native);
 
             GenerateClassFields(@class, GenerateClassInternalsField, true);
-            if (Options.GenerateVirtualTables && @class.IsDynamic)
-                GenerateVTablePointers(@class);
+            if (@class.IsGenerated)
+            {
+                if (Options.GenerateVirtualTables && @class.IsDynamic)
+                    GenerateVTablePointers(@class);
 
-            var functions = GatherClassInternalFunctions(@class);
+                var functions = GatherClassInternalFunctions(@class);
 
-            foreach (var function in functions)
-            {
-                GenerateInternalFunction(function);
+                foreach (var function in functions)
+                {
+                    GenerateInternalFunction(function);
+                }
             }
 
             typePrinter.PopContext();
@@ -583,7 +583,7 @@ namespace CppSharp.Generators.CSharp
             if (@class != null && @class.HasBaseClass)
                 Write("new ");
 
-            WriteLine("struct Internal");
+            WriteLine("partial struct Internal");
         }
 
         private void GenerateStructMarshalingProperties(Class @class)
diff --git a/src/Generator/Generators/Template.cs b/src/Generator/Generators/Template.cs
index f5b8c567..9748f745 100644
--- a/src/Generator/Generators/Template.cs
+++ b/src/Generator/Generators/Template.cs
@@ -247,17 +247,12 @@ namespace CppSharp.Generators
 
     public abstract class Template : ITextGenerator
     {
-        [Flags]
-        public enum Order
-        {
-            First = 1,
-            InBetween = 2,
-            Last = 4
-        }
-
         public Driver Driver { get; private set; }
         public DriverOptions Options { get; private set; }
-        public TranslationUnit TranslationUnit { get; private set; }
+        public List<TranslationUnit> TranslationUnits { get; private set; }
+
+        public TranslationUnit TranslationUnit { get { return TranslationUnits[0]; } }
+
         public IDiagnosticConsumer Log
         {
             get { return Driver.Diagnostics; }
@@ -268,16 +263,16 @@ namespace CppSharp.Generators
 
         public abstract string FileExtension { get; }
 
-        protected Template(Driver driver, TranslationUnit unit)
+        protected Template(Driver driver, IEnumerable<TranslationUnit> units)
         {
             Driver = driver;
             Options = driver.Options;
-            TranslationUnit = unit;
+            TranslationUnits = new List<TranslationUnit>(units);
             RootBlock = new Block();
             ActiveBlock = RootBlock;
         }
 
-        public abstract void Process(Order order);
+        public abstract void Process();
 
         public string Generate()
         {
diff --git a/tests/CSharpTemp/AnotherUnit.cpp b/tests/CSharpTemp/AnotherUnit.cpp
new file mode 100644
index 00000000..f49e6d30
--- /dev/null
+++ b/tests/CSharpTemp/AnotherUnit.cpp
@@ -0,0 +1,5 @@
+#include "AnotherUnit.h"
+
+void functionInAnotherUnit()
+{
+}
diff --git a/tests/CSharpTemp/AnotherUnit.h b/tests/CSharpTemp/AnotherUnit.h
new file mode 100644
index 00000000..95c5d84f
--- /dev/null
+++ b/tests/CSharpTemp/AnotherUnit.h
@@ -0,0 +1,3 @@
+#include "../Tests.h"
+
+void DLL_API functionInAnotherUnit();
diff --git a/tests/CSharpTemp/CSharpTemp.Tests.cs b/tests/CSharpTemp/CSharpTemp.Tests.cs
index 5dac132e..ec7e2b92 100644
--- a/tests/CSharpTemp/CSharpTemp.Tests.cs
+++ b/tests/CSharpTemp/CSharpTemp.Tests.cs
@@ -138,4 +138,10 @@ public class CSharpTempTests : GeneratorTestFixture
         methodsWithDefaultValues.DefaultValueType();
         methodsWithDefaultValues.DefaultIntAssignedAnEnum();
     }
+
+    [Test]
+    public void TestGenerationOfAnotherUnitInSameFile()
+    {
+        AnotherUnit.FunctionInAnotherUnit();
+    }
 }
\ No newline at end of file
diff --git a/tests/CSharpTemp/CSharpTemp.cs b/tests/CSharpTemp/CSharpTemp.cs
index 91525525..0798d0b9 100644
--- a/tests/CSharpTemp/CSharpTemp.cs
+++ b/tests/CSharpTemp/CSharpTemp.cs
@@ -69,6 +69,7 @@ namespace CppSharp.Tests
             driver.TranslationUnitPasses.AddPass(new TestAttributesPass());
             driver.Options.MarshalCharAsManagedChar = true;
             driver.Options.GenerateDefaultValuesForArguments = true;
+            driver.Options.GenerateSingleCSharpFile = true;
         }
 
         public override void Preprocess(Driver driver, ASTContext ctx)