|  | 
|  | 1 | +import useEvent from 'rc-util/lib/hooks/useEvent'; | 
|  | 2 | +import useLayoutEffect from 'rc-util/lib/hooks/useLayoutEffect'; | 
|  | 3 | +import isEqual from 'rc-util/lib/isEqual'; | 
|  | 4 | +import * as React from 'react'; | 
|  | 5 | +import { unstable_batchedUpdates } from 'react-dom'; | 
|  | 6 | + | 
|  | 7 | +export type Selector<ContextProps, SelectorValue = ContextProps> = ( | 
|  | 8 | +  value: ContextProps, | 
|  | 9 | +) => SelectorValue; | 
|  | 10 | + | 
|  | 11 | +export type Trigger<ContextProps> = (value: ContextProps) => void; | 
|  | 12 | + | 
|  | 13 | +export type Listeners<ContextProps> = Set<Trigger<ContextProps>>; | 
|  | 14 | + | 
|  | 15 | +export interface Context<ContextProps> { | 
|  | 16 | +  getValue: () => ContextProps; | 
|  | 17 | +  listeners: Listeners<ContextProps>; | 
|  | 18 | +} | 
|  | 19 | + | 
|  | 20 | +export interface ContextSelectorProviderProps<T> { | 
|  | 21 | +  value: T; | 
|  | 22 | +  children?: React.ReactNode; | 
|  | 23 | +} | 
|  | 24 | + | 
|  | 25 | +export interface SelectorContext<ContextProps> { | 
|  | 26 | +  Context: React.Context<Context<ContextProps>>; | 
|  | 27 | +  Provider: React.ComponentType<ContextSelectorProviderProps<ContextProps>>; | 
|  | 28 | +} | 
|  | 29 | + | 
|  | 30 | +export function createContext<ContextProps>( | 
|  | 31 | +  defaultContext?: ContextProps, | 
|  | 32 | +): SelectorContext<ContextProps> { | 
|  | 33 | +  const Context = React.createContext<Context<ContextProps>>(defaultContext as any); | 
|  | 34 | + | 
|  | 35 | +  const Provider = ({ value, children }: ContextSelectorProviderProps<ContextProps>) => { | 
|  | 36 | +    const valueRef = React.useRef(value); | 
|  | 37 | +    valueRef.current = value; | 
|  | 38 | + | 
|  | 39 | +    const [context] = React.useState<Context<ContextProps>>(() => ({ | 
|  | 40 | +      getValue: () => valueRef.current, | 
|  | 41 | +      listeners: new Set(), | 
|  | 42 | +    })); | 
|  | 43 | + | 
|  | 44 | +    useLayoutEffect(() => { | 
|  | 45 | +      unstable_batchedUpdates(() => { | 
|  | 46 | +        context.listeners.forEach(listener => { | 
|  | 47 | +          listener(value); | 
|  | 48 | +        }); | 
|  | 49 | +      }); | 
|  | 50 | +    }, [value]); | 
|  | 51 | + | 
|  | 52 | +    return <Context.Provider value={context}>{children}</Context.Provider>; | 
|  | 53 | +  }; | 
|  | 54 | + | 
|  | 55 | +  return { Context, Provider }; | 
|  | 56 | +} | 
|  | 57 | + | 
|  | 58 | +/** e.g. useSelect(userContext, user => user.name) => user.name */ | 
|  | 59 | +export function useContextSelector<ContextProps, SelectorValue>( | 
|  | 60 | +  holder: SelectorContext<ContextProps>, | 
|  | 61 | +  selector: Selector<ContextProps, SelectorValue>, | 
|  | 62 | +): SelectorValue; | 
|  | 63 | + | 
|  | 64 | +/** e.g. useSelect(userContext, ['name', 'age']) => user { name, age } */ | 
|  | 65 | +export function useContextSelector<ContextProps, SelectorValue extends Partial<ContextProps>>( | 
|  | 66 | +  holder: SelectorContext<ContextProps>, | 
|  | 67 | +  selector: (keyof ContextProps)[], | 
|  | 68 | +): SelectorValue; | 
|  | 69 | + | 
|  | 70 | +/** e.g. useSelect(userContext, 'name') => user.name */ | 
|  | 71 | +export function useContextSelector<ContextProps, PropName extends keyof ContextProps>( | 
|  | 72 | +  holder: SelectorContext<ContextProps>, | 
|  | 73 | +  selector: PropName, | 
|  | 74 | +): ContextProps[PropName]; | 
|  | 75 | + | 
|  | 76 | +export function useContextSelector<ContextProps, SelectorValue>( | 
|  | 77 | +  holder: SelectorContext<ContextProps>, | 
|  | 78 | +  selector: Selector<ContextProps, any> | (keyof ContextProps)[] | keyof ContextProps, | 
|  | 79 | +) { | 
|  | 80 | +  const eventSelector = useEvent<Selector<ContextProps, SelectorValue>>( | 
|  | 81 | +    typeof selector === 'function' | 
|  | 82 | +      ? selector | 
|  | 83 | +      : ctx => { | 
|  | 84 | +          if (!Array.isArray(selector)) { | 
|  | 85 | +            return ctx[selector]; | 
|  | 86 | +          } | 
|  | 87 | + | 
|  | 88 | +          const obj = {} as SelectorValue; | 
|  | 89 | +          selector.forEach(key => { | 
|  | 90 | +            (obj as any)[key] = ctx[key]; | 
|  | 91 | +          }); | 
|  | 92 | +          return obj; | 
|  | 93 | +        }, | 
|  | 94 | +  ); | 
|  | 95 | +  const context = React.useContext(holder?.Context); | 
|  | 96 | +  const { listeners, getValue } = context || {}; | 
|  | 97 | + | 
|  | 98 | +  const valueRef = React.useRef<SelectorValue>(); | 
|  | 99 | +  valueRef.current = eventSelector(context ? getValue() : null); | 
|  | 100 | +  const [, forceUpdate] = React.useState({}); | 
|  | 101 | + | 
|  | 102 | +  useLayoutEffect(() => { | 
|  | 103 | +    if (!context) { | 
|  | 104 | +      return; | 
|  | 105 | +    } | 
|  | 106 | + | 
|  | 107 | +    function trigger(nextValue: ContextProps) { | 
|  | 108 | +      const nextSelectorValue = eventSelector(nextValue); | 
|  | 109 | +      if (!isEqual(valueRef.current, nextSelectorValue, true)) { | 
|  | 110 | +        forceUpdate({}); | 
|  | 111 | +      } | 
|  | 112 | +    } | 
|  | 113 | + | 
|  | 114 | +    listeners.add(trigger); | 
|  | 115 | + | 
|  | 116 | +    return () => { | 
|  | 117 | +      listeners.delete(trigger); | 
|  | 118 | +    }; | 
|  | 119 | +  }, [context]); | 
|  | 120 | + | 
|  | 121 | +  return valueRef.current; | 
|  | 122 | +} | 
0 commit comments