godot-host-test/addons/GodotDI/Editor/Plugin.cs

#if TOOLS
using System.Linq;
using System.Reflection;
using Godot;
using Array = Godot.Collections.Array;

namespace GodotHostTest.GodotDI.Editor;

[Tool]
public partial class Plugin : EditorPlugin
{
    /// <summary>
    ///     Amount of debounce before updating the injections
    /// </summary>
    private const double DebounceDelay = 0.1;

    private const BindingFlags InjectBindingFlags =
        BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic;

    private ScopeInspectorPlugin? _plugin;
    private SceneTreeTimer? _sceneTreeTimer;

    private bool _treeChangesRegistered;

    public override void _EnterTree()
    {
        GD.Print("Initializing Godot DI");

        // add the inspector plugin to provide a better interface for the "scope" property
        _plugin = new ScopeInspectorPlugin();
        AddInspectorPlugin(_plugin);

        // update injected nodes when the scene changes
        SceneChanged += OnSceneChanged;
    }

    public override void _ExitTree()
    {
        GD.Print("Deinitializing Godot DI");

        SceneChanged -= OnSceneChanged;

        UnregisterTreeChanges();

        RemoveInspectorPlugin(_plugin);
        _plugin?.Dispose();
        _plugin = null;
    }

    private void OnTreeChanged()
    {
        var root = EditorInterface.Singleton.GetEditedSceneRoot();
        OnSceneChanged(root);
    }

    private void OnSceneChanged(Node? sceneRoot)
    {
        // register to tree changes, in case the registration wasn't possible before
        RegisterTreeChanges();

        // don't do anything if there's no scene root
        if (sceneRoot == null) return;

        // apply debounce
        if (_sceneTreeTimer != null)
        {
            _sceneTreeTimer.SetTimeLeft(DebounceDelay);
            return;
        }

        _sceneTreeTimer = GetTree().CreateTimer(DebounceDelay);
        _sceneTreeTimer.Timeout += () => ActualOnSceneChanged(sceneRoot);
    }

    private void ActualOnSceneChanged(Node sceneRoot)
    {
        // mark debounce as completed
        _sceneTreeTimer = null;

        var sceneNodes = sceneRoot.FindChildren("*").Where(node => node.Owner == sceneRoot).Append(sceneRoot).ToList();
        GD.Print($"Updating injection for root {sceneRoot.Name}");

        // look for the service source
        var serviceSource = sceneNodes.Where(child => child.IsServiceSource())
            .Single(out var failureReason);

        switch (failureReason)
        {
            case EnumerableExtensions.SingleFailureReason.LessThanOne:
                return;
            case EnumerableExtensions.SingleFailureReason.MoreThanOne:
                GD.PushWarning("Found more than one service source owned by the scene.");
                return;
            default:
            {
                // find all nodes that need injection
                var nodesNeedingInjection = sceneNodes.Where(node =>
                {
                    var type = node.GetCSharpScriptType();
                    return type != null && (
                        type.GetFields(InjectBindingFlags).Any(field =>
                            field.GetCustomAttributes(typeof(InjectAttribute), true).Length > 0) ||
                        type.GetMethods(InjectBindingFlags).Any(field =>
                            field.GetCustomAttributes(typeof(InjectAttribute), true).Length > 0)
                    );
                }).ToList();

                // add all nodes needed to the scope
                var injectedNodes = serviceSource!.Get(ServiceSource.PropertyName.InjectedNodes).AsGodotArray();
                Array? newArray = null;
                foreach (var node in nodesNeedingInjection
                             .Where(node => !injectedNodes.Contains(node)))
                {
                    newArray ??= new Array(injectedNodes);
                    newArray.Add(node);
                    GD.Print($"Added node {node.Name} to scope {serviceSource.GetPath()}.");
                }

                // remove all nodes that are no longer needing injection
                foreach (var node in injectedNodes
                             .Select(n => (Node?)n.AsGodotObject())
                             .Where(node => !nodesNeedingInjection.Contains(node)))
                {
                    if (node == null) continue;
                    newArray ??= new Array(injectedNodes);
                    newArray.Remove(node);
                    GD.Print($"Removed node {node.Name} to scope {serviceSource.GetPath()}.");
                }

                // apply changes to the scene if necessary
                if (newArray != null)
                {
                    GD.Print("Apply changes.");
                    var undoRedo = GetUndoRedo();
                    undoRedo.CreateAction("Update injected nodes");
                    undoRedo.AddDoProperty(serviceSource, ServiceSource.PropertyName.InjectedNodes, newArray);
                    undoRedo.AddUndoProperty(serviceSource, ServiceSource.PropertyName.InjectedNodes, injectedNodes);
                    undoRedo.CommitAction();
                }

                break;
            }
        }
    }

    /// <summary>
    ///     Attach to the TreeChanged event of the SceneTree if it's not already registered and if there's a scene tree to
    ///     attach to from the currently edited scene.
    /// </summary>
    private void RegisterTreeChanges()
    {
        if (_treeChangesRegistered) return;
        var sceneRoot = EditorInterface.Singleton.GetEditedSceneRoot();
        if (sceneRoot == null) return;
        var tree = sceneRoot.GetTree();
        tree.TreeChanged += OnTreeChanged;
        _treeChangesRegistered = true;
    }

    private void UnregisterTreeChanges()
    {
        var sceneRoot = EditorInterface.Singleton.GetEditedSceneRoot();
        if (!_treeChangesRegistered || sceneRoot == null) return;
        var tree = sceneRoot.GetTree();
        tree.TreeChanged -= OnTreeChanged;
        SceneChanged -= OnSceneChanged;
    }
}
#endif