canto/Assets/Scripts/UI/Elements/RecycleView/RecycleVirtualizationController.cs

using System;
using System.Collections.Generic;
using System.Linq;
using KitsuneCafe.System.Collections;
using R3;
using UnityEngine;
using UnityEngine.UIElements;

namespace KitsuneCafe.UI
{

  interface IRecycleVirtualizationContainer : IVirtualizationController
  {
    ILayout Layout { get; set; }
  }

  public class RecycleVirtualizationController : IRecycleVirtualizationContainer, IDisposable
  {
    internal record ItemGeometry(float Position, float Size);

    public abstract record UpdateRequest;
    public record FullUpdate : UpdateRequest;
    public record PartialUpdate(int Index) : UpdateRequest;

    private ICollectionDataSource dataSource;

    public ICollectionDataSource DataSource
    {
      get => dataSource;
      set
      {
        if (dataSource != value)
        {
          dataSource = value;
          Setup();
        }
      }
    }

    private VisualElement container;
    public VisualElement Container
    {
      get => container;
      set
      {
        if (container != value)
        {
          container = value;
          Setup();
        }
      }
    }

    private ILayout layout;
    public ILayout Layout
    {
      get => layout;
      set
      {
        if (layout != value)
        {
          layout = value;
          Setup();
        }
      }
    }

    public FlowDirection Direction => Layout.Direction;

    private readonly Queue<VisualElement> itemPool = new();
    private readonly BiDictionary<int, VisualElement> visibleItems = new();
    private float parentSize = 0;
    private float lastKnownScrollOffset = 0;

    private readonly int itemCountBuffer = 2;

    private readonly Subject<UpdateRequest> requestSubject = new();
    private IDisposable subscriptions;

    public void Setup()
    {
      Dispose();


      if (DataSource == null || DataSource.Length == 0 || Layout == null || Container == null)
      {
        return;
      }

      var full = requestSubject.Where(req => req is FullUpdate)
        .ThrottleLastFrame(1)
        .Subscribe(_ => UpdateVisibleItems());

      var partial = requestSubject.Where(req => req is PartialUpdate)
        .Select(req => (PartialUpdate)req)
        .ChunkFrame(1)
        .Where(requests => requests.Any())
        .Subscribe(requests =>
        {
          var index = requests.Min(req => req.Index);
          UpdateFromIndex(index);
        });

      subscriptions = Disposable.Combine(full, partial);
      // subscriptions = Observable.Merge(full, partial)
      // // var updateSubscription = requestSubject.ThrottleLastFrame(1)
      //   .Subscribe(request =>
      //   {
      //     if (request is PartialUpdate req)
      //     {
      //       UpdateFromIndex(req.Index);
      //     }
      //     else { UpdateVisibleItems(); }
      //   });

      Layout.Update(DataSource.Length, Container);
      ScheduleUpdate();
    }

    public void OnScrolled(float scrollOffset)
    {
      lastKnownScrollOffset = scrollOffset;
      ScheduleUpdate();
    }

    public void OnParentSizeChanged(Vector2 size)
    {
      parentSize = Direction switch
      {
        FlowDirection.Vertical => size.y,
        FlowDirection.Horizontal => size.x,
        _ => throw new NotImplementedException(),
      };

      if (!float.IsNaN(parentSize))
      {
        ScheduleUpdate();
      }
    }

    private void ScheduleUpdate()
    {
      requestSubject.OnNext(new FullUpdate());
    }

    private void ScheduleUpdate(int index)
    {
      requestSubject.OnNext(new PartialUpdate(index));
    }

    private void UpdateSize(VisualElement element, StyleLength size)
    {
      if (Layout is null)
      {
        element.style.width = 0;
        element.style.height = 0;
      }
      else if (Direction == FlowDirection.Vertical)
      {
        element.style.width = StyleKeyword.Auto;
        element.style.height = size;
      }
      else if (Direction == FlowDirection.Horizontal)
      {
        element.style.height = StyleKeyword.Auto;
        element.style.width = size;
      }
    }

    private VisualElement RecycleItem()
    {
      var item = itemPool.Dequeue();
      item.UnregisterCallback<GeometryChangedEvent, int>(OnItemGeometryChanged);

      return item;
    }

    private VisualElement CreateItem(int index)
    {
      var item = DataSource.CreateItem();

      if (Layout is FixedLayout)
      {
        float itemSize = Layout.GetItemSize(index);
        UpdateSize(item, itemSize);
      }

      return item;
    }

