ModspecModelGenerator.cs

/*
 * Copyright (c) 2025 Brill Power.
 *
 * SPDX-License-Identifier: Apache-2.0
 */
using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Linq;
using System.Text;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.Text;
using Modspec.Model.Extensions;

namespace Modspec.Model.Generation;

/// <summary>
/// Generates typed proxy classes from Modspec JSON schemata.
/// </summary>
[Generator]
public class ModelGenerator : IIncrementalGenerator
{
    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        var pipeline = context.AdditionalTextsProvider
            .Where(static (file) => Path.GetFileName(file.Path).EndsWith("json"))
            .Select(static (model, cancellationToken) =>
            {
                string path = model.Path;
                ModelCompiler.TryGenerate(path, out string? code);
                return (path, code);
            })
            .Where(static (pair) => !String.IsNullOrEmpty(pair.code));

        context.RegisterSourceOutput(pipeline,
            static (context, pair) =>
            {
                string newName = Path.GetFileName(pair.path).Replace(".json", ".generated.cs");
                context.AddSource(newName, SourceText.From(pair.code!, Encoding.UTF8));
            });
    }

    private class ModelCompiler
    {
        public static bool TryGenerate(string path, [NotNullWhen(true)] out string? result)
        {
            result = default;
            Schema? schema;
            using (FileStream stream = new FileStream(path, FileMode.Open, FileAccess.Read))
            {
                // this will fail loudly on a parse error
                schema = Schema.GetSchema(stream);
            }

            using (StringWriter mainWriter = new StringWriter())
            using (StringWriter appendixWriter = new StringWriter())
            {
                mainWriter.WriteLine($"""
// This code has been generated by a tool.
// Do not modify it. Your changes will be overwritten.
#nullable enable
using System;
using System.Buffers.Binary;
using System.Collections.Generic;
using System.ComponentModel;
using System.ComponentModel.DataAnnotations;
using System.Text;
using System.Threading.Tasks;
using Modspec.Model;

namespace {schema.Name};

""");
                mainWriter.WriteLine($"public class {schema.Name}Client");
                mainWriter.WriteLine("{");
                mainWriter.WriteLine("\tprivate readonly IModbusClient _client;");
                List<string> bufferInitialisers = [];
                List<string> fieldInitialisers = [];
                List<ConstructorParameter> constructorParams = [];
                bool hasChangeDetection = false;
                WriteGroups(schema.Groups, mainWriter, appendixWriter, bufferInitialisers, fieldInitialisers, constructorParams, ref hasChangeDetection);

                foreach (RepeatingGroup repeatingGroup in schema.RepeatingGroups)
                {
                    List<string> groupBufferInitialisers = [];
                    List<string> groupFieldInitialisers = [];
                    List<ConstructorParameter> groupConstructorParams = [];
                    string repeatingGroupName = Pluralise(repeatingGroup.Name);
                    string repeatingGroupFieldName = ToFieldName(repeatingGroupName);
                    // backing list
                    mainWriter.WriteLine($"\tprivate readonly List<{repeatingGroup.Name}> {repeatingGroupFieldName};");
                    mainWriter.WriteLine();
                    // public read-only property to expose
                    mainWriter.WriteLine($"\tpublic IReadOnlyList<{repeatingGroup.Name}> {repeatingGroupName} => {repeatingGroupFieldName};");
                    mainWriter.WriteLine();
                    string repeatingGroupCountParamName = $"{repeatingGroup.Count.Name}Count";
                    ConstructorParameter repeatingGroupCountParam = new(repeatingGroupCountParamName, repeatingGroup.Count.MaxValue);
                    if (!constructorParams.Contains(repeatingGroupCountParam))
                    {
                        constructorParams.Add(repeatingGroupCountParam);
                    }
                    mainWriter.WriteLine($"\tpublic class {repeatingGroup.Name}");
                    mainWriter.WriteLine("\t{");
                    mainWriter.WriteLine("\t\tprivate readonly IModbusClient _client;");
                    mainWriter.WriteLine("\t\tprivate readonly int _offset;"); // the offset to the base offset for this particular repeated element
                    mainWriter.WriteLine();
                    bool groupHasChangeDetection = false;
                    WriteGroups(repeatingGroup.Groups, mainWriter, appendixWriter, groupBufferInitialisers, groupFieldInitialisers, groupConstructorParams, ref groupHasChangeDetection, "\t", "_offset + ");
                    if (groupHasChangeDetection)
                    {
                        hasChangeDetection = true;
                        mainWriter.WriteLine("\t\tprivate readonly BitfieldChangedCallback? _onBitfieldChanged;");
                    }
                    groupFieldInitialisers.Add("_offset = offset;");
                    if (groupHasChangeDetection)
                    {
                        groupFieldInitialisers.Add("_onBitfieldChanged = onBitfieldChanged;");
                    }
                    fieldInitialisers.Add($"{repeatingGroupFieldName} = new List<{repeatingGroup.Name}>();");
                    fieldInitialisers.Add($"for (int i = 0; i < {repeatingGroupCountParam.Name}; i++)");
                    fieldInitialisers.Add("{");
                    string repeatingGroupParameterRefs = String.Empty;
                    if (groupConstructorParams.Count > 0)
                    {
                        repeatingGroupParameterRefs = ", " + String.Join(", ", groupConstructorParams.Select(cp => cp.Name));
                    }
                    string changeDetectionRef = groupHasChangeDetection ? ", onBitfieldChanged" : "";
                    fieldInitialisers.Add($"\t{repeatingGroupFieldName}.Add(new {repeatingGroup.Name}(client, i * {repeatingGroup.Every}{repeatingGroupParameterRefs}{changeDetectionRef}));");
                    fieldInitialisers.Add("}");
                    groupConstructorParams.Insert(0, new("offset", UInt16.MaxValue));
                    WriteFieldsAndConstructor(repeatingGroup.Name, mainWriter, groupBufferInitialisers, groupFieldInitialisers, groupConstructorParams, "\t", groupHasChangeDetection);
                    mainWriter.WriteLine("\t}");
                    mainWriter.WriteLine();
                    // need to expose inner constructor parameters to parent client constructor
                    foreach (ConstructorParameter constructorParam in groupConstructorParams)
                    {
                        if (constructorParam.Name == "offset")
                        {
                            continue;
                        }
                        if (!constructorParams.Contains(constructorParam))
                        {
                            constructorParams.Add(constructorParam);
                        }
                    }
                }

                if (hasChangeDetection)
                {
                    mainWriter.WriteLine("\tprivate readonly BitfieldChangedCallback? _onBitfieldChanged;");
                }
                WriteFieldsAndConstructor(schema.Name + "Client", mainWriter, bufferInitialisers, fieldInitialisers, constructorParams, hasChangeDetection: hasChangeDetection);
                mainWriter.WriteLine("}");
                mainWriter.WriteLine();

                result = mainWriter.ToString() + appendixWriter.ToString();
                return true;
            }
        }

        private static void WriteFieldsAndConstructor(string name, StringWriter mainWriter, List<string> bufferInitialisers, List<string> fieldInitialisers, List<ConstructorParameter> constructorParams, string indent = "", bool hasChangeDetection = false)
        {
            mainWriter.Write($"\t{indent}public {name}(IModbusClient client");
            if (constructorParams.Count > 0)
            {
                mainWriter.Write(", ");
                mainWriter.Write(String.Join(", ", constructorParams.Select(cp => $"int {cp.Name}")));
            }
            if (hasChangeDetection)
            {
                mainWriter.Write(", BitfieldChangedCallback? onBitfieldChanged = null");
            }
            mainWriter.WriteLine(")");
            mainWriter.WriteLine($"{indent}\t{{");
            mainWriter.WriteLine($"{indent}\t\t_client = client;");
            if (hasChangeDetection)
            {
                mainWriter.WriteLine($"{indent}\t\t_onBitfieldChanged = onBitfieldChanged;");
            }
            foreach ((string constructorParamName, int maxCount) in constructorParams)
            {
                mainWriter.WriteLine($"{indent}\t\tif ({constructorParamName} > {maxCount}) throw new ArgumentException(\"{constructorParamName} is greater than the maximum permitted value ({maxCount}).\", \"{constructorParamName}\");");
            }
            foreach (string bufferInitialiser in bufferInitialisers)
            {
                mainWriter.WriteLine($"{indent}\t\t{bufferInitialiser}");
            }
            foreach (string fieldInitialiser in fieldInitialisers)
            {
                mainWriter.WriteLine($"{indent}\t\t{fieldInitialiser}");
            }
            mainWriter.WriteLine($"{indent}\t}}");
        }

        private static void WriteGroups(IReadOnlyCollection<Group> groups, StringWriter mainWriter, StringWriter appendixWriter, List<string> bufferInitialisers, List<string> fieldInitialisers, List<ConstructorParameter> constructorParams, ref bool hasChangeDetection, string indent = "", string readOffsetField = "")
        {
            foreach (Group group in groups)
            {
                string bufferName = $"_buffer{group.Name}";
                string previousName = $"_previous{group.Name}";
                mainWriter.WriteLine($"{indent}\tprivate readonly Memory<byte> {bufferName};");
                // pre-scan to determine if this group has bitfield points with levels
                bool groupHasLevels = group.Points.Any(p =>
                    p.Type.IsBitfield() && p.Symbols is not null &&
                    p.Symbols.Any(s => s.Level.HasValue && s.Level.Value != Level.None));
                if (groupHasLevels)
                {
                    mainWriter.WriteLine($"{indent}\tprivate readonly Memory<byte> {previousName};");
                    hasChangeDetection = true;
                }
                int maxOffset = 0;
                string bufferSize = String.Empty;
                List<BitfieldPointInfo> groupBitfieldPoints = [];
                for (int i = 0; i < group.Points.Count; i++)
                {
                    Point point = group.Points[i];
                    if (point.Count is not null && i < group.Points.Count - 1)
                    {
                        // this check allows us to do optimisation on array reads (i.e. reading only
                        // supplied count of elements, rather than max size of array)
                        throw new InvalidOperationException($"An array must be the last (or only) element in a group.");
                    }
                    WritePoint(point, bufferName, group.Table, mainWriter, appendixWriter, fieldInitialisers, constructorParams, ref maxOffset, ref bufferSize, indent, groupBitfieldPoints);
                }
                if (String.IsNullOrEmpty(bufferSize))
                {
                    bufferSize = $"{maxOffset}";
                }
                bufferInitialisers.Add($"{bufferName} = new byte[{bufferSize}];");
                if (groupHasLevels)
                {
                    bufferInitialisers.Add($"{previousName} = new byte[{bufferSize}];");
                }
                mainWriter.WriteLine();
                // generate Read method for this group
                mainWriter.WriteLine($"\t{indent}public async ValueTask Read{group.Name}Async()");
                mainWriter.WriteLine($"\t{indent}{{");
                // note dependency between table name and Read...Async method on IModbusClient
                mainWriter.WriteLine($"\t\t{indent}await _client.Read{group.Table}Async({readOffsetField}{group.BaseRegister}, {bufferName});");
                if (groupBitfieldPoints.Count > 0)
                {
                    mainWriter.WriteLine($"\t\t{indent}if (_onBitfieldChanged is not null)");
                    mainWriter.WriteLine($"\t\t{indent}{{");
                    mainWriter.WriteLine($"\t\t\t{indent}Check{group.Name}();");
                    mainWriter.WriteLine($"\t\t{indent}}}");
                }
                mainWriter.WriteLine($"\t{indent}}}");
                mainWriter.WriteLine();
                mainWriter.WriteLine($"\t{indent}public void Read{group.Name}()");
                mainWriter.WriteLine($"\t{indent}{{");
                // note dependency between table name and Read... method on IModbusClient
                mainWriter.WriteLine($"\t\t{indent}_client.Read{group.Table}({readOffsetField}{group.BaseRegister}, {bufferName}.Span);");
                if (groupBitfieldPoints.Count > 0)
                {
                    mainWriter.WriteLine($"\t\t{indent}if (_onBitfieldChanged is not null)");
                    mainWriter.WriteLine($"\t\t{indent}{{");
                    mainWriter.WriteLine($"\t\t\t{indent}Check{group.Name}();");
                    mainWriter.WriteLine($"\t\t{indent}}}");
                }
                mainWriter.WriteLine($"\t{indent}}}");
                mainWriter.WriteLine();
                // generate Check method for groups with leveled bitfields
                if (groupBitfieldPoints.Count > 0)
                {
                    mainWriter.WriteLine($"\t{indent}private void Check{group.Name}()");
                    mainWriter.WriteLine($"\t{indent}{{");
                    mainWriter.WriteLine($"\t\t{indent}Span<byte> current = {bufferName}.Span;");
                    mainWriter.WriteLine($"\t\t{indent}Span<byte> previous = {previousName}.Span;");
                    foreach (BitfieldPointInfo bp in groupBitfieldPoints)
                    {
                        mainWriter.WriteLine($"\t\t{indent}if (!current.Slice({bp.Offset}, {bp.SizeInBytes}).SequenceEqual(previous.Slice({bp.Offset}, {bp.SizeInBytes})))");
                        mainWriter.WriteLine($"\t\t{indent}{{");
                        mainWriter.WriteLine($"\t\t\t{indent}{bp.PointName} oldValue = ({bp.PointName}){bp.ReadMethod}(previous.Slice({bp.Offset}, {bp.SizeInBytes}));");
                        mainWriter.WriteLine($"\t\t\t{indent}_onBitfieldChanged!(\"{bp.PointName}\", oldValue, {bp.PointName}, {bp.PointName}.GetLevel());");
                        mainWriter.WriteLine($"\t\t{indent}}}");
                    }
                    mainWriter.WriteLine($"\t\t{indent}current.CopyTo(previous);");
                    mainWriter.WriteLine($"\t{indent}}}");
                    mainWriter.WriteLine();
                }
            }
        }

        private static void WritePoint(Point point, string bufferName, Table table, StringWriter mainWriter, StringWriter appendixWriter, List<string> fieldInitialisers, List<ConstructorParameter> constructorParams, ref int maxOffset, ref string bufferSize, string indent = "", List<BitfieldPointInfo>? groupBitfieldPoints = null)
        {
            string type;
            string readMethod;
            switch (point.Type)
            {
                case PointType.Enum16:
                case PointType.Bitfield16:
                case PointType.UInt16:
                case PointType.Acc16:
                    type = "ushort";
                    readMethod = "BinaryPrimitives.ReadUInt16BigEndian";
                    break;
                case PointType.Enum32:
                case PointType.Bitfield32:
                case PointType.UInt32:
                case PointType.Acc32:
                    type = "uint";
                    readMethod = "BinaryPrimitives.ReadUInt32BigEndian";
                    break;
                case PointType.Enum64:
                case PointType.Bitfield64:
                case PointType.UInt64:
                case PointType.Acc64:
                    type = "ulong";
                    readMethod = "BinaryPrimitives.ReadUInt64BigEndian";
                    break;
                case PointType.String:
                    type = "string";
                    readMethod = "ModbusString.ReadNullTerminatedString";
                    break;
                case PointType.Int16:
                    type = "short";
                    readMethod = "BinaryPrimitives.ReadInt16BigEndian";
                    break;
                case PointType.Int32:
                    type = "int";
                    readMethod = "BinaryPrimitives.ReadInt32BigEndian";
                    break;
                case PointType.Int64:
                    type = "long";
                    readMethod = "BinaryPrimitives.ReadInt64BigEndian";
                    break;
                case PointType.Float32:
                    type = "float";
                    readMethod = "BinaryPrimitives.ReadSingleBigEndian";
                    break;
                case PointType.Float64:
                    type = "double";
                    readMethod = "BinaryPrimitives.ReadDoubleBigEndian";
                    break;
                case PointType.Padding:
                    maxOffset += point.SizeInBytes;
                    return;
                default:
                    throw new NotSupportedException($"Register {point.Name} at position {maxOffset} of type {point.Type} is not supported.");
            }

            if (table == Table.Coils || table == Table.DiscreteInputs)
            {
                if (point.Type.IsBitfield())
                {
                    // when representing coils/discrete inputs as enums, we should read them
                    // as little-endian
                    readMethod = readMethod.Replace("BigEndian", "LittleEndian");
                }
                else
                {
                    throw new NotSupportedException($"Register {point.Name} at position {maxOffset} is invalid; coil and discrete input tables only support bitfields.");
                }
            }

            string propertyType = type;
            string transform = String.Empty;
            if ((point.ScaleFactor ?? 1) != 1 || (point.Offset ?? 0) != 0)
            {
                // all scaled properties are doubles; this could be refined (if, say, an integer offset)
                // but this will do for now
                propertyType = "double";
                if (point.ScaleFactor != 1)
                {
                    transform = $" * {point.ScaleFactor}";
                }
                if ((point.Offset ?? 0) != 0)
                {
                    transform += $" + {point.Offset}";
                }
            }
            else if (point.Symbols is not null && point.Symbols.Count > 0 && point.Type.IsEnumOrBitfield())
            {
                // enums and bitmasks
                // generate a type for the enum at the end of the file
                bool isFlags = point.Type.IsBitfield();
                if (isFlags)
                {
                    appendixWriter.WriteLine("[Flags]");
                }
                if (!String.IsNullOrEmpty(point.Description))
                {
                    appendixWriter.WriteLine($"[Description(\"{point.Description}\")]");
                }
                appendixWriter.WriteLine($"public enum {point.Name} : {type}");
                appendixWriter.WriteLine("{");
                // enum members
                Dictionary<Level, ulong> masksByLevel = [];
                foreach (Symbol symbol in point.Symbols)
                {
                    string cast = String.Empty;
                    if (symbol.Value >= 31)
                    {
                        // cast to correct type before shift
                        cast = $"({type})";
                    }
                    string value = isFlags ? $"{cast}1 << {symbol.Value}" : $"{symbol.Value}";
                    appendixWriter.WriteLine($"\t{symbol.Name} = {value},");
                    if (isFlags && symbol.Level.HasValue && symbol.Level.Value != Level.None)
                    {
                        if (!masksByLevel.TryGetValue(symbol.Level.Value, out ulong mask))
                        {
                            mask = 0;
                        }
                        masksByLevel[symbol.Level.Value] = mask |= 1UL << symbol.Value;
                    }
                }
                appendixWriter.WriteLine("}");
                appendixWriter.WriteLine();
                if (isFlags && masksByLevel.Count > 0)
                {
                    groupBitfieldPoints?.Add(new BitfieldPointInfo(point.Name, maxOffset, point.SizeInBytes, readMethod));
                    appendixWriter.WriteLine($"public static class {point.Name}Extensions");
                    appendixWriter.WriteLine("{");
                    appendixWriter.WriteLine($"\tpublic static Level GetLevel(this {point.Name} self)");
                    appendixWriter.WriteLine("\t{");
                    foreach (Level level in Enum.GetValues<Level>().Skip(1).Reverse())
                    {
                        if (masksByLevel.TryGetValue(level, out ulong mask))
                        {
                            appendixWriter.WriteLine($"\t\tif ((self & ({point.Name})0x{mask:x}) != 0)");
                            appendixWriter.WriteLine("\t\t{");
                            appendixWriter.WriteLine($"\t\t\treturn Level.{level};");
                            appendixWriter.WriteLine("\t\t}");
                        }
                    }
                    appendixWriter.WriteLine("\t\treturn Level.None;");
                    appendixWriter.WriteLine("\t}");
                    appendixWriter.WriteLine("}");
                    appendixWriter.WriteLine();
                }
                if (point.Count is null)
                {
                    // read must cast to enum type
                    // (for arrays, we will do this in the transform lambda below)
                    readMethod = $"({point.Name})" + readMethod;
                }
                propertyType = point.Name;
            }

            if (point.Count is not null)
            {
                // array handling
                // (we have already checked that this is the last point in the group)
                propertyType = $"ModbusArray<{type}, {propertyType}>";
                if (point.Type.IsEnumOrBitfield())
                {
                    transform = $"x => ({point.Name})x";
                }
                else
                {
                    transform = $"x => x" + transform;
                }
                string fieldName = ToFieldName(point.Name);
                mainWriter.WriteLine();
                mainWriter.WriteLine($"{indent}\tprivate readonly {propertyType} {fieldName};"); // field
                string constructorParamName = $"{point.Count.Name}Count";
                ConstructorParameter constructorParam = new(constructorParamName, point.Count.MaxValue);
                if (!constructorParams.Contains(constructorParam))
                {
                    constructorParams.Add(constructorParam);
                }
                fieldInitialisers.Add($"{fieldName} = new {propertyType}({bufferName}.Slice({maxOffset}, {point.SizeInBytes} * {constructorParamName}), {point.SizeInBytes}, {readMethod}, {transform});");
                readMethod = $"ref {fieldName}";
                propertyType = $"ref readonly {propertyType}";
                bufferSize = $"{maxOffset} + ({point.SizeInBytes} * {constructorParamName})";
            }
            else
            {
                readMethod = $"{readMethod}({bufferName}.Span.Slice({maxOffset}, {point.SizeInBytes})){transform}";
            }

            mainWriter.WriteLine();
            if (!String.IsNullOrEmpty(point.Label))
            {
                mainWriter.WriteLine($"{indent}\t[DisplayName(\"{point.Label}\")]");
            }
            if (!String.IsNullOrEmpty(point.Description))
            {
                mainWriter.WriteLine($"{indent}\t[Description(\"{point.Description}\")]");
            }
            if (point.MinValue.HasValue && point.MaxValue.HasValue)
            {
                mainWriter.WriteLine($"{indent}\t[Range({point.MinValue}, {point.MaxValue})]");
            }
            if (point.Type == PointType.String)
            {
                mainWriter.WriteLine($"{indent}\t[MaxLength({point.Length})]");
            }
            mainWriter.WriteLine($"{indent}\tpublic {propertyType} {point.Name}");
            mainWriter.WriteLine($"{indent}\t{{");
            mainWriter.WriteLine($"{indent}\t\tget {{ return {readMethod}; }}");
            mainWriter.WriteLine($"{indent}\t}}");
            maxOffset += point.SizeInBytes * (point.Count?.MaxValue ?? 1);
        }

        private static string ToFieldName(string name)
        {
            Span<char> result = stackalloc char[name.Length + 1];
            name.CopyTo(result.Slice(1));
            result[0] = '_';
            result[1] = Char.ToLower(name[0]);
            return new string(result);
        }

        private static string Pluralise(string self)
        {
            switch (self)
            {
                case string x when x.EndsWith("um"):
                    return String.Create(self.Length - 1, self, (neue, old) =>
                    {
                        ReadOnlySpan<char> chars = old;
                        chars.Slice(0, chars.Length - 2).CopyTo(neue);
                        neue[neue.Length - 1] = 'a';
                    });
                case string x when x.EndsWith("us"):
                    return String.Create(self.Length - 1, self, (neue, old) =>
                    {
                        ReadOnlySpan<char> chars = old;
                        chars.Slice(0, chars.Length - 2).CopyTo(neue);
                        neue[neue.Length - 1] = 'i';
                    });
                case string x when x.EndsWith("y") && !x.EndsWith("ay"):
                    return String.Create(self.Length + 2, self, (neue, old) =>
                    {
                        ReadOnlySpan<char> chars = old;
                        chars.Slice(0, chars.Length - 1).CopyTo(neue);
                        ReadOnlySpan<char> ies = "ies";
                        ies.CopyTo(neue.Slice(neue.Length - 3));
                    });
                default:
                    return String.Create(self.Length + 1, self, (neue, old) =>
                    {
                        ReadOnlySpan<char> chars = old;
                        chars.CopyTo(neue);
                        neue[neue.Length - 1] = 's';
                    });
            }
        }
    }

    private record ConstructorParameter(string Name, int Count);
    private record BitfieldPointInfo(string PointName, int Offset, int SizeInBytes, string ReadMethod);
}