MeWeb/src/MeWeb.Core/Knowledge/UpstreamTally.cs

using System.Buffers.Binary;

namespace MeWeb.Core.Knowledge;

public sealed class UpstreamTally
{
  readonly Dictionary<uint, Dictionary<uint, double>> weights = [];
  readonly Dictionary<uint, int> seen = [];

  public int Collectors { get; private set; }

  public int Add(Stream mrt)
  {
    var counts = new Dictionary<uint, Dictionary<uint, int>>();
    var header = new byte[12];
    var body = new byte[1 << 16];
    while (Sources.Fill(mrt, header))
    {
      var type = BinaryPrimitives.ReadUInt16BigEndian(header.AsSpan(4));
      var subtype = BinaryPrimitives.ReadUInt16BigEndian(header.AsSpan(6));
      var length = (int)BinaryPrimitives.ReadUInt32BigEndian(header.AsSpan(8));
      if (body.Length < length) body = new byte[length];
      if (!Sources.Fill(mrt, body.AsSpan(0, length))) break;
      if (type != 13 || subtype is not (2 or 4 or 8 or 10)) continue;
      Rib(body.AsSpan(0, length), subtype is 8 or 10, counts);
    }
    foreach (var (origin, neighbours) in counts)
    {
      var total = (double)neighbours.Values.Sum();
      if (!weights.TryGetValue(origin, out var sum)) weights[origin] = sum = [];
      foreach (var (neighbour, count) in neighbours) sum[neighbour] = sum.GetValueOrDefault(neighbour) + count / total;
      seen[origin] = seen.GetValueOrDefault(origin) + 1;
    }
    Collectors++;
    return counts.Count;
  }

  public Dictionary<uint, List<uint>> Upstreams(int keep = 4, double share = 0.05)
  {
    var result = new Dictionary<uint, List<uint>>();
    foreach (var (origin, neighbours) in weights)
    {
      var collectors = (double)seen[origin];
      var chosen = neighbours.Where(n => n.Value / collectors >= share).OrderByDescending(n => n.Value).ThenBy(n => n.Key).Take(keep).Select(n => n.Key).ToList();
      if (chosen.Count > 0) result[origin] = chosen;
    }
    return result;
  }

  static void Rib(ReadOnlySpan<byte> rib, bool addPath, Dictionary<uint, Dictionary<uint, int>> counts)
  {
    if (rib.Length < 5) return;
    var at = 4;
    var bits = rib[at++];
    at += (bits + 7) / 8;
    if (at + 2 > rib.Length) return;
    var entries = BinaryPrimitives.ReadUInt16BigEndian(rib[at..]);
    at += 2;
    var head = addPath ? 10 : 6;
    for (var e = 0; e < entries && at + head + 2 <= rib.Length; e++)
    {
      at += head;
      var size = BinaryPrimitives.ReadUInt16BigEndian(rib[at..]);
      at += 2;
      if (at + size > rib.Length) return;
      Attributes(rib.Slice(at, size), counts);
      at += size;
    }
  }

  static void Attributes(ReadOnlySpan<byte> attributes, Dictionary<uint, Dictionary<uint, int>> counts)
  {
    var at = 0;
    while (at + 3 <= attributes.Length)
    {
      var flags = attributes[at];
      var type = attributes[at + 1];
      var extended = (flags & 0x10) != 0;
      if (extended && at + 4 > attributes.Length) return;
      var length = extended ? BinaryPrimitives.ReadUInt16BigEndian(attributes[(at + 2)..]) : attributes[at + 2];
      var start = at + (extended ? 4 : 3);
      if (start + length > attributes.Length) return;
      if (type == 2)
      {
        Path(attributes.Slice(start, length), counts);
        return;
      }
      at = start + length;
    }
  }

  static void Path(ReadOnlySpan<byte> path, Dictionary<uint, Dictionary<uint, int>> counts)
  {
    uint last = 0;
    uint before = 0;
    var distinct = 0;
    var endsInSet = false;
    var at = 0;
    while (at + 2 <= path.Length)
    {
      var segment = path[at];
      var count = path[at + 1];
      at += 2;
      if (at + count * 4 > path.Length) return;
      for (var i = 0; i < count; i++)
      {
        var number = BinaryPrimitives.ReadUInt32BigEndian(path[(at + i * 4)..]);
        if (segment != 2)
        {
          endsInSet = true;
          continue;
        }
        endsInSet = false;
        if (number == last) continue;
        before = last;
        last = number;
        distinct++;
      }
      at += count * 4;
    }
    if (endsInSet || distinct < 3 || Reserved(before) || Reserved(last)) return;
    if (!counts.TryGetValue(last, out var neighbours)) counts[last] = neighbours = [];
    neighbours[before] = neighbours.GetValueOrDefault(before) + 1;
  }

  static bool Reserved(uint number) => number is 0 or 23456 or (>= 64496 and <= 131071) or >= 4200000000;
}