utils.ts

import type { Table, Table_Internal } from './types/Table'
import type { NoInfer, RowData, Updater } from './types/type-utils'
import type { TableFeatures } from './types/TableFeatures'
import type { TableState, TableState_All } from './types/TableState'

export function functionalUpdate<T>(updater: Updater<T>, input: T): T {
  return typeof updater === 'function'
    ? (updater as (i: T) => T)(input)
    : updater
}

export function shallowEqual<T>(a: T, b: T): boolean {
  if (Object.is(a, b)) return true
  if (
    typeof a !== 'object' ||
    typeof b !== 'object' ||
    a === null ||
    b === null
  ) {
    return false
  }
  const keysA = Object.keys(a)
  const keysB = Object.keys(b)
  if (keysA.length !== keysB.length) return false
  for (const key of keysA) {
    if (!Object.is((a as any)[key], (b as any)[key])) return false
  }
  return true
}

export function noop() {}

export function makeStateUpdater<
  TFeatures extends TableFeatures,
  K extends (string & {}) | keyof TableState_All | keyof TableState<TFeatures>,
>(key: K, instance: Table<TFeatures, any>) {
  return (updater: Updater<TableState<any>[K & keyof TableState<any>]>) => {
    instance.baseStore.setState(
      <TTableState extends TableState_All>(old: TTableState) => {
        return {
          ...old,
          [key]: functionalUpdate(updater, (old as any)[key]),
        }
      },
    )
  }
}

type AnyFunction = (...args: any) => any

export function isFunction<T extends AnyFunction>(d: any): d is T {
  return d instanceof Function
}

export function isNumberArray(d: any): d is Array<number> {
  return Array.isArray(d) && d.every((val) => typeof val === 'number')
}

export function flattenBy<TNode>(
  arr: Array<TNode>,
  getChildren: (item: TNode) => Array<TNode>,
) {
  const flat: Array<TNode> = []

  const recurse = (subArr: Array<TNode>) => {
    subArr.forEach((item) => {
      flat.push(item)
      const children = getChildren(item)
      if (children.length) {
        recurse(children)
      }
    })
  }

  recurse(arr)

  return flat
}

export const $internalMemoFnMeta = Symbol('memoFnMeta')
export type MemoFnMeta = { originalArgsLength?: number }

/**
 * @internal
 */
function setMemoFnMeta(fn: Function, meta: MemoFnMeta) {
  Object.defineProperty(fn, $internalMemoFnMeta, { value: meta })
}

/**
 * @internal
 */
export function getMemoFnMeta(fn: any): MemoFnMeta | null {
  return (typeof fn === 'function' && fn[$internalMemoFnMeta]) ?? null
}

interface MemoOptions<TDeps extends ReadonlyArray<any>, TDepArgs, TResult> {
  fn: (...args: NoInfer<TDeps>) => TResult
  memoDeps?: (depArgs?: TDepArgs) => [...TDeps] | undefined
  onAfterCompare?: (depsChanged: boolean) => void
  onAfterUpdate?: (result: TResult) => void
  onBeforeCompare?: () => void
  onBeforeUpdate?: () => void
}

export const memo = <TDeps extends ReadonlyArray<any>, TDepArgs, TResult>({
  fn,
  memoDeps,
  onAfterCompare,
  onAfterUpdate,
  onBeforeCompare,
  onBeforeUpdate,
}: MemoOptions<TDeps, TDepArgs, TResult>): ((
  depArgs?: TDepArgs,
) => TResult) => {
  let deps: Array<any> | undefined = []
  let result: TResult | undefined

  const memoizedFn = (depArgs?: TDepArgs): TResult => {
    onBeforeCompare?.()
    const newDeps = memoDeps?.(depArgs)
    const depsChanged =
      !newDeps ||
      newDeps.length !== deps?.length ||
      newDeps.some((dep: any, index: number) => deps?.[index] !== dep)
    onAfterCompare?.(depsChanged)

    if (!depsChanged) {
      return result!
    }

    deps = newDeps

    onBeforeUpdate?.()
    result = fn(...(newDeps ?? ([] as any)))
    onAfterUpdate?.(result)

    return result
  }

  setMemoFnMeta(memoizedFn, { originalArgsLength: fn.length })

  return memoizedFn
}

