Học từ Telegram #5: Store tự viết và diff danh sách có key

import { useLayoutEffect, useReducer, useRef } from 'react';

/**
 * Store toàn cục tối giản theo ý tưởng teactn của Telegram Web A (lib/teact/teactn.tsx):
 * - `setGlobal` chỉ đổi tham chiếu state; việc báo cho component được gom lại, chạy
 *   MỘT lần trong requestAnimationFrame kế tiếp (nhiều action trong một frame = một lần tính lại).
 * - Mỗi component đăng ký một "container" với hàm `mapState`. Khi flush, chạy lại mọi
 *   mapState và so sánh nông (shallow) với kết quả cũ, chỉ component có props đổi mới render.
 * - Công cụ debug: cảnh báo mapState trả NaN, cảnh báo flush chậm, đếm số lần render.
 *
 * Khác bản gốc: một mapState ném lỗi không làm dừng việc cập nhật các container khác,
 * và id container không bị sinh lại ở mỗi lần render (xem bài viết).
 */

type Mapper<S, P> = (state: S) => P;

interface Container<S> {
  id: number;
  name: string;
  mapState: Mapper<S, Record<string, unknown>>;
  mapped: Record<string, unknown>;
  forceUpdate: () => void;
  updates: number;
}

export interface StoreDebugEvent {
  type: 'nan' | 'slow' | 'error' | 'flush';
  message: string;
}

export interface StoreOptions {
  /** ngưỡng cảnh báo flush chậm (ms), Telegram dùng 7ms */
  slowFlushMs?: number;
  /** nơi nhận cảnh báo debug, mặc định console.warn */
  onDebug?: (event: StoreDebugEvent) => void;
  /** hàm lên lịch flush, mặc định requestAnimationFrame (để test có thể truyền hàm đồng bộ) */
  schedule?: (cb: () => void) => void;
}

export function shallowEqual(a: Record<string, unknown>, b: Record<string, unknown>) {
  if (a === b) return true;
  const aKeys = Object.keys(a);
  if (aKeys.length !== Object.keys(b).length) return false;
  return aKeys.every((key) => Object.is(a[key], b[key]));
}

/** Liệt kê props khác nhau, để biết "vì sao component này render lại". */
export function diffProps(a: Record<string, unknown>, b: Record<string, unknown>) {
  const keys = new Set([...Object.keys(a), ...Object.keys(b)]);
  return [...keys].filter((key) => !Object.is(a[key], b[key]));
}