    private VisualElement GetItem(int index)
    {
      VisualElement item = itemPool.Count > 0 ? RecycleItem() : CreateItem(index);

      item.RegisterCallback<GeometryChangedEvent, int>(OnItemGeometryChanged, index);

      DataSource.BindItem(item, index);
      visibleItems[index] = item;
      Container.Add(item);

      var position = Layout.GetItemPosition(index);
      if (Direction == FlowDirection.Horizontal)
      {
        item.style.left = position;
      }
      else
      {
        item.style.top = position;
      }


      if (Layout is DynamicLayout layout)
      {
        var size = Direction switch
        {
          FlowDirection.Vertical => item.resolvedStyle.height,
          FlowDirection.Horizontal => item.resolvedStyle.width,
          _ => throw new NotImplementedException()
        };

        if (!float.IsNaN(size))
        {
          layout.SetMeasuredItemSize(index, size);
        }
      }

      return item;
    }

    private void OnItemGeometryChanged(GeometryChangedEvent evt, int index)
    {
      if (Layout is DynamicLayout layout)
      {
        var newSize = Direction == FlowDirection.Horizontal ? evt.newRect.width : evt.newRect.height;
        layout.SetMeasuredItemSize(index, newSize);
        UpdateSize(Container, Layout.ContentSize);
        ScheduleUpdate(index);
      }
    }

    private Range GetDynamicItems(DynamicLayout layout, float scrollOffset)
    {
      var first = layout.GetFirstVisibleIndex(scrollOffset);
      var count = 1;
      var position = layout.GetItemPosition(first);
      var width = position + parentSize;

      foreach (var item in layout.Enumerate(first))
      {
        if (item.Position > width) { break; }
        count += 1;
        position = item.Position;
      }

      return new Range(first, Math.Min(DataSource.Length, first + count));
    }

    private Range GetFixedItems(FixedLayout layout, float scrollOffset)
    {
      var size = layout.ItemSize + layout.GutterSize;
      var first = Mathf.FloorToInt(scrollOffset / (layout.ItemSize + layout.GutterSize));
      var count = Mathf.CeilToInt(parentSize / size) + itemCountBuffer;
      var last = Math.Min(DataSource.Length, first + count);

      return new Range(first, last);
    }

    private void UpdateVisibleItems()
    {
      if (DataSource == null || DataSource.Length == 0 || Mathf.Approximately(parentSize, 0))
      {
        return;
      }


      var scrollOffset = lastKnownScrollOffset;

      var range = Layout switch
      {
        FixedLayout fl => GetFixedItems(fl, scrollOffset),
        DynamicLayout dl => GetDynamicItems(dl, scrollOffset),
        _ => throw new NotImplementedException(),
      };

      var firstIndex = range.Start.Value;
      var lastIndex = range.End.Value;

      var itemsToRemove = new List<int>();
      foreach (var (index, item) in visibleItems)
      {
        if (index < firstIndex || index >= lastIndex)
        {
          item.Value.UnregisterCallback<GeometryChangedEvent, int>(OnItemGeometryChanged);
          DataSource.UnbindItem(item, index);
          itemPool.Enqueue(item);
          Container.Remove(item);
          itemsToRemove.Add(index);
        }
      }

      foreach (int index in itemsToRemove)
      {
        visibleItems.Remove(index);
      }

      for (int i = firstIndex; i < lastIndex; i++)
      {
        if (!visibleItems.ContainsKey(i))
        {
          var item = GetItem(i);

          float position = Layout.GetItemPosition(i);
          if (Direction == FlowDirection.Horizontal)
          {
            item.style.left = position;
          }
          else
          {
            item.style.top = position;
          }
        }
      }
    }

    private void UpdateFromIndex(int index)
    {
      foreach (var (i, item) in visibleItems)
      {
        if (i >= index)
        {
          float position = Layout.GetItemPosition(i);
          if (Direction == FlowDirection.Horizontal)
          {
            item.Value.style.left = position;
          }
          else
          {
            item.Value.style.top = position;
          }
        }
      }
    }

    public float GetContentSize()
    {
      return Layout.ContentSize;
    }

    private void UnregisterCallback(VisualElement element)
    {
      element.UnregisterCallback<GeometryChangedEvent, int>(OnItemGeometryChanged);
    }

    public void Dispose()
    {
      foreach (var item in visibleItems.Values)
      {
        UnregisterCallback(item);
      }

      foreach (var item in itemPool)
      {
        UnregisterCallback(item);
      }


      lastKnownScrollOffset = 0;

      subscriptions?.Dispose();
      Container.Clear();
      itemPool.Clear();
      visibleItems.Clear();
    }

    ~RecycleVirtualizationController()
    {
      Dispose();
    }
  }
}