natsdotnet/src/NATS.Server/Internal/SubjectTree/Parts.cs

// Go reference: server/stree/parts.go, server/stree/util.go
namespace NATS.Server.Internal.SubjectTree;

/// <summary>
/// Subject tokenization helpers and match logic for the ART.
/// </summary>
internal static class Parts
{
    // For subject matching.
    internal const byte Pwc = (byte)'*';
    internal const byte Fwc = (byte)'>';
    internal const byte Tsep = (byte)'.';

    /// <summary>
    /// No pivot available sentinel value (DEL character).
    /// </summary>
    internal const byte NoPivot = 127;

    /// <summary>
    /// Returns the pivot byte at the given position, or NoPivot if past end.
    /// Go reference: server/stree/util.go:pivot
    /// </summary>
    internal static byte Pivot(ReadOnlySpan<byte> subject, int pos)
    {
        if (pos >= subject.Length) return NoPivot;
        return subject[pos];
    }

    /// <summary>
    /// Returns the length of the common prefix between two byte spans.
    /// Go reference: server/stree/util.go:commonPrefixLen
    /// </summary>
    internal static int CommonPrefixLen(ReadOnlySpan<byte> s1, ReadOnlySpan<byte> s2)
    {
        var limit = Math.Min(s1.Length, s2.Length);
        int i = 0;
        for (; i < limit; i++)
        {
            if (s1[i] != s2[i]) break;
        }
        return i;
    }

    /// <summary>
    /// Copy bytes helper.
    /// </summary>
    internal static byte[] CopyBytes(ReadOnlySpan<byte> src)
    {
        if (src.Length == 0) return [];
        return src.ToArray();
    }

    /// <summary>
    /// Break a filter subject into parts based on wildcards (pwc '*' and fwc '>').
    /// Go reference: server/stree/parts.go:genParts
    /// </summary>
    internal static ReadOnlyMemory<byte>[] GenParts(ReadOnlySpan<byte> filter)
    {
        var parts = new List<ReadOnlyMemory<byte>>();
        // We work on a copy since ReadOnlyMemory needs a backing array
        var filterArr = filter.ToArray();
        var filterMem = new ReadOnlyMemory<byte>(filterArr);
        int start = 0;
        int e = filterArr.Length - 1;

        for (int i = 0; i < filterArr.Length; i++)
        {
            if (filterArr[i] == Tsep)
            {
                // See if next token is pwc. Either internal or end pwc.
                if (i < e && filterArr[i + 1] == Pwc && ((i + 2 <= e && filterArr[i + 2] == Tsep) || i + 1 == e))
                {
                    if (i > start)
                    {
                        parts.Add(filterMem.Slice(start, i + 1 - start));
                    }
                    parts.Add(filterMem.Slice(i + 1, 1));
                    i++; // Skip pwc
                    if (i + 2 <= e)
                    {
                        i++; // Skip next tsep from next part too.
                    }
                    start = i + 1;
                }
                else if (i < e && filterArr[i + 1] == Fwc && i + 1 == e)
                {
                    if (i > start)
                    {
                        parts.Add(filterMem.Slice(start, i + 1 - start));
                    }
                    parts.Add(filterMem.Slice(i + 1, 1));
                    i++; // Skip fwc
                    start = i + 1;
                }
            }
            else if (filterArr[i] == Pwc || filterArr[i] == Fwc)
            {
                // Wildcard must be at the start or preceded by tsep.
                int prev = i - 1;
                if (prev >= 0 && filterArr[prev] != Tsep)
                {
                    continue;
                }

                // Wildcard must be at the end or followed by tsep.
                int next = i + 1;
                if (next == e || (next < e && filterArr[next] != Tsep))
                {
                    continue;
                }

                // Full wildcard must be terminal.
                if (filterArr[i] == Fwc && i < e)
                {
                    break;
                }

                // We start with a pwc or fwc.
                parts.Add(filterMem.Slice(i, 1));
                if (i + 1 <= e)
                {
                    i++; // Skip next tsep from next part too.
                }
                start = i + 1;
            }
        }

        if (start < filterArr.Length)
        {
            // Check to see if we need to eat a leading tsep.
            if (filterArr[start] == Tsep)
            {
                start++;
            }
            parts.Add(filterMem[start..]);
        }

        return [.. parts];
    }

    /// <summary>
    /// Match parts against a fragment (prefix for nodes or suffix for leaves).
    /// Go reference: server/stree/parts.go:matchParts
    /// </summary>
    internal static (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchPartsAgainstFragment(
        ReadOnlyMemory<byte>[] parts, ReadOnlySpan<byte> frag)
    {
        int lf = frag.Length;
        if (lf == 0)
        {
            return (parts, true);
        }

        int si = 0;
        int lpi = parts.Length - 1;

        for (int i = 0; i < parts.Length; i++)
        {
            if (si >= lf)
            {
                return (parts[i..], true);
            }

            var part = parts[i].Span;
            int lp = part.Length;

            // Check for pwc or fwc place holders.
            if (lp == 1)
            {
                if (part[0] == Pwc)
                {
                    var index = frag[si..].IndexOf(Tsep);
                    // We are trying to match pwc and did not find our tsep.
                    if (index < 0)
                    {
                        if (i == lpi)
                        {
                            return ([], true);
                        }
                        return (parts[i..], true);
                    }
                    si += index + 1;
                    continue;
                }
                else if (part[0] == Fwc)
                {
                    return ([], true);
                }
            }

            int end = Math.Min(si + lp, lf);
            // If part is bigger than the remaining fragment, adjust to a portion of the part.
            var partToCompare = part;
            if (si + lp > end)
            {
                // Frag is smaller than part itself.
                partToCompare = part[..(end - si)];
            }

            if (!partToCompare.SequenceEqual(frag[si..end]))
            {
                return (parts, false);
            }

            // If we still have a portion of the fragment left, update and continue.
            if (end < lf)
            {
                si = end;
                continue;
            }

            // If we matched a partial, do not move past current part
            // but update the part to what was consumed.
            if (end < si + lp)
            {
                if (end >= lf)
                {
                    // Create a copy with the current part trimmed.
                    var newParts = new ReadOnlyMemory<byte>[parts.Length - i];
                    Array.Copy(parts, i, newParts, 0, newParts.Length);
                    newParts[0] = parts[i][(lf - si)..];
                    return (newParts, true);
                }
                else
                {
                    i++;
                }
                return (parts[i..], true);
            }

            if (i == lpi)
            {
                return ([], true);
            }

            // If we are here we are not the last part which means we have a wildcard
            // gap, so we need to match anything up to next tsep.
            si += part.Length;
        }

        return (parts, false);
    }
}