export function createGlobalStore<S>(initial: S, options: StoreOptions = {}) {
  const {
    slowFlushMs = 7,
    onDebug = (event) => console.warn(`[store] ${event.message}`),
    schedule = (cb) => requestAnimationFrame(cb),
  } = options;

  let current = initial;
  let flushScheduled = false;
  let nextContainerId = 1;
  let flushCount = 0;
  const containers = new Map<number, Container<S>>();
  const listeners = new Set<() => void>();

  /**
   * Chạy mapState có bắt lỗi: lỗi trong một component không được làm sập cả app.
   * Chỉ báo lỗi khi flush; trong lúc render thì im lặng giữ props cũ (báo ở đó dễ thành vòng lặp
   * nếu chính việc báo lỗi làm component render lại).
   */
  function safeMap(container: Container<S>, report = false) {
    try {
      return container.mapState(current);
    } catch (error) {
      if (report) onDebug({ type: 'error', message: `${container.name}: mapState lỗi (${String(error)})` });
      return undefined;
    }
  }

  function flush() {
    flushScheduled = false;
    flushCount++;
    const startedAt = performance.now();

    for (const container of containers.values()) {
      const next = safeMap(container, true);
      // Bản gốc `return` ở đây, nên MỘT mapper lỗi làm mọi container phía sau không được cập nhật.
      if (!next) continue;

      if (Object.values(next).some((value) => Number.isNaN(value))) {
        // NaN !== NaN: store so sánh bằng !== (như bản gốc) sẽ render thừa mỗi lần flush.
        // shallowEqual ở đây dùng Object.is nên không bị, nhưng NaN thường là dấu hiệu tính sai.
        onDebug({ type: 'nan', message: `${container.name}: mapState trả về NaN (với so sánh !== sẽ render lại mỗi lần)` });
      }

      if (!shallowEqual(container.mapped, next)) {
        container.mapped = next;
        container.updates++;
        container.forceUpdate();
      }
    }

    const duration = performance.now() - startedAt;
    if (duration > slowFlushMs) {
      onDebug({ type: 'slow', message: `flush chậm: ${duration.toFixed(1)}ms cho ${containers.size} container` });
    }
    listeners.forEach((listener) => listener());
  }

  /**
   * Đổi state. Mặc định gom theo frame; `sync: true` flush ngay (Telegram dùng cho cử chỉ
   * vuốt back trên iOS, nơi giao diện phải đổi đồng bộ với animation của hệ điều hành).
   */
  function setGlobal(next: S, { sync = false } = {}) {
    if (next === current) return;
    current = next;
    if (sync) {
      flush();
      return;
    }
    if (!flushScheduled) {
      flushScheduled = true;
      schedule(flush);
    }
  }

  function getGlobal() {
    return current;
  }

  /** Tạo action: reducer nhận state hiện tại, trả state mới (hoặc undefined nếu không đổi). */
  function action<A extends unknown[]>(reducer: (state: S, ...args: A) => S | void) {
    return (...args: A) => {
      const next = reducer(current, ...args);
      if (next) setGlobal(next);
    };
  }

  /**
   * Hook nối component với store. `mapState` có thể là hàm inline: kết quả được so sánh nông,
   * nên trả về object mới mỗi lần cũng không sao, miễn các giá trị bên trong giữ tham chiếu.
   */
  function useGlobal<P extends Record<string, unknown>>(mapState: Mapper<S, P>, name = 'Component'): P {
    const [, forceUpdate] = useReducer((n: number) => n + 1, 0);
    const containerRef = useRef<Container<S> | null>(null);

    // Tạo container một lần cho mỗi instance. (Bản gốc gọi generateIdFor() ở mỗi lần render.)
    if (!containerRef.current) {
      containerRef.current = {
        id: nextContainerId++,
        name,
        mapState: mapState as Mapper<S, Record<string, unknown>>,
        mapped: {},
        forceUpdate,
        updates: 0,
      };
      containerRef.current.mapped = safeMap(containerRef.current) ?? {};
    }
    const container = containerRef.current;
    // Luôn dùng mapState mới nhất: nó có thể đóng gói props của component (vd. chatId), nên
    // khi component render lại vì props đổi thì phải tính lại (bản gốc làm khi ownProps đổi).
    if (container.mapState !== mapState) {
      container.mapState = mapState as Mapper<S, Record<string, unknown>>;
      const fresh = safeMap(container);
      if (fresh && !shallowEqual(container.mapped, fresh)) container.mapped = fresh;
    }

    useLayoutEffect(() => {
      containers.set(container.id, container);
      // state có thể đã đổi giữa lúc render và lúc đăng ký
      const latest = safeMap(container);
      if (latest && !shallowEqual(container.mapped, latest)) {
        container.mapped = latest;
        forceUpdate();
      }
      return () => {
        containers.delete(container.id);
      };
    }, [container]);

    return container.mapped as P;
  }

  /** Bảng "ai render bao nhiêu lần", giống double-click debug của Telegram. */
  function getRenderStats() {
    return [...containers.values()]
      .map(({ name, updates }) => ({ name, updates }))
      .sort((a, b) => b.updates - a.updates);
  }

  return {
    getGlobal,
    setGlobal,
    action,
    useGlobal,
    getRenderStats,
    getFlushCount: () => flushCount,
    /** gọi sau mỗi lần flush (dùng cho công cụ debug, devtools) */
    onFlush(listener: () => void) {
      listeners.add(listener);
      return () => {
        listeners.delete(listener);
      };
    },
  };
}

export type GlobalStore<S> = ReturnType<typeof createGlobalStore<S>>;

/**
 * Cập nhật một danh sách DOM có key với ít thao tác DOM nhất, theo ý tưởng `teactFastList`
 * trong renderer tự viết của Telegram Web A (lib/teact/teact-dom.ts):
 *
 * 1. Xoá các node có key không còn trong danh sách mới.
 * 2. Duyệt danh sách mới. Các item MỚI nằm liền nhau được gom vào một DocumentFragment và
 *    chèn một lần (tải thêm 50 tin = 1 lần chèn thay vì 50).
 * 3. Node cũ chỉ bị di chuyển khi vị trí lệch VÀ `order` của nó đổi. Không có `order`
 *    thì không bao giờ di chuyển: Telegram dùng chế độ này khi vị trí hiển thị do CSS quyết định
 *    (vd. `top: i * 72px` trong danh sách chat), thứ tự trong DOM không còn quan trọng.
 */

export interface DomOps {
  inserted: number;
  moved: number;
  removed: number;
  /** số lần chèn bằng fragment (mỗi lần có thể chứa nhiều node) */
  fragments: number;
}

export interface KeyedListOptions<T> {
  getKey: (item: T) => string;
  /** "thứ tự" của item; đổi thì node mới được phép di chuyển trong DOM */
  getOrder?: (item: T, index: number) => number | string | undefined;
  create: (item: T) => HTMLElement;
  update?: (node: HTMLElement, item: T) => void;
}

interface Rendered {
  node: HTMLElement;
  order: number | string | undefined;
}

const renderedByParent = new WeakMap<HTMLElement, Map<string, Rendered>>();

function adoptChildren(parent: HTMLElement) {
  const adopted = new Map<string, Rendered>();
  for (const child of Array.from(parent.children) as HTMLElement[]) {
    if (child.dataset.key !== undefined) adopted.set(child.dataset.key, { node: child, order: undefined });
    else child.remove();
  }
  return adopted;
}

