react/compiler/packages/react-forget-runtime/src/index.ts

/**
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

"use no forget";

import * as React from "react";

const {
  // @ts-ignore
  __SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED: {
    ReactCurrentDispatcher,
  },
  useRef,
  useEffect,
} = React;

type MemoCache = Array<number | typeof $empty>;

const $empty = Symbol.for("react.memo_cache_sentinel");
/**
 * DANGER: this hook is NEVER meant to be called directly!
 **/
export function unstable_useMemoCache(size: number) {
  return React.useState(() => {
    const $ = new Array(size);
    for (let ii = 0; ii < size; ii++) {
      $[ii] = $empty;
    }
    return $;
  })[0];
}

export function $read(memoCache: MemoCache, index: number) {
  const value = memoCache[index];
  if (value === $empty) {
    throw new Error("useMemoCache: read before write");
  }
  return value;
}

const LazyGuardDispatcher: { [key: string]: (...args: Array<any>) => any } = {};
[
  "readContext",
  "useCallback",
  "useContext",
  "useEffect",
  "useImperativeHandle",
  "useInsertionEffect",
  "useLayoutEffect",
  "useMemo",
  "useReducer",
  "useRef",
  "useState",
  "useDebugValue",
  "useDeferredValue",
  "useTransition",
  "useMutableSource",
  "useSyncExternalStore",
  "useId",
  "unstable_isNewReconciler",
  "getCacheSignal",
  "getCacheForType",
  "useCacheRefresh",
].forEach((name) => {
  LazyGuardDispatcher[name] = () => {
    throw new Error(
      `[React] Unexpected React hook call (${name}) from a React Forget compiled function. ` +
        "Check that all hooks are called directly and named according to convention ('use[A-Z]') "
    );
  };
});

let originalDispatcher: unknown = null;

// Allow guards are not emitted for useMemoCache
LazyGuardDispatcher["useMemoCache"] = (count: number) => {
  if (originalDispatcher == null) {
    throw new Error(
      "React Forget internal invariant violation: unexpected null dispatcher"
    );
  } else {
    return (originalDispatcher as any).useMemoCache(count);
  }
};

enum GuardKind {
  PushGuardContext = 0,
  PopGuardContext = 1,
  PushExpectHook = 2,
  PopExpectHook = 3,
}

function setCurrent(newDispatcher: any) {
  ReactCurrentDispatcher.current = newDispatcher;
  return ReactCurrentDispatcher.current;
}

const guardFrames: Array<unknown> = [];

/**
 * When `enableEmitHookGuards` is set, this does runtime validation
 * of the no-conditional-hook-calls rule.
 * As Forget needs to statically understand which calls to move out of
 * conditional branches (i.e. Forget cannot memoize the results of hook
 * calls), its understanding of "the rules of React" are more restrictive.
 * This validation throws on unsound inputs at runtime.
 *
 * Components should only be invoked through React as Forget could memoize
 * the call to AnotherComponent, introducing conditional hook calls in its
 * compiled output.
 * ```js
 * function Invalid(props) {
 *  const myJsx = AnotherComponent(props);
 *  return <div> { myJsx } </div>;
 * }
 *
 * Hooks must be named as hooks.
 * ```js
 * const renamedHook = useState;
 * function Invalid() {
 *   const [state, setState] = renamedHook(0);
 * }
 * ```
 *
 * Hooks must be directly called.
 * ```
 * function call(fn) {
 *  return fn();
 * }
 * function Invalid() {
 *   const result = call(useMyHook);
 * }
 * ```
 */
export function $dispatcherGuard(kind: GuardKind) {
  const curr = ReactCurrentDispatcher.current;
  if (kind === GuardKind.PushGuardContext) {
    // Push before checking invariant or errors
    guardFrames.push(curr);

    if (guardFrames.length === 1) {
      // save if we're the first guard on the stack
      originalDispatcher = curr;
    }

    if (curr === LazyGuardDispatcher) {
      throw new Error(
        `[React] Unexpected call to custom hook or component from a React Forget compiled function. ` +
          "Check that (1) all hooks are called directly and named according to convention ('use[A-Z]') " +
          "and (2) components are returned as JSX instead of being directly invoked."
      );
    }
    setCurrent(LazyGuardDispatcher);
  } else if (kind === GuardKind.PopGuardContext) {
    // Pop before checking invariant or errors
    const lastFrame = guardFrames.pop();

    if (lastFrame == null) {
      throw new Error(
        "React Forget internal error: unexpected null in guard stack"
      );
    }
    if (guardFrames.length === 0) {
      originalDispatcher = null;
    }
    setCurrent(lastFrame);
  } else if (kind === GuardKind.PushExpectHook) {
    // ExpectHooks could be nested, so we save the current dispatcher
    // for the matching PopExpectHook to restore.
    guardFrames.push(curr);
    setCurrent(originalDispatcher);
  } else if (kind === GuardKind.PopExpectHook) {
    const lastFrame = guardFrames.pop();
    if (lastFrame == null) {
      throw new Error(
        "React Forget internal error: unexpected null in guard stack"
      );
    }
    setCurrent(lastFrame);
  } else {
    throw new Error("Forget internal error: unreachable block" + kind);
  }
}

export function $reset($: MemoCache) {
  for (let ii = 0; ii < $.length; ii++) {
    $[ii] = $empty;
  }
}

export function $makeReadOnly() {
  throw new Error("TODO: implement $makeReadOnly in react-forget-runtime");
}

/**
 * Instrumentation to count rerenders in React components
 */
export const renderCounterRegistry: Map<
  string,
  Set<{ count: number }>
> = new Map();
export function clearRenderCounterRegistry() {
  for (const counters of renderCounterRegistry.values()) {
    counters.forEach((counter) => {
      counter.count = 0;
    });
  }
}

function registerRenderCounter(name: string, val: { count: number }) {
  let counters = renderCounterRegistry.get(name);
  if (counters == null) {
    counters = new Set();
    renderCounterRegistry.set(name, counters);
  }
  counters.add(val);
}

function removeRenderCounter(name: string, val: { count: number }): void {
  const counters = renderCounterRegistry.get(name);
  if (counters == null) {
    return;
  }
  counters.delete(val);
}

export function useRenderCounter(name: string): void {
  const val = useRef<{ count: number }>(null);

  if (val.current != null) {
    val.current.count += 1;
  }
  useEffect(() => {
    // Not counting initial render shouldn't be a problem
    if (val.current == null) {
      const counter = { count: 0 };
      registerRenderCounter(name, counter);
      // @ts-ignore
      val.current = counter;
    }
    return () => {
      if (val.current !== null) {
        removeRenderCounter(name, val.current);
      }
    };
  });
}