Question: Some libraries cannot seemingly clear data in custom derivate...

[onzag]

EDIT: I figured it out, the issue was not in the walker but a rather strange implementation in those libraries that were using a global cache (guess that's why globals aren't too good), you could replicate it with react dom itself in special cases; I'd still not mind if you could check this code, just in case there's something blatant, feel free to close the question, your codebase helped me a lot.

I think what I wrote also may work as a guide for those that need to create a custom async walker, mine just has the mere basics so it's simpler to understand.

==== Original message

Back in the days of react 15 before suspense was even a thing I had come with a workaround to get promises SSR to work, it was hacky but it worked, the process was different, today I am updating that project to react 17 to use some libraries and I've been reading this source code as inspiration.

So I created this custom function based on what I understood from this library, however, there's an issue, some libraries, such as emotion-js renders (using the standard ReactDOM) differs after passing through the walker, and I can't tell why, it's as if, they refuse to be called twice; but I tested with reactDOM and they work being called twice producing the same output, my mechanism must be causing it.

You may notice the usage of getDerivedServerSideStateFromProps that is basically used to setup SSR, this is similar to NextJS, I manage to inject that state later both in server and client.

The resulting HTML hydration is equal, I get no warnings of wrong HTML, it's simply the style tags coming from emotion-js which I assume, it's somehow writing by hand.

====

Can you spot any issues?...

import uuid from "uuid";
import React from "react";

const {
  ReactCurrentDispatcher
} = (React as any).__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED;

function noop(): void { }

let currentContextMap: Map<any, any> = null;

function readContextValue(context: any) {
  const value = currentContextMap.get(context);
  if (value !== undefined) {
    return value
  }

  // Return default if context has no value yet
  return context._currentValue;
}

function readContext(context: any, _: void | number | boolean) {
  return readContextValue(context)
}

function useContext(context: any, _: void | number | boolean) {
  return readContextValue(context)
}

function useState<S>(
  initialState: (() => S) | S
): [S, any] {
  return [typeof initialState === "function" ? (initialState as any)() : initialState, noop];
}

function useReducer<S, I, A>(
  reducer: any,
  initialArg: any,
  init?: any
): [S, any] {
  return [initialArg, noop];
}

function useMemo<T>(memoValue: () => T, deps: Array<any>): T {
  return memoValue();
}

function useRef<T>(initialValue: T): { current: T } {
  return {
    current: initialValue,
  }
}

function useOpaqueIdentifier(): string {
  // doesn't matter much since we don't use this anywhere for practical usage
  return "R:" + uuid.v1();
}

function useCallback<T>(callback: T, deps: Array<any>): T {
  return callback;
}

function useMutableSource<Source, Snapshot>(
  source: any,
  getSnapshot: any,
  _subscribe: any
): Snapshot {
  return getSnapshot(source._source)
}

function useTransition(): [any, boolean] {
  return [noop, false]
}

function useDeferredValue<T>(input: T): T {
  return input
}

export const Dispatcher = {
  readContext,
  useContext,
  useMemo,
  useReducer,
  useRef,
  useState,
  useCallback,
  useMutableSource,
  useTransition,
  useDeferredValue,
  useOpaqueIdentifier,
  useId: useOpaqueIdentifier,
  // ignore useLayout effect completely as usage of it will be caught
  // in a subsequent render pass
  useLayoutEffect: noop,
  // useImperativeHandle is not run in the server environment
  useImperativeHandle: noop,
  // Effects are not run in the server environment.
  useEffect: noop,
  // Debugging effect
  useDebugValue: noop
}

const symbolFor = Symbol.for;
const REACT_ELEMENT_TYPE = symbolFor('react.element');
const REACT_PORTAL_TYPE = symbolFor('react.portal');
const REACT_FRAGMENT_TYPE = symbolFor('react.fragment');
const REACT_STRICT_MODE_TYPE = symbolFor('react.strict_mode');
const REACT_PROFILER_TYPE = symbolFor('react.profiler');
const REACT_PROVIDER_TYPE = symbolFor('react.provider');
const REACT_CONTEXT_TYPE = symbolFor('react.context');
const REACT_CONCURRENT_MODE_TYPE = Symbol.for('react.concurrent_mode');
const REACT_FORWARD_REF_TYPE = symbolFor('react.forward_ref');
const REACT_SUSPENSE_TYPE = symbolFor('react.suspense');
const REACT_MEMO_TYPE = symbolFor('react.memo');
const REACT_LAZY_TYPE = symbolFor('react.lazy');

function isConstructed(element: any) {
  return element.type && element.type.prototype && element.type.prototype.isReactComponent;
}

