【Unity】Source GeneratorでMessagePipeのセットアップを行う拡張メソッドを自動生成する

はじめに

MessagePipeというUnityで利用できるメッセージングライブラリがあります。 とても便利なライブラリで自分もよく個人的な開発で使用させていただいているのですが、このライブラリはDI前提の作りになっています。 例えば、以下のようなメッセージMoveInputをMessagePipe上でやり取りしたいという状況を考えます。

public static class SampleMessages
{
    public struct MoveInput
    {
        public UnityEngine.Vector2 Move;
    }
}

メッセージのPublish/Subscribeを行うためには、使用しているDIコンテナでメッセージを配信するためのセットアップを行う必要があります。下記はDIコンテナとしてVContainerを使用した時の例です。

using MessagePipe;
using VContainer;
using VContainer.Unity;

public class SampleLifetimeScope : LifetimeScope
{
    protected override void Configure(IContainerBuilder builder)
    {
        var options = builder.RegisterMessagePipe();
        builder.RegisterMessageBroker<SampleMessages.MoveInput>(options);
    }
}

そしてセットアップされたDIコンテナからのインジェクションを受けることで、メッセージの送受信を行いたいオブジェクトはIPublisher/ISubscriberのインスタンスを取得することができます。以下は方向キーの入力を受け付けて移動するスクリプトの実装例です。SampleInputが入力メッセージであるMoveInputを送信し、SamplePlayerMoveInputを受信して移動処理を行なっています。

using MessagePipe;
using UnityEngine;
using VContainer;

public class SampleInput : MonoBehaviour
{
    [Inject] private IPublisher<SampleMessages.MoveInput> sampleMessagePublisher;

    private void Update()
    {
        var move = new Vector2(Input.GetAxisRaw("Horizontal"), Input.GetAxisRaw("Vertical"));
        // 方向入力を送信
        sampleMessagePublisher.Publish(new SampleMessages.MoveInput { Move = move });
    }
}
using Cysharp.Threading.Tasks;
using MessagePipe;
using UnityEngine;
using VContainer;

public class SamplePlayer : MonoBehaviour
{
    [SerializeField] private float speed = 1f;

    [Inject] private ISubscriber<SampleMessages.MoveInput> messageDSubscriber;
    private Vector2 move;

    private void Start()
    {
        messageDSubscriber
            // 方向入力を受信
            .Subscribe(x => move = x.Move)
            .AddTo(destroyCancellationToken);
    }

    private void Update()
    {
        transform.position += new Vector3(move.x, 0f, move.y) * (speed * Time.deltaTime);
    }
}

このように、MessagePipeを用いることでメッセージの送信側であるSampleInputとメッセージの受信側であるSamplePlayerはお互いのことを一切知らずにメッセージの型であるMoveInputのみを介してやりとりすることができます。
ただこの送受信するメッセージが増えてくると、コンテナへの登録が徐々に大変になってきます。

public static class SampleMessages
{
    public struct MessageA
    {
    }

    public struct MessageB
    {
        public float Value;
    }

    public class MessageC
    {
    }

    public class MessageD
    {
        public bool Value;
    }

    public record MessageE
    {
    }
}
using MessagePipe;
using VContainer;
using VContainer.Unity;

public class SampleLifetimeScope : LifetimeScope
{
    protected override void Configure(IContainerBuilder builder)
    {
        var options = builder.RegisterMessagePipe();
        builder.RegisterMessageBroker<SampleMessages.MessageA>(options);
        builder.RegisterMessageBroker<SampleMessages.MessageB>(options);
        builder.RegisterMessageBroker<SampleMessages.MessageC>(options);
        builder.RegisterMessageBroker<SampleMessages.MessageD>(options);
        builder.RegisterMessageBroker<SampleMessages.MessageE>(options);
    }
}

人力でいちいち登録する処理を記述するのはただ大変なだけではなく、メッセージの登録をうっかり忘れたり間違えたりすることによって本来発行されているはずのメッセージを受け取れないということも発生し得ます。 このようなボイラープレートは何らかの手段で自動生成した方が効率的で間違いを減らせそうです。 そこで今回は、Source Generatorを使い自動でMessageBroker登録用の拡張メソッドを生成しようと思います。

仕様を考える

先ほどのSampleMessagesを例に考えてみます。

public static class SampleMessages
{
    public struct MessageA
    {
    }

    public struct MessageB
    {
        public float Value;
    }

    public class MessageC
    {
    }

    public class MessageD
    {
        public bool Value;
    }

    public record MessageE
    {
    }
}

staticなクラス内に使用するメッセージの型宣言を列挙していますが、このstaticクラスを1つのスコープとみた場合クラス内のメッセージを一括で登録できるようにすれば便利そうです。

