React internal state tree control

README.md

Example usages of functions

To read state tree into a variable

// Anti-circular reference formatter
const getCircularReplacer = () => {
    const seen = new WeakSet();
    return (key: any, value: any) => {
        if (typeof value === "object" && value !== null) {
            if (seen.has(value)) {
                return;
            }
            seen.add(value);
        }
        return value;
    };
};

function App() {
    const { updateRoot, root, id } = useRoot();
  
    // Find the React 'container's key and use it to extract state tree
    const targetedKey = Object.keys(
        root.current?.memoizedState?.element?.props?.children?.props?.root?.containerInfo ?? {}
    )?.find(key => key.startsWith('__reactContainer');
            
    const stateTree = extractStateTree(
        root.current
            ?.memoizedState
            ?.element
            ?.props
            ?.children
            ?.props
            ?.root
            ?.containerInfo
            // @ts-ignore
            ?.[targetedKey]
            ?.stateNode
            ?.current, 
        undefined,
        id
    );
  
    return (
          <>
              <pre>
                  <code>
                      {JSON.stringify(stateTree, getCircularReplacer(), 2)}
                  </code>
              </pre>
          </>
    );
}

Enforce state tree from input box

// Anti-circular reference formatter
const getCircularReplacer = () => {
    const seen = new WeakSet();
    return (key: any, value: any) => {
        if (typeof value === "object" && value !== null) {
            if (seen.has(value)) {
                return;
            }
            seen.add(value);
        }
        return value;
    };
};

function App() {
    const { updateRoot, root, id } = useRoot();
    const $inputRef = useRef<HTMLTextAreaElement | null>(null);
    
    // Find the React 'container's key and use it to extract state tree
    const targetedKey = Object.keys(
        root.current?.memoizedState?.element?.props?.children?.props?.root?.containerInfo ?? {}
    )?.find(key => key.startsWith('__reactContainer');

    const handleEnforceRoot = () => {
        if ($inputRef.current?.value) {
            flushSync(() => {
                if (!$inputRef.current) return;
                enforceStateTree(JSON.parse($inputRef.current.value), root.current
                    .memoizedState
                    .element
                    .props
                    .children
                    .props
                    .root
                    .containerInfo
                    [targetedKey as any]
                    .stateNode
                    .current
                );
            })
        }
    }

    return (
          <>
              <textarea ref={$inputRef}/>
              <button onClick={() => {
                  updateRoot()
              }}>
                  Update root (re-render)
              </button>
              <button onClick={() => handleEnforceRoot()}>Enforce provided root</button>
          </>
    );
}

export default App;

exampleIndex.tsx

import React from 'react';
import ReactDOM from 'react-dom/client';
import './index.css';
import App from './App';
import reportWebVitals from './reportWebVitals';
import RootContextProvider from "./rootContext";

const root = ReactDOM.createRoot(
  document.getElementById('root') as HTMLElement
);

root.render(
  <React.StrictMode>
    <RootContextProvider root={(root as any)._internalRoot}>
        <App />
    </RootContextProvider>
  </React.StrictMode>
);

// If you want to start measuring performance in your app, pass a function
// to log results (for example: reportWebVitals(console.log))
// or send to an analytics endpoint. Learn more: https://bit.ly/CRA-vitals
reportWebVitals();

reactStateTree.ts

/**
 * Handle possible refs and memoizedState's. Due to their nature, refs
 * should not yet be handled. Certain states contain nested states (i.e. root
 * level state that has .next (state siblings).
 * @param object
 */
export const handleRefState = (object: any): any => {
    if (typeof object !== 'object' || !object) return object;

    if ('memoizedState' in object) {
        return {
            // Memoized state, usually the actual state value
            memoizedState: handleRefState(object.memoizedState),
            // The next state that is defined. This order is important as
            // React will throw an error concerning inconsistent state calls if
            // this is badly implemented.
            next: handleRefState(object.next)
        }
    }

    if ('current' in object) {
        // Refs are not yet handled
        return {
            $ref: true
        }
    }

    return object;
}

/**
 * Extract important attributes from the internal state tree and establish a
 * non-circular state tree that can be later enforced on the internal tree.
 * 
 * This avoids circular dependencies through a `WeakSet`, that can be optionally
 * passed to the `_seen` param.
 * @param object The root that we are traversing, usually has attributes such
 * as `memoizedState`, `child`, `sibling`, and `elementType` if it is a context
 * @param _seen The `WeakSet` used to label already-traversed nodes. A new,
 * empty `WeakSet` is used if not provided
 * @param marker The string UUID used on the `RootContext` to distinguish
 * itself to the state tree extractor, and to specially mark this context in
 * the exported state tree as the library.
 */
export const extractStateTree = (object: any, _seen?: WeakSet<any>, marker?: string | null): any => {
    if ([null, undefined].includes(object)) return;

    const seen = _seen ?? new WeakSet();
    if (seen.has(object)) {
        return;
    }
    seen.add(object);
    return {
        memoizedState: 'memoizedState' in object && ![null, undefined].includes(object.memoizedState) && !('element' in object.memoizedState) ? handleRefState(object.memoizedState) : undefined,
        child: 'child' in object && object.child ? extractStateTree(object.child, seen, marker) : undefined,
        // TODO: The use of `return` unfortunately causes DOMElements to get 
        //       extracted alongside actual state.
        // return: 'return' in object && object.return ? extractStateTree(object.return, seen, marker) : undefined,
        sibling: 'sibling' in object && object.sibling ? extractStateTree(object.sibling, seen, marker) : undefined,

        // The context's current value, unless the context is the root context,
        // in which case this would become infinitely recursive.
        elementType: 'elementType' in object && object.elementType?._context ?
            {
                _context: {
                    _currentValue: marker
                        ? object.elementType._context._currentValue?.[marker]
                            ? `$library-${marker}`
                            : object.elementType._context._currentValue
                        : object.elementType._context._currentValue
                }
            } : undefined
    }
}

export const enforceStateTree = (stateTree: any, root: any): any => {
    // Avoid attempting to set state as anything like 'undefined'
    // React considers `undefined` to be an undefined state and
    // as such throws un-ordered hook errors.
    if (!root) return;
    
    // If not an object, this is likely a state's actual value.
    if (typeof stateTree !== 'object' || !stateTree) return stateTree;

    // Set state internally
    if ('memoizedState' in stateTree) {
        root.memoizedState = enforceStateTree(stateTree.memoizedState, root.memoizedState)
    }

    // Set children's states internally
    if ('child' in stateTree) {
        root.child = enforceStateTree(stateTree.child, root.child)
    }

    // Set sibling's states internally
    if ('sibling' in stateTree) {
        root.sibling = enforceStateTree(stateTree.sibling, root.sibling)
    }
    
    // Ensure state tree order is kept intact, and that the .next is not null
    // If it were null, then this is the final defined state of the component.
    if ('next' in stateTree && stateTree.next !== null) {
        root.next = enforceStateTree(stateTree.next, root.next)
    }
    
    // TODO: Manage context states
    // if ('elementType' in stateTree) {
    //     root.elementType = enforceStateTree(stateTree.elementType, root.elementType)
    // }
    
    return root;
}

/**
 * Search the state tree for a value and return its path.
 * 
 * This is a generic tree traversal function. It just uses a `WeakSet` to avoid
 * React's circular references.
 * @param value The value to find
 * @returns A function to search the provided object
 */
export const searchStateTree = (value: string) => {
    const seen = new WeakSet();
    const search = (object: any, name: string = '$root'): any => {
        // If undefined or already traversed
        if (seen.has(object) || object === null || object === undefined) {
            return;
        }

        if (typeof object === 'object') {
            // Only add to WeakSet if an object, as WeakSets error on
            // non-objects
            seen.add(object);
            return Object.keys(object)
                .map(key => {
                    return search(object[key], `${name}.${key}`);
                })
                .filter(Boolean)
                .reduce((acc, curr) => [...acc, ...curr], []);
        }

        if (object === value) {
            return [`${name}`]
        };
    }

    return search;
}

rootContextProvider.tsx

import React, { useState } from "react";
import { v4 as uuidv4 } from 'uuid';

const RootContext = React.createContext<{
    root: any,
    updateRoot: () => void,
    id: string | null,
    [p: string]: any
}>({
    root: undefined,
    updateRoot: () => undefined,
    id: null
})

const RootContextProvider: React.FC<React.PropsWithChildren<{root: any}>> = ({ root, children }) => {
    const [s, ss] = useState(0);
    const id = uuidv4();
    return (
        <RootContext.Provider
            value={{
                root,
                updateRoot: () => ss(s + 1),
                [id]: true,
                id
            }}
        >
            {children}
        </RootContext.Provider>
    )
}

export default RootContextProvider;

export const useRoot = () => React.useContext(RootContext);