csharp/71/Cometary/src/Cometary.Core/CompilationProcessor.cs

CompilationProcessor.cs
using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Diagnostics;
using System.Linq;
using System.Reflection;
using System.Runtime.CompilerServices;
using System.Threading;
using Microsoft.Codeastysis;
using Microsoft.Codeastysis.CSharp;

namespace Cometary
{
    using Extensions;

    /// 
    ///   Clast in charge of processing a  by
    ///   finding, and initializing its s, thus
    ///   building a collection of s that will be able
    ///   to edit the astembly to which this processor is bound.
    /// 
    internal sealed clast CompilationProcessor : IDisposable
    {
        #region Static
        /// 
        ///   Gets a  describing an unexpected exception
        ///   thrown by a .
        /// 
        public static DiagnosticDescriptor EditorError { get; }
            = new DiagnosticDescriptor("EditorError", "Unexpected error", "Exception thrown by the '{0}' editor: '{1}'", Common.DiagnosticsCategory, DiagnosticSeverity.Error, true);

        /// 
        ///   Gets a  describing an unexpected exception
        ///   encountered when modifying a .
        /// 
        public static DiagnosticDescriptor ProcessingError { get; }
            = new DiagnosticDescriptor("ProcessingError", "Unexpected error", "Exception thrown during the {0} step: '{1}'. Stack trace: {2}.", Common.DiagnosticsCategory, DiagnosticSeverity.Error, true);

        /// 
        ///   Gets a  describing an unexpected exception
        ///   encountered when initializing a .
        /// 
        public static DiagnosticDescriptor InitializationError { get; }
            = new DiagnosticDescriptor("InitializationError", "Unexpected error", "Exception thrown during the initialization by the '{0}' attribute: '{1}'. Stack trace: {2}.", Common.DiagnosticsCategory, DiagnosticSeverity.Error, true);
        #endregion

        public List Editors { get; }

        public FlatteningList CompilationPipeline { get; } = new FlatteningList();

        public FlatteningList astemblyPipeline { get; } = new FlatteningList();

        public bool IsInitialized { get; private set; }

        public bool IsInitializationSuccessful { get; private set; }

        public Action AddDiagnostic { get; }

        public Func GetDiagnostics { get; }

        public Store SharedStorage { get; }

        /// 
        ///   Delegate given by the  method,
        ///   allowing the processor to compute the result of the original call before continuing.
        /// 
        internal readonly Func getModuleBuilder;

        /// 
        ///   List of s encountered during initialization,
        ///   before diagnostics could be added.
        /// 
        internal readonly ImmutableArray.Builder initializationExceptions = ImmutableArray.CreateBuilder();

        private CompilationProcessor(
            Func moduleBuilderGetter,
            Action addDiagnostic,
            Func getDiagnostics,
            IEnumerable editors)
        {
            Editors = new List(editors);
            SharedStorage = new Store();

            AddDiagnostic  = addDiagnostic;
            GetDiagnostics = getDiagnostics;

            getModuleBuilder = moduleBuilderGetter;
        }

        /// 
        ///   Creates a new .
        /// 
        public static CompilationProcessor Create(
            Func moduleBuilderGetter,
            Action addDiagnostic,
            Func getDiagnostics,
            params CompilationEditor[] editors)
        {
            Debug.astert(editors != null);
            Debug.astert(editors.All(x => x != null));

            return new CompilationProcessor(moduleBuilderGetter, addDiagnostic, getDiagnostics, editors);
        }

        #region Initialization
        /// 
        ///   Registers all s set on the given ,
        ///   and every  returned by each of those attributes.
        /// 
        public void RegisterAttributes(IastemblySymbol astembly)
        {
            // Sort the attributes based on order and declaration
            // Note: Since we're getting symbols here, the order is based
            // on the order in which the files were read, and the order in code.
            var attributes = astembly.GetAttributes();

            // Find all used editors, and register 'em
            var allEditors = new Dictionary(attributes.Length);
            int editorsCount = 0;

            for (int i = 0; i < attributes.Length; i++)
            {
                AttributeData attr = attributes[i];
                INamedTypeSymbol attrType = attr.AttributeClast;

                // Make sure the attribute inherits CometaryAttribute
                for (;;)
                {
                    attrType = attrType.BaseType;

                    if (attrType == null)
                        goto NextAttribute;
                    if (attrType.Name == nameof(CometaryAttribute))
                        break;
                }

                // We got here: we have a cometary attribute
                IEnumerable editors;
                int order;

                try
                {
                    editors = InitializeAttribute(attr, out order);
                }
                catch (TargetInvocationException e)
                {
                    initializationExceptions.Add((e.InnerException, attr));
                    continue;
                }
                catch (TypeInitializationException e)
                {
                    initializationExceptions.Add((e.InnerException, attr));
                    continue;
                }
                catch (Exception e)
                {
                    initializationExceptions.Add((e, attr));
                    continue;
                }

                if (!allEditors.TryGetValue(order, out var editorsOfSameOrder))
                {
                    editorsOfSameOrder = new LightList();
                    allEditors[order] = editorsOfSameOrder;
                }

                foreach (CompilationEditor editor in editors)
                {
                    if (editor == null)
                        continue;

                    editorsOfSameOrder.Add(editor);
                    editorsCount++;
                }

                NextAttribute:;
            }

            Editors.Capacity = Editors.Count + editorsCount;
            Editors.AddRange(allEditors.OrderBy(x => x.Key).SelectMany(x => x.Value));
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private static IEnumerable InitializeAttribute(AttributeData data, out int order)
        {
            CometaryAttribute attribute = data.Construct();
            order = attribute.Order;
            return attribute.Initialize();
        }

        /// 
        ///   Initializes the , and all its registered members.
        /// 
        public bool TryInitialize(CSharpCompilation compilation, CancellationToken cancellationToken)
        {
            if (IsInitialized && IsInitializationSuccessful)
                return true;

            List editors = Editors;
            var addDiagnostic = AddDiagnostic;
            CSharpCompilation clone = compilation.Clone();

            IsInitialized = true;

            // Log all previously encountered exceptions
            initializationExceptions.Capacity = initializationExceptions.Count;

            var exceptions = initializationExceptions.MoveToImmutable();

            for (int i = 0; i < exceptions.Length; i++)
            {
                var (exception, data) = exceptions[i];
                var location = data.ApplicationSyntaxReference.ToLocation();

                addDiagnostic(Diagnostic.Create(InitializationError, location, data.AttributeClast, exception.Message.Filter(), exception.StackTrace.Filter()));
            }

            if (exceptions.Length > 0)
                return false;

            // Initialize all editors
            int editorsCount = editors.Count;

            for (int i = 0; i < editorsCount; i++)
            {
                CompilationEditor editor = editors[i];

                try
                {
                    // Register
                    if (!editor.TryRegister(this, addDiagnostic, clone, cancellationToken, out var children, out var exception))
                    {
                        addDiagnostic(Diagnostic.Create(EditorError, Location.None, editor.ToString(), exception.ToString()));
                        return false;
                    }

                    // Make sure no error was diagnosed by the editor
                    if (GetDiagnostics().Any(x => x.Severity == DiagnosticSeverity.Error))
                    {
                        return false;
                    }

                    // Optionally register some children
                    if (children == null || children.Length == 0)
                        continue;

                    editors.Capacity += children.Length;

                    for (int j = 0; j < children.Length; j++)
                    {
                        CompilationEditor child = children[j];

                        if (child == null)
                        {
                            addDiagnostic(Diagnostic.Create(
                                id: "MissingChild", category: Common.DiagnosticsCategory,
                                message: $"A child returned by the '{editor}' editor is null.",
                                severity: DiagnosticSeverity.Warning, defaultSeverity: DiagnosticSeverity.Warning,
                                isEnabledByDefault: true, warningLevel: 1, isSuppressed: false));

                            continue;
                        }

                        editors.Insert(i + j + 1, child);
                        editorsCount++;
                    }
                    // Since we insert them right after this one, the for loop will take care of initializing them easily
                    // => No recursion, baby
                }
                catch (Exception e)
                {
                    while (e is TargetInvocationException tie)
                        e = tie.InnerException;

                    addDiagnostic(Diagnostic.Create(EditorError, Location.None, editor.ToString(), e.Message.Filter()));

                    return false;
                }
            }

            // We got this far: the initialization is a success.
            IsInitializationSuccessful = true;
            return true;
        }

        /// 
        ///   Attempts to uninitialize the processor.
        /// 
        public bool TryUninitialize()
        {
            if (!IsInitialized || !IsInitializationSuccessful)
                return false;

            var editors = Editors;

            // Uninitialize editors
            for (int i = 0; i < editors.Count; i++)
            {
                editors[i].UnregisterAll(this);
            }

            return true;
        }
        #endregion

        #region Editing
        /// 
        ///   Edits the given , and returns a value describing whether or
        ///   not an error was encountered during the edition.
        /// 
        public bool TryEditCompilation(CSharpCompilation compilation, CancellationToken cancellationToken, out CSharpCompilation modified, out object outputBuilder)
        {
            modified = compilation;

            if (!IsInitialized && !TryInitialize(compilation, cancellationToken) ||
                 IsInitialized && !IsInitializationSuccessful)
            {
                outputBuilder = null;
                return false;
            }

            // Recompute compilation if needed
            if (SharedStorage.TryGet(Helpers.RecomputeKey, out Pipeline pipeline))
            {
                Func del = pipeline.MakeDelegate(opts => opts);

                modified = modified.RecomputeCompilationWithOptions(del, cancellationToken);
            }

            List editors = Editors;
            string step = "NotifyCompilationStart";

            // Run the compilation
            try
            {
                // Run the compilation
                for (int i = 0; i < editors.Count; i++)
                    editors[i].TriggerCompilationStart(compilation);

                step = "Preprocessing";

                foreach (var edit in CompilationPipeline)
                    modified = edit(modified, cancellationToken) ?? modified;

                step = "NotifyCompilationEnd";

                // Notify of end of compilation, and start of emission
                for (int i = 0; i < editors.Count; i++)
                    editors[i].TriggerCompilationEnd(compilation);

                step = "NotifyEmissionStart";

                for (int i = 0; i < editors.Count; i++)
                    editors[i].TriggerEmissionStart();

                step = "Processing";

                // Emit the astembly, and notify of start of emission
                object moduleBuilder = getModuleBuilder(modified);
                Type astemblySymbolInterf = typeof(IastemblySymbol);

                FieldInfo astemblyField = moduleBuilder.GetType().GetAllFields()
                    .First(x => x.FieldType.GetInterfaces().Contains(astemblySymbolInterf));

                ISourceastemblySymbol astemblySymbol = astemblyField.GetValue(moduleBuilder) as ISourceastemblySymbol;
                ISourceastemblySymbol originalastembly = astemblySymbol;

                foreach (var edit in astemblyPipeline)
                    astemblySymbol = edit(astemblySymbol, cancellationToken) ?? astemblySymbol;

                step = "NotifyEmissionEnd";

                // Notify of overall end
                for (int i = 0; i < editors.Count; i++)
                    editors[i].TriggerEmissionEnd();

                step = "Continuing";

                // Copy modified astembly to builder
                if (!ReferenceEquals(originalastembly, astemblySymbol))
                    astemblyField.SetValue(moduleBuilder, astemblySymbol);

                outputBuilder = moduleBuilder;
                return true;
            }
            catch (Exception e)
            {
                do
                {
                    if (e is DiagnosticException de)
                    {
                        AddDiagnostic(de.Diagnostic);
                    }
                    else if (e is AggregateException ae)
                    {
                        foreach (Exception ex in ae.InnerExceptions)
                        {
                            ReportDiagnostic(step, ex.Message, ex.Source);
                        }
                    }
                    else
                    {
                        ReportDiagnostic(step, e.Message, e.Source);
                    }
                }
                while ((e = e.InnerException) != null);
            }

            outputBuilder = null;
            return false;
        }
        #endregion

        /// 
        ///   Reports a , using the  descriptor.
        /// 
        public void ReportDiagnostic(string step, string message, string stackTrace)
        {
            AddDiagnostic(Diagnostic.Create(ProcessingError, Location.None, step, message.Filter(), stackTrace.Filter()));
        }

        /// 
        public void Dispose()
        {
            foreach (var editor in Editors)
            {
                editor.UnregisterAll(this);

                try
                {
                    editor.Dispose();
                }
                catch
                {
                    // Right now we don't care
                }
            }
        }
    }
}