Refine DataSource id to use actual key type instead of a wide string type

Summary: Current implementation uses type `string` as a key for indexing items stored in datasource. However, users can provide any key as an index which means that the type of index item can be anything, not only string. This diff introduces a more refined types for the key. It adds another requirement to provide a key property to a generic which is used to infer the index type.

Reviewed By: mweststrate, aigoncharov

Differential Revision: D31895751

fbshipit-source-id: 19ba907bd6f35df87e3fa442db5fc5cec6af174d

desktop/flipper-plugin/src/data-source/DataSource.tsx

@@ -77,12 +77,14 @@ type OutputChange =
       newCount: number;
     };
 
-export type DataSourceOptions<T, K extends keyof T> = {
+export type DataSourceOptionKey<K extends PropertyKey> = {
   /**
    * If a key is set, the given field of the records is assumed to be unique,
    * and it's value can be used to perform lookups and upserts.
    */
   key?: K;
+};
+export type DataSourceOptions = {
   /**
    * The maximum amount of records that this DataSource will store.
    * If the limit is exceeded, the oldest records will automatically be dropped to make place for the new ones
@@ -90,16 +92,19 @@ export type DataSourceOptions<T, K extends keyof T> = {
   limit?: number;
 };
 
-export function createDataSource<T, KEY extends keyof T = any>(
+export function createDataSource<T, Key extends keyof T>(
   initialSet: readonly T[],
-  options?: DataSourceOptions<T, KEY>,
-): DataSource<T>;
-export function createDataSource<T>(initialSet?: readonly T[]): DataSource<T>;
-export function createDataSource<T, KEY extends keyof T>(
+  options: DataSourceOptions & DataSourceOptionKey<Key>,
+): DataSource<T, T[Key] extends string | number ? T[Key] : never>;
+export function createDataSource<T>(
+  initialSet?: readonly T[],
+  options?: DataSourceOptions,
+): DataSource<T, never>;
+export function createDataSource<T, Key extends keyof T>(
   initialSet: readonly T[] = [],
-  options?: DataSourceOptions<T, KEY>,
-): DataSource<T> {
-  const ds = new DataSource<T>(options?.key);
+  options?: DataSourceOptions & DataSourceOptionKey<Key>,
+): DataSource<T, Key> {
+  const ds = new DataSource<T, Key>(options?.key);
   if (options?.limit !== undefined) {
     ds.limit = options.limit;
   }
@@ -107,15 +112,15 @@ export function createDataSource<T, KEY extends keyof T>(
   return ds as any;
 }
 
-export class DataSource<T> {
+export class DataSource<T extends any, KeyType = never> {
   private nextId = 0;
   private _records: Entry<T>[] = [];
-  private _recordsById: Map<string, T> = new Map();
+  private _recordsById: Map<KeyType, T> = new Map();
   /**
    * @readonly
    */
   public keyAttribute: keyof T | undefined;
-  private idToIndex: Map<string, number> = new Map();
+  private idToIndex: Map<KeyType, number> = new Map();
 
   // if we shift the window, we increase shiftOffset to correct idToIndex results, rather than remapping all values
   private shiftOffset = 0;
@@ -131,11 +136,11 @@ export class DataSource<T> {
    *
    * Additional views can created through the fork method.
    */
-  public readonly view: DataSourceView<T>;
+  public readonly view: DataSourceView<T, KeyType>;
 
   constructor(keyAttribute: keyof T | undefined) {
     this.keyAttribute = keyAttribute;
-    this.view = new DataSourceView<T>(this);
+    this.view = new DataSourceView<T, KeyType>(this);
   }
 
   public get size() {
@@ -154,22 +159,22 @@ export class DataSource<T> {
     return unwrap(this._records[index]);
   }
 
-  public has(key: string) {
+  public has(key: KeyType) {
     this.assertKeySet();
     return this._recordsById.has(key);
   }
 
-  public getById(key: string): T | undefined {
+  public getById(key: KeyType): T | undefined {
     this.assertKeySet();
     return this._recordsById.get(key);
   }
 
-  public keys(): IterableIterator<string> {
+  public keys(): IterableIterator<KeyType> {
     this.assertKeySet();
     return this._recordsById.keys();
   }
 
-  public entries(): IterableIterator<[string, T]> {
+  public entries(): IterableIterator<[KeyType, T]> {
     this.assertKeySet();
     return this._recordsById.entries();
   }
@@ -198,7 +203,7 @@ export class DataSource<T> {
    * Returns the index of a specific key in the *records* set.
    * Returns -1 if the record wansn't found
    */
-  public getIndexOfKey(key: string): number {
+  public getIndexOfKey(key: KeyType): number {
     this.assertKeySet();
     const stored = this.idToIndex.get(key);
     return stored === undefined ? -1 : stored + this.shiftOffset;
@@ -328,7 +333,7 @@ export class DataSource<T> {
    *
    * Warning: this operation can be O(n) if a key is set
    */
-  public deleteByKey(keyValue: string): boolean {
+  public deleteByKey(keyValue: KeyType): boolean {
     this.assertKeySet();
     const index = this.getIndexOfKey(keyValue);
     if (index === -1) {
@@ -394,7 +399,7 @@ export class DataSource<T> {
    * Returns a fork of this dataSource, that shares the source data with this dataSource,
    * but has it's own FSRW pipeline, to allow multiple views on the same data
    */
-  public fork(): DataSourceView<T> {
+  public fork(): DataSourceView<T, KeyType> {
     throw new Error(
       'Not implemented. Please contact oncall if this feature is needed',
     );
@@ -408,7 +413,7 @@ export class DataSource<T> {
     }
   }
 
-  private getKey(value: T): string;
+  private getKey(value: T): KeyType;
   private getKey(value: any): any {
     this.assertKeySet();
     const key = value[this.keyAttribute!];
@@ -418,7 +423,7 @@ export class DataSource<T> {
     throw new Error(`Invalid key value: '${key}'`);
   }
 
-  private storeIndexOfKey(key: string, index: number) {
+  private storeIndexOfKey(key: KeyType, index: number) {
     // de-normalize the index, so that on  later look ups its corrected again
     this.idToIndex.set(key, index - this.shiftOffset);
   }
@@ -452,8 +457,8 @@ function unwrap<T>(entry: Entry<T>): T {
   return entry?.value;
 }
 
-class DataSourceView<T> {
-  public readonly datasource: DataSource<T>;
+class DataSourceView<T, KeyType> {
+  public readonly datasource: DataSource<T, KeyType>;
   private sortBy: undefined | ((a: T) => Primitive) = undefined;
   private reverse: boolean = false;
   private filter?: (value: T) => boolean = undefined;
@@ -474,7 +479,7 @@ class DataSourceView<T> {
    */
   private _output: Entry<T>[] = [];
 
-  constructor(datasource: DataSource<T>) {
+  constructor(datasource: DataSource<T, KeyType>) {
     this.datasource = datasource;
   }
 

desktop/flipper-plugin/src/data-source/DataSourceRendererStatic.tsx

@@ -16,7 +16,7 @@ type DataSourceProps<T extends object, C> = {
   /**
    * The data source to render
    */
-  dataSource: DataSource<T>;
+  dataSource: DataSource<T, T[keyof T]>;
   /**
    * additional context that will be passed verbatim to the itemRenderer, so that it can be easily memoized
    */
@@ -32,7 +32,9 @@ type DataSourceProps<T extends object, C> = {
   defaultRowHeight: number;
   onKeyDown?: React.KeyboardEventHandler<HTMLDivElement>;
   onUpdateAutoScroll?(autoScroll: boolean): void;
-  emptyRenderer?: null | ((dataSource: DataSource<T>) => React.ReactElement);
+  emptyRenderer?:
+    | null
+    | ((dataSource: DataSource<T, T[keyof T]>) => React.ReactElement);
 };
 
 /**

desktop/flipper-plugin/src/data-source/DataSourceRendererVirtual.tsx

@@ -39,7 +39,7 @@ type DataSourceProps<T extends object, C> = {
   /**
    * The data source to render
    */
-  dataSource: DataSource<T>;
+  dataSource: DataSource<T, T[keyof T]>;
   /**
    * Automatically scroll if the user is near the end?
    */
@@ -66,7 +66,9 @@ type DataSourceProps<T extends object, C> = {
     offset: number,
   ): void;
   onUpdateAutoScroll?(autoScroll: boolean): void;
-  emptyRenderer?: null | ((dataSource: DataSource<T>) => React.ReactElement);
+  emptyRenderer?:
+    | null
+    | ((dataSource: DataSource<T, T[keyof T]>) => React.ReactElement);
 };
 
 /**

desktop/flipper-plugin/src/data-source/__tests__/datasource-basics.node.tsx

@@ -34,7 +34,7 @@ function unwrap<T>(array: readonly {value: T}[]): readonly T[] {
   return array.map((entry) => entry.value);
 }
 
-function rawOutput<T>(ds: DataSource<T>): readonly T[] {
+function rawOutput<T>(ds: DataSource<T, T[keyof T]>): readonly T[] {
   // @ts-ignore
   const output = ds.view._output;
   return unwrap(output);
@@ -60,7 +60,7 @@ test('can create a datasource', () => {
 });
 
 test('can create a keyed datasource', () => {
-  const ds = createDataSource<Todo>([eatCookie], {key: 'id'});
+  const ds = createDataSource<Todo, 'id'>([eatCookie], {key: 'id'});
   expect(ds.records()).toEqual([eatCookie]);
 
   ds.append(drinkCoffee);
@@ -110,7 +110,7 @@ test('can create a keyed datasource', () => {
 });
 
 test('throws on invalid keys', () => {
-  const ds = createDataSource<Todo>([eatCookie], {key: 'id'});
+  const ds = createDataSource<Todo, 'id'>([eatCookie], {key: 'id'});
   expect(() => {
     ds.append({id: '', title: 'test'});
   }).toThrow(`Invalid key value: ''`);
@@ -120,7 +120,7 @@ test('throws on invalid keys', () => {
 });
 
 test('throws on update causing duplicate key', () => {
-  const ds = createDataSource<Todo>([eatCookie, submitBug], {key: 'id'});
+  const ds = createDataSource<Todo, 'id'>([eatCookie, submitBug], {key: 'id'});
   expect(() => {
     ds.update(0, {id: 'bug', title: 'oops'});
   }).toThrow(
@@ -129,7 +129,7 @@ test('throws on update causing duplicate key', () => {
 });
 
 test('removing invalid keys', () => {
-  const ds = createDataSource<Todo>([eatCookie], {key: 'id'});
+  const ds = createDataSource<Todo, 'id'>([eatCookie], {key: 'id'});
   expect(ds.deleteByKey('trash')).toBe(false);
   expect(() => {
     ds.delete(1);
@@ -263,7 +263,7 @@ test('filter + sort', () => {
 });
 
 test('filter + sort + index', () => {
-  const ds = createDataSource<Todo>([eatCookie, drinkCoffee, submitBug], {
+  const ds = createDataSource<Todo, 'id'>([eatCookie, drinkCoffee, submitBug], {
     key: 'id',
   });
 
@@ -315,7 +315,7 @@ test('filter + sort + index', () => {
 });
 
 test('filter', () => {
-  const ds = createDataSource<Todo>([eatCookie, drinkCoffee, submitBug], {
+  const ds = createDataSource<Todo, 'id'>([eatCookie, drinkCoffee, submitBug], {
     key: 'id',
   });
 
@@ -448,7 +448,7 @@ test('reverse with sorting', () => {
 });
 
 test('reset', () => {
-  const ds = createDataSource<Todo>([submitBug, drinkCoffee, eatCookie], {
+  const ds = createDataSource<Todo, 'id'>([submitBug, drinkCoffee, eatCookie], {
     key: 'id',
   });
   ds.view.setSortBy('title');
@@ -462,7 +462,7 @@ test('reset', () => {
 });
 
 test('clear', () => {
-  const ds = createDataSource<Todo>([submitBug, drinkCoffee, eatCookie], {
+  const ds = createDataSource<Todo, 'id'>([submitBug, drinkCoffee, eatCookie], {
     key: 'id',
   });
   ds.view.setSortBy('title');
@@ -484,10 +484,10 @@ test('clear', () => {
 
 function testEvents<T>(
   initial: T[],
-  op: (ds: DataSource<T>) => void,
+  op: (ds: DataSource<T, T[keyof T]>) => void,
   key?: keyof T,
 ): any[] {
-  const ds = createDataSource<T>(initial, {key});
+  const ds = createDataSource<T, keyof T>(initial, {key});
   const events: any[] = [];
   ds.view.setListener((e) => events.push(e));
   op(ds);

desktop/flipper-plugin/src/data-source/__tests__/datasource-perf.node.tsx

@@ -32,7 +32,7 @@ function generateTodos(amount: number): Todo[] {
 
 const defaultFilter = (t: Todo) => !t.done;
 
-type DataSourceish = DataSource<Todo> | FakeDataSource<Todo>;
+type DataSourceish = DataSource<Todo, Todo[keyof Todo]> | FakeDataSource<Todo>;
 
 // NOTE: this run in jest, which is not optimal for perf, but should give some idea
 // make sure to use the `yarn watch` script in desktop root, so that the garbage collector is exposed

desktop/flipper-plugin/src/data-source/index.tsx

@@ -7,7 +7,12 @@
  * @format
  */
 
-export {DataSource, createDataSource, DataSourceOptions} from './DataSource';
+export {
+  DataSource,
+  createDataSource,
+  DataSourceOptions,
+  DataSourceOptionKey,
+} from './DataSource';
 export {
   DataSourceRendererVirtual,
   DataSourceVirtualizer,

desktop/flipper-plugin/src/state/createDataSource.tsx

@@ -11,13 +11,11 @@ import {
   DataSource,
   createDataSource as baseCreateDataSource,
   DataSourceOptions as BaseDataSourceOptions,
+  DataSourceOptionKey as BaseDataSourceOptionKey,
 } from '../data-source/index';
 import {registerStorageAtom} from '../plugin/PluginBase';
 
-type CreateDataSourceOptions<T, K extends keyof T> = BaseDataSourceOptions<
-  T,
-  K
-> & {
+type DataSourceOptions = BaseDataSourceOptions & {
   /**
    * Should this state persist when exporting a plugin?
    * If set, the dataSource will be saved / loaded under the key provided
@@ -25,16 +23,21 @@ type CreateDataSourceOptions<T, K extends keyof T> = BaseDataSourceOptions<
   persist?: string;
 };
 
-export function createDataSource<T, KEY extends keyof T = any>(
+export function createDataSource<T, Key extends keyof T>(
   initialSet: readonly T[],
-  options: CreateDataSourceOptions<T, KEY>,
-): DataSource<T>;
-export function createDataSource<T>(initialSet?: readonly T[]): DataSource<T>;
-export function createDataSource<T, KEY extends keyof T>(
+  options: DataSourceOptions & BaseDataSourceOptionKey<Key>,
+): DataSource<T, T[Key] extends string | number ? T[Key] : never>;
+export function createDataSource<T>(
+  initialSet?: readonly T[],
+  options?: DataSourceOptions,
+): DataSource<T, never>;
+export function createDataSource<T, Key extends keyof T>(
   initialSet: readonly T[] = [],
-  options?: CreateDataSourceOptions<T, KEY>,
-): DataSource<T> {
-  const ds = baseCreateDataSource(initialSet, options);
+  options?: DataSourceOptions & BaseDataSourceOptionKey<Key>,
+): DataSource<T, T[Key] extends string | number ? T[Key] : never> {
+  const ds = options
+    ? baseCreateDataSource(initialSet, options)
+    : baseCreateDataSource(initialSet);
   registerStorageAtom(options?.persist, ds);
   return ds;
 }

desktop/flipper-plugin/src/ui/DataList.tsx

@@ -47,7 +47,7 @@ type DataListBaseProps<Item> = {
   /**
    * Items to display. Per item at least a title and unique id should be provided
    */
-  items: DataSource<Item> | readonly Item[];
+  items: DataSource<Item, Item[keyof Item]> | readonly Item[];
   /**
    * Custom render function. By default the component will render the `title` in bold and description (if any) below it
    */
@@ -75,10 +75,12 @@ export type DataListProps<Item> = DataListBaseProps<Item> &
   // Some table props are set by DataList instead, so override them
   Omit<DataTableProps<Item>, 'records' | 'dataSource' | 'columns' | 'onSelect'>;
 
-export const DataList: (<T>(props: DataListProps<T>) => React.ReactElement) & {
+export const DataList: (<T extends object>(
+  props: DataListProps<T>,
+) => React.ReactElement) & {
   Item: React.FC<DataListItemProps>;
 } = Object.assign(
-  function <T>({
+  function <T extends object>({
     onSelect: baseOnSelect,
     selection,
     className,
@@ -151,7 +153,7 @@ export const DataList: (<T>(props: DataListProps<T>) => React.ReactElement) & {
 
     return (
       <Layout.Container style={style} className={className} grow>
-        <DataTable<any>
+        <DataTable<T>
           {...tableProps}
           tableManagerRef={tableManagerRef}
           records={Array.isArray(items) ? items : undefined!}

desktop/flipper-plugin/src/ui/data-table/DataTable.tsx

@@ -69,7 +69,9 @@ interface DataTableBaseProps<T = any> {
   tableManagerRef?: RefObject<DataTableManager<T> | undefined>; // Actually we want a MutableRefObject, but that is not what React.createRef() returns, and we don't want to put the burden on the plugin dev to cast it...
   onCopyRows?(records: T[]): string;
   onContextMenu?: (selection: undefined | T) => React.ReactElement;
-  onRenderEmpty?: null | ((dataSource?: DataSource<T>) => React.ReactElement);
+  onRenderEmpty?:
+    | null
+    | ((dataSource?: DataSource<T, T[keyof T]>) => React.ReactElement);
 }
 
 export type ItemRenderer<T> = (
@@ -80,7 +82,7 @@ export type ItemRenderer<T> = (
 
 type DataTableInput<T = any> =
   | {
-      dataSource: DataSource<T>;
+      dataSource: DataSource<T, T[keyof T]>;
       records?: undefined;
       recordsKey?: undefined;
     }
@@ -558,7 +560,9 @@ DataTable.defaultProps = {
 } as Partial<DataTableProps<any>>;
 
 /* eslint-disable react-hooks/rules-of-hooks */
-function normalizeDataSourceInput<T>(props: DataTableInput<T>): DataSource<T> {
+function normalizeDataSourceInput<T>(
+  props: DataTableInput<T>,
+): DataSource<T, T[keyof T] | never> {
   if (props.dataSource) {
     return props.dataSource;
   }
@@ -578,7 +582,10 @@ function normalizeDataSourceInput<T>(props: DataTableInput<T>): DataSource<T> {
 }
 /* eslint-enable */
 
-function syncRecordsToDataSource<T>(ds: DataSource<T>, records: readonly T[]) {
+function syncRecordsToDataSource<T>(
+  ds: DataSource<T, T[keyof T] | never>,
+  records: readonly T[],
+) {
   const startTime = Date.now();
   ds.clear();
   // TODO: optimize in the case we're only dealing with appends or replacements

desktop/flipper-plugin/src/ui/data-table/DataTableManager.tsx

@@ -93,7 +93,7 @@ type DataManagerActions<T> =
   | Action<'setAutoScroll', {autoScroll: boolean}>;
 
 type DataManagerConfig<T> = {
-  dataSource: DataSource<T>;
+  dataSource: DataSource<T, T[keyof T]>;
   defaultColumns: DataTableColumn<T>[];
   scope: string;
   onSelect: undefined | ((item: T | undefined, items: T[]) => void);
@@ -279,11 +279,11 @@ export type DataTableManager<T> = {
   toggleColumnVisibility(column: keyof T): void;
   sortColumn(column: keyof T, direction?: SortDirection): void;
   setSearchValue(value: string): void;
-  dataSource: DataSource<T>;
+  dataSource: DataSource<T, T[keyof T]>;
 };
 
 export function createDataTableManager<T>(
-  dataSource: DataSource<T>,
+  dataSource: DataSource<T, T[keyof T]>,
   dispatch: DataTableDispatch<T>,
   stateRef: MutableRefObject<DataManagerState<T>>,
 ): DataTableManager<T> {
@@ -400,7 +400,7 @@ function addColumnFilter<T>(
 }
 
 export function getSelectedItem<T>(
-  dataSource: DataSource<T>,
+  dataSource: DataSource<T, T[keyof T]>,
   selection: Selection,
 ): T | undefined {
   return selection.current < 0
@@ -409,7 +409,7 @@ export function getSelectedItem<T>(
 }
 
 export function getSelectedItems<T>(
-  dataSource: DataSource<T>,
+  dataSource: DataSource<T, T[keyof T]>,
   selection: Selection,
 ): T[] {
   return [...selection.items]

desktop/flipper-plugin/src/ui/data-table/TableContextMenu.tsx

@@ -26,7 +26,7 @@ import {textContent} from '../../utils/textContent';
 const {Item, SubMenu} = Menu;
 
 export function tableContextMenuFactory<T>(
-  datasource: DataSource<T>,
+  datasource: DataSource<T, T[keyof T]>,
   dispatch: DataTableDispatch<T>,
   selection: Selection,
   columns: DataTableColumn<T>[],

desktop/flipper-plugin/src/utils/createTablePlugin.tsx

@@ -19,7 +19,7 @@ import {createDataSource} from '../state/createDataSource';
 
 type PluginResult<Raw, Row> = {
   plugin(client: PluginClient<Record<string, Raw | {}>>): {
-    rows: DataSource<Row>;
+    rows: DataSource<Row, Row[keyof Row]>;
   };
   Component(): React.ReactElement;
 };
@@ -75,9 +75,9 @@ export function createTablePlugin<
   function plugin(
     client: PluginClient<Record<Method, Raw> & Record<ResetMethod, {}>, {}>,
   ) {
-    const rows = createDataSource<Row>([], {
+    const rows = createDataSource<Row, keyof Row>([], {
       persist: 'rows',
-      key: props.key,
+      key: props.key as keyof Row | undefined,
     });
     const selection = createState<undefined | Row>(undefined);
     const isPaused = createState(false);

desktop/plugins/public/network/types.tsx

@@ -33,7 +33,7 @@ export interface Request {
   insights?: Insights;
 }
 
-export type Requests = DataSource<Request>;
+export type Requests = DataSource<Request, never>;
 
 export type SerializedRequest = Omit<
   Request,