interface TableMemoOptions<
  TFeatures extends TableFeatures,
  TDeps extends ReadonlyArray<any>,
  TDepArgs,
  TResult,
> extends MemoOptions<TDeps, TDepArgs, TResult> {
  feature?: keyof TFeatures & string
  fnName: string
  objectId?: string
  onAfterUpdate?: () => void
  table: Table_Internal<TFeatures, any>
}

const pad = (str: number | string, num: number) => {
  str = String(str)
  while (str.length < num) {
    str = ' ' + str
  }
  return str
}

export function tableMemo<
  TFeatures extends TableFeatures,
  TDeps extends ReadonlyArray<any>,
  TDepArgs,
  TResult,
>({
  feature,
  fnName,
  objectId,
  onAfterUpdate,
  table,
  ...memoOptions
}: TableMemoOptions<TFeatures, TDeps, TDepArgs, TResult>) {
  let beforeCompareTime: number
  let afterCompareTime: number
  let startCalcTime: number
  let endCalcTime: number
  let runCount = 0
  let debug: boolean | undefined
  let debugCache: boolean | undefined

  if (process.env.NODE_ENV === 'development') {
    const { debugCache: _debugCache, debugAll } = table.options
    debugCache = _debugCache
    const { parentName } = getFunctionNameInfo(fnName, '.')

    const debugByParent =
      // @ts-expect-error
      table.options[
        `debug${(parentName != 'table' ? parentName + 's' : parentName).replace(
          parentName,
          parentName.charAt(0).toUpperCase() + parentName.slice(1),
        )}`
      ]
    const debugByFeature = feature
      ? // @ts-expect-error
        table.options[
          `debug${feature.charAt(0).toUpperCase() + feature.slice(1)}`
        ]
      : false

    debug = debugAll || debugByParent || debugByFeature
  }

  function logTime(time: number, depsChanged: boolean) {
    const runType =
      runCount === 0
        ? '(1st run)'
        : depsChanged
          ? '(rerun #' + runCount + ')'
          : '(cache)'
    runCount++

    console.groupCollapsed(
      `%c⏱ ${pad(`${time.toFixed(1)} ms`, 12)} %c${runType}%c ${fnName}%c ${objectId ? `(${fnName.split('.')[0]}Id: ${objectId})` : ''}`,
      `font-size: .6rem; font-weight: bold; ${
        depsChanged
          ? `color: hsl(
        ${Math.max(0, Math.min(120 - Math.log10(time) * 60, 120))}deg 100% 31%);`
          : ''
      } `,
      `color: ${runCount < 2 ? '#FF00FF' : '#FF1493'}`,
      'color: #666',
      'color: #87CEEB',
    )
    console.info({
      feature,
      state: table.store.state,
      deps: memoOptions.memoDeps?.toString(),
    })
    console.trace()
    console.groupEnd()
  }

  const debugOptions =
    process.env.NODE_ENV === 'development'
      ? {
          onBeforeCompare: () => {
            if (debugCache) {
              beforeCompareTime = performance.now()
            }
          },
          onAfterCompare: (depsChanged: boolean) => {
            if (debugCache) {
              afterCompareTime = performance.now()
              const compareTime =
                Math.round((afterCompareTime - beforeCompareTime) * 100) / 100
              if (!depsChanged) {
                logTime(compareTime, depsChanged)
              }
            }
          },
          onBeforeUpdate: () => {
            if (debug) {
              startCalcTime = performance.now()
            }
          },
          onAfterUpdate: () => {
            if (debug) {
              endCalcTime = performance.now()
              const executionTime =
                Math.round((endCalcTime - startCalcTime) * 100) / 100
              logTime(executionTime, true)
            }
            queueMicrotask(() => onAfterUpdate?.())
          },
        }
      : {
          onAfterUpdate: () => {
            queueMicrotask(() => onAfterUpdate?.())
          },
        }

  return memo({
    ...memoOptions,
    ...debugOptions,
  })
}

export interface API<TDeps extends ReadonlyArray<any>, TDepArgs> {
  fn: (...args: any) => any
  memoDeps?: (depArgs?: any) => [...any] | undefined
}