export function renderKeyedList<T>(parent: HTMLElement, items: T[], options: KeyedListOptions<T>): DomOps {
  const { getKey, getOrder, create, update } = options;
  const ops: DomOps = { inserted: 0, moved: 0, removed: 0, fragments: 0 };
  // Lần đầu (hoặc DOM do code khác dựng): nhận các node đã có `data-key` làm trạng thái cũ.
  const previous = renderedByParent.get(parent) ?? adoptChildren(parent);
  const newKeys = new Set(items.map(getKey));

  // 1. Xoá node không còn, ghi lại vị trí (theo thứ tự DOM) của các node còn lại.
  const remaining = new Map<string, Rendered & { index: number }>();
  let remainingIndex = 0;
  for (const [key, rendered] of previous) {
    if (!newKeys.has(key)) {
      rendered.node.remove();
      ops.removed++;
    }
  }
  // thứ tự DOM thật sau khi xoá (Map giữ thứ tự chèn, nhưng node có thể đã bị di chuyển)
  for (const child of Array.from(parent.children) as HTMLElement[]) {
    const key = child.dataset.key;
    const rendered = key !== undefined ? previous.get(key) : undefined;
    if (key !== undefined && rendered && newKeys.has(key)) {
      remaining.set(key, { ...rendered, index: remainingIndex++ });
    }
  }

  const next = new Map<string, Rendered>();
  let queue: { item: T; key: string; index: number }[] = [];
  let queueIndex = 0;
  let preservedIndex = 0;

  const flushQueue = () => {
    if (queue.length === 0) return;
    const reference = parent.children[queueIndex] ?? null;
    if (queue.length === 1) {
      const { item, key, index } = queue[0];
      const node = createNode(item, key, index);
      parent.insertBefore(node, reference);
    } else {
      const fragment = document.createDocumentFragment();
      queue.forEach(({ item, key, index }) => fragment.appendChild(createNode(item, key, index)));
      parent.insertBefore(fragment, reference);
      ops.fragments++;
    }
    ops.inserted += queue.length;
    queue = [];
  };

  const createNode = (item: T, key: string, index: number) => {
    const node = create(item);
    node.dataset.key = key;
    next.set(key, { node, order: getOrder?.(item, index) });
    return node;
  };

  // 2 + 3. Duyệt danh sách mới.
  items.forEach((item, i) => {
    const key = getKey(item);
    const info = remaining.get(key);
    if (!info) {
      if (queue.length === 0) queueIndex = i;
      queue.push({ item, key, index: i });
      return;
    }
    flushQueue();

    const order = getOrder?.(item, i);
    const shouldMove = info.index !== preservedIndex && info.order !== order;
    const isMovingDown = shouldMove && preservedIndex > info.index;
    // Khi node đi xuống, các node phía sau "dồn lên" nên chỉ số giữ chỗ vẫn tăng.
    if (!shouldMove || isMovingDown) preservedIndex++;

    if (shouldMove) {
      const reference = parent.children[isMovingDown ? i + 1 : i] ?? null;
      if (reference !== info.node) {
        parent.insertBefore(info.node, reference);
        ops.moved++;
      }
    }
    update?.(info.node, item);
    next.set(key, { node: info.node, order });
  });
  flushQueue();

  renderedByParent.set(parent, next);
  return ops;
}

/** Cách ngây thơ để so sánh: dựng lại toàn bộ danh sách mỗi lần. */
export function renderByRebuild<T>(parent: HTMLElement, items: T[], options: KeyedListOptions<T>): DomOps {
  const removed = parent.children.length;
  const fragment = document.createDocumentFragment();
  items.forEach((item) => {
    const node = options.create(item);
    node.dataset.key = options.getKey(item);
    fragment.appendChild(node);
  });
  parent.replaceChildren(fragment);
  renderedByParent.delete(parent);
  return { inserted: items.length, moved: 0, removed, fragments: 1 };
}

/** Cách "có key nhưng không khéo": giữ node cũ nhưng appendChild lại tất cả theo thứ tự mới. */
export function renderByReappend<T>(parent: HTMLElement, items: T[], options: KeyedListOptions<T>): DomOps {
  const ops: DomOps = { inserted: 0, moved: 0, removed: 0, fragments: 0 };
  const byKey = new Map<string, HTMLElement>();
  for (const child of Array.from(parent.children) as HTMLElement[]) {
    if (child.dataset.key !== undefined) byKey.set(child.dataset.key, child);
  }
  const keys = new Set(items.map(options.getKey));
  for (const [key, node] of byKey) {
    if (!keys.has(key)) {
      node.remove();
      ops.removed++;
    }
  }
  items.forEach((item) => {
    const key = options.getKey(item);
    let node = byKey.get(key);
    if (!node) {
      node = options.create(item);
      node.dataset.key = key;
      ops.inserted++;
    } else {
      options.update?.(node, item);
      ops.moved++;
    }
    parent.appendChild(node);
  });
  renderedByParent.delete(parent);
  return ops;
}