public static void RegisterSampleMessages(this IContainerBuilder builder, MessagePipeOptions options)
{
    builder.RegisterMessageBroker<MessageA>(options);
    builder.RegisterMessageBroker<MessageB>(options);
    builder.RegisterMessageBroker<MessageC>(options);
    builder.RegisterMessageBroker<MessageD>(options);
    builder.RegisterMessageBroker<MessageE>(options);
}

このstaticクラスをpartialで宣言し、何らかの属性を付加することでSource Generator側でフックして拡張メソッドを追加するという方針で実装してみます。属性の名前はここではRegistrableMessagesとします。

実装

まず、生成する属性およびpartialクラスのT4テンプレートを作成します。

<#@ template debug="false" hostspecific="false" language="C#" linePragmas="false" #>
#pragma warning disable CS8669
#pragma warning disable CS8625
using System;
namespace RegisterMessageBrokerGenerator
{
    [AttributeUsage(AttributeTargets.Class)]
    internal class RegistrableMessagesAttribute : Attribute
    {
    }
}
<#@ template debug="false" hostspecific="false" language="C#" linePragmas="false" #>
#pragma warning disable CS8669
#pragma warning disable CS8625
using MessagePipe;
using VContainer;

<# if (!string.IsNullOrEmpty(Property.Namespace)) { #>
namespace <#= Property.Namespace #>
{
<# } #>
    public static partial class <#= Property.ClassName #>
    {
        public static void Register<#= Property.ClassName #>(this IContainerBuilder builder, MessagePipeOptions options)
        {
<# foreach (var message in Property.Messages) { #>
            builder.RegisterMessageBroker<<#= message #>>(options);
<# } #>
        }
    }
<# if (!string.IsNullOrEmpty(Property.Namespace)) { #>
}
<# } #>

次に、SourceGeneratorの実装を示します。

using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp.Syntax;

namespace RegisterMessageBrokerGenerator
{
    [Generator(LanguageNames.CSharp)]
    public class SourceGenerator : ISourceGenerator
    {
        public void Initialize(GeneratorInitializationContext context)
        {
            context.RegisterForSyntaxNotifications(SyntaxContextReceiver.Create);
        }

        public void Execute(GeneratorExecutionContext context)
        {
            var registrableAttr = new RegistrableMessagesAttributeTemplate().TransformText();
            context.AddSource("RegistrableMessagesAttribute.g.cs", registrableAttr);

            if (context.SyntaxContextReceiver is not SyntaxContextReceiver receiver ||
                receiver.ClassDeclarations.Count == 0)
            {
                return;
            }

            foreach (var classDeclaration in receiver.ClassDeclarations)
            {
                var innerClassDeclarations = classDeclaration.Members.OfType<TypeDeclarationSyntax>().ToArray();
                var property = new MessagesProperty
                {
                    Namespace = GetNamespaceFrom(classDeclaration),
                    ClassName = classDeclaration.Identifier.ToString(),
                    Messages = innerClassDeclarations.Select(x => x.Identifier.ToString()).ToArray()
                };

                var generatingClass = new MessagesTemplate(property).TransformText();
                context.AddSource($"{classDeclaration.Identifier}.g.cs", generatingClass);
            }
        }

        private static string GetNamespaceFrom(SyntaxNode s) => s.Parent switch
        {
            NamespaceDeclarationSyntax namespaceDeclarationSyntax => namespaceDeclarationSyntax.Name.ToString(),
            null => string.Empty,
            _ => GetNamespaceFrom(s.Parent)
        };
    }

    internal class SyntaxContextReceiver : ISyntaxContextReceiver
    {
        internal static ISyntaxContextReceiver Create()
        {
            return new SyntaxContextReceiver();
        }

        public HashSet<TypeDeclarationSyntax> ClassDeclarations { get; } = new();

        public void OnVisitSyntaxNode(GeneratorSyntaxContext context)
        {
            var node = context.Node;

            // 対象nodeが型宣言でない場合は処理しない
            if (node is not (ClassDeclarationSyntax
                or StructDeclarationSyntax
                or RecordDeclarationSyntax
                or InterfaceDeclarationSyntax)) return;

            var typeSyntax = (TypeDeclarationSyntax)node;
            // 対象nodeが属性を持たない場合は処理しない
            if (typeSyntax.AttributeLists.Count <= 0) return;
            var attr = typeSyntax.AttributeLists.SelectMany(x => x.Attributes)
                .FirstOrDefault(x =>
                {
                    // 属性がRegistrableMessagesの場合は処理する
                    var hasRegistrable = x.Name.ToString() is "RegistrableMessages" or "RegistrableMessagesAttribute";
                    return hasRegistrable;
                });
            if (attr != null)
            {
                ClassDeclarations.Add(typeSyntax);
            }
        }
    }
}

これをビルドしてRoslynAnalyzerのタグをつけた上でUnityに追加することで、動作してくれるようになります。

参考

2022年のC# (Incremental) Source Generator開発手法

SourceGeneratorでらくらく(でもない)ソースコード自動生成