export type APIObject<TDeps extends ReadonlyArray<any>, TDepArgs> = Record<
  string,
  API<TDeps, TDepArgs>
>

/**
 * Assumes that a function name is in the format of `parentName_fnKey` and returns the `fnKey` and `fnName` in the format of `parentName.fnKey`.
 */
export function getFunctionNameInfo(
  staticFnName: string,
  splitBy: '_' | '.' = '_',
) {
  const [parentName, fnKey] = staticFnName.split(splitBy)
  const fnName = `${parentName}.${fnKey}`
  return { fnKey, fnName, parentName } as {
    fnKey: string
    fnName: string
    parentName: string
  }
}

/**
 * Assigns Table API methods directly to the table instance.
 * Unlike row/cell/column/header, the table is a singleton so methods are assigned directly.
 */
export function assignTableAPIs<
  TFeatures extends TableFeatures,
  TData extends RowData,
  TDeps extends ReadonlyArray<any>,
  TDepArgs,
>(
  feature: keyof TFeatures & string,
  table: Table_Internal<TFeatures, TData>,
  apis: APIObject<TDeps, NoInfer<TDepArgs>>,
): void {
  for (const [staticFnName, { fn, memoDeps }] of Object.entries(apis)) {
    const { fnKey, fnName } = getFunctionNameInfo(staticFnName)

    ;(table as Record<string, any>)[fnKey] = memoDeps
      ? tableMemo({
          memoDeps,
          fn,
          fnName,
          table,
          feature,
        })
      : fn
  }
}

export interface PrototypeAPI<TDeps extends ReadonlyArray<any>, TDepArgs> {
  fn: (self: any, ...args: any) => any
  memoDeps?: (self: any, depArgs?: any) => [...any] | undefined
}

export type PrototypeAPIObject<
  TDeps extends ReadonlyArray<any>,
  TDepArgs,
> = Record<string, PrototypeAPI<TDeps, TDepArgs>>

/**
 * Assigns API methods to a prototype object for memory-efficient method sharing.
 * All instances created with this prototype will share the same method references.
 *
 * For memoized methods, the memo state is lazily created and stored on each instance.
 * This provides the best of both worlds: shared method code + per-instance caching.
 */
export function assignPrototypeAPIs<
  TFeatures extends TableFeatures,
  TData extends RowData,
  TDeps extends ReadonlyArray<any>,
  TDepArgs,
>(
  feature: keyof TFeatures & string,
  prototype: Record<string, any>,
  table: Table_Internal<TFeatures, TData>,
  apis: PrototypeAPIObject<TDeps, NoInfer<TDepArgs>>,
): void {
  for (const [staticFnName, { fn, memoDeps }] of Object.entries(apis)) {
    const { fnKey, fnName } = getFunctionNameInfo(staticFnName)

    if (memoDeps) {
      // For memoized methods, create a function that lazily initializes
      // the memo on first access and stores it on the instance
      const memoKey = `_memo_${fnKey}`

      prototype[fnKey] = function (this: any, ...args: Array<any>) {
        // Lazily create memo on first access for this instance
        if (!this[memoKey]) {
          const self = this
          this[memoKey] = tableMemo({
            memoDeps: () => memoDeps(self),
            fn: (...deps) => fn(self, ...deps),
            fnName,
            objectId: self.id,
            table,
            feature,
          })
        }
        return this[memoKey](...args)
      }
    } else {
      // Non-memoized methods just call the static function with `this`
      prototype[fnKey] = function (this: any, ...args: Array<any>) {
        return fn(this, ...args)
      }
    }

    setMemoFnMeta(prototype[fnKey], { originalArgsLength: fn.length })
  }
}

/**
 * Looks to run the memoized function with the builder pattern on the object if it exists, otherwise fallback to the static method passed in.
 */
export function callMemoOrStaticFn<
  TObject extends Record<string, any>,
  TStaticFn extends AnyFunction,
>(
  obj: TObject,
  fnKey: string,
  staticFn: TStaticFn,
  ...args: Parameters<TStaticFn> extends [any, ...infer Rest] ? Rest : never
): ReturnType<TStaticFn> {
  return (
    (obj[fnKey] as Function | undefined)?.(...args) ?? staticFn(obj, ...args)
  )
}