function getType(element: any) {
  switch (element.$$typeof) {
    case REACT_PORTAL_TYPE:
      return REACT_PORTAL_TYPE
    case REACT_ELEMENT_TYPE:
    case REACT_MEMO_TYPE:                       // memo are treated as element because they are a memo of something
      switch (element.type) {
        case REACT_CONCURRENT_MODE_TYPE:
          return REACT_CONCURRENT_MODE_TYPE
        case REACT_FRAGMENT_TYPE:
          return REACT_FRAGMENT_TYPE
        case REACT_PROFILER_TYPE:
          return REACT_PROFILER_TYPE
        case REACT_STRICT_MODE_TYPE:
          return REACT_STRICT_MODE_TYPE
        case REACT_SUSPENSE_TYPE:
          return REACT_SUSPENSE_TYPE

        default: {
          switch (element.type && element.type.$$typeof) {
            case REACT_LAZY_TYPE:
              return REACT_LAZY_TYPE
            case REACT_MEMO_TYPE:
              return REACT_MEMO_TYPE
            case REACT_CONTEXT_TYPE:
              return REACT_CONTEXT_TYPE
            case REACT_PROVIDER_TYPE:
              return REACT_PROVIDER_TYPE
            case REACT_FORWARD_REF_TYPE:
              return REACT_FORWARD_REF_TYPE
            default:
              return REACT_ELEMENT_TYPE
          }
        }
      }

    default:
      return undefined
  }
}

export async function walkReactTree(element: any, contextMap: Map<any, any> = new Map()): Promise<void> {
  if (element === null || typeof element === "undefined") {
    return;
  }
  if (Array.isArray(element)) {
    await Promise.all(element.map((c) => walkReactTree(c, contextMap)));
  }

  const type = getType(element);
  const isBaseType = type === REACT_ELEMENT_TYPE && typeof element.type === "string";
  if (type === REACT_ELEMENT_TYPE && !isBaseType) {
    const isComponentType = isConstructed(element);

    if (isComponentType) {
      const instance = new element.type(element.props);

      if (instance.state === undefined) {
        instance.state = null
      }

      if (element.type.getDerivedStateFromProps) {
        const newState = element.type.getDerivedStateFromProps(instance.props, instance.state);
        if (newState) {
          instance.state = Object.assign({}, instance.state, newState)
        }
      } else if (typeof instance.componentWillMount === "function") {
        instance.componentWillMount()
      } else if (typeof instance.UNSAFE_componentWillMount === "function") {
        instance.UNSAFE_componentWillMount()
      }

      instance._isMounted = true;

      if (element.type.getDerivedServerSideStateFromProps) {
        const newState = await element.type.getDerivedServerSideStateFromProps(instance.props, instance.state);
        if (newState) {
          instance.state = Object.assign({}, instance.state, newState)
        }
      }

      try {
        const childComponent = instance.render();
        await walkReactTree(childComponent, contextMap);
      } catch (err) {
        if (typeof element.type.getDerivedStateFromError === "function") {
          const newState = element.type.getDerivedStateFromError(err);
          if (newState) {
            instance.state = Object.assign({}, instance.state, newState)
          }

          const childComponent = instance.render();
          await walkReactTree(childComponent, contextMap);
        } else {
          throw err;
        }
      }

      if (
        typeof instance.getDerivedStateFromProps !== "function" &&
        (typeof instance.componentWillMount === "function" ||
          typeof instance.UNSAFE_componentWillMount === "function") &&
        typeof instance.componentWillUnmount === 'function'
      ) {
        try {
          instance.componentWillUnmount()
        } catch (_err) {}
      }

      instance._isMounted = false;

    } else {
      const originalDispatcher = ReactCurrentDispatcher.current;

      currentContextMap = contextMap;
      ReactCurrentDispatcher.current = Dispatcher;
      const childComponent = element.type(element.props);

      currentContextMap = null;
      ReactCurrentDispatcher.current = originalDispatcher;

      await walkReactTree(childComponent, contextMap);
    }
  } else if (type === REACT_PROVIDER_TYPE) {
    const newContextMap = new Map(contextMap);
    newContextMap.set(element.type._context, element.props.value);
    await walkReactTree(element.props.children, newContextMap);
  } else if (type === REACT_CONTEXT_TYPE) {
    const contextualValue = contextMap.get(element.type._context);
    const apparentChild = element.props.children(contextualValue);
    await walkReactTree(apparentChild, contextMap);
  } else if (
    isBaseType ||
    type === REACT_FRAGMENT_TYPE ||
    type === REACT_SUSPENSE_TYPE ||
    type === REACT_STRICT_MODE_TYPE ||
    type === REACT_PROFILER_TYPE
  ) {
    await walkReactTree(element.props.children, contextMap);
  } else if (type === REACT_FORWARD_REF_TYPE) {
    const originalDispatcher = ReactCurrentDispatcher.current;

    currentContextMap = contextMap;
    ReactCurrentDispatcher.current = Dispatcher;
    const childComponent = element.type.render(element.props, element.ref);

    currentContextMap = null;
    ReactCurrentDispatcher.current = originalDispatcher;

    await walkReactTree(childComponent, contextMap);
  } else if (type === REACT_MEMO_TYPE) {
    const elementMemoed = element.type;
    await walkReactTree(elementMemoed, contextMap);
  }
}