memoization.ts

type None = false
const none = false


type Equality<A> = (left: A, right: A) => boolean

function areSameReference<A>(left: A, right: A): boolean {
  return left === right
}

function arraysAreEqual<Elem>(left: Elem[], right: Elem[], elemEquality: Equality<Elem> = areSameReference): boolean {
  if (left.length !== right.length) {
    return false
  }
  return left.every((leftElem, index) => {
    const rightElem = right[index]
    return elemEquality(leftElem, rightElem)
  })
}

type MemoizeEquality<Params extends any[], Result> = Partial<
  Readonly<{
    paramsAreEqual: Equality<Params>
    resultsAreEqual: Equality<Result>
  }>
>

function memoize<Params extends any[], Result>(fn: (...params: Params) => Result, memoizeEquality: MemoizeEquality<Params, Result> = {}): (...params: Params) => Result {
  const paramsAreEqual: Equality<Params> = memoizeEquality.paramsAreEqual || arraysAreEqual
  const resultsAreEqual: Equality<Result> = memoizeEquality.resultsAreEqual || areSameReference
  type LastCall = Readonly<{
    params: Params
    result: Result
  }>
  let lastCall: LastCall | None = none
  return (...params: Params): Result => {
    if (lastCall !== none && paramsAreEqual(lastCall.params, params)) {
      return lastCall.result
    } else {
      const result = fn(...params)
      if (lastCall !== none && resultsAreEqual(lastCall.result, result)) {
        // If the results are the same, return the *exact* same result
        // to make sure that all other memoized functions work as efficiently
        // as possible:
        lastCall = { params, result: lastCall.result }
        return lastCall.result
      } else {
        lastCall = { params, result }
        return result
      }
    }
  }
}

type Func<P, R> = (p: P) => R 

function flatMap<P, R1, R2>(subject: Func<P, R1>, fn: (r: R1) => Func<P, R2>): Func<P, R2> {
  return (p: P): R2 => {
    const r1 = subject(p)
    return fn(r1)(p)
  }
}

function map<P, R1, R2>(subject: Func<P, R1>, fn: (r: R1) => R2): Func<P, R2> {
  return (p: P): R2 => {
    const r1 = subject(p)
    return fn(r1)
  }
}

function combine<P, Args extends any[], R>(combinator: (...args: Args) => R, ...selectors: { [K in keyof Args]: (p: P) => Args[K] }): (p: P) => R {
  return (p: P): R => {
    const args = selectors.map(s => s(p)) as Args
    return combinator(...args)
  }
}

type C = {
  a: number
  b: string
  c: boolean
}

const aSelector = memoize<[C], number>(c => c.a)
const bSelector = memoize<[C], string>(c => c.b)
const bLengthPlusA = flatMap(aSelector, memoize(a => {
  console.log('New a', a)
  return map(bSelector, memoize(b => {
    console.log('New b', a, c.b)
    return c.b + a
  }))
}))

let c: C = {
  a: 4,
  b: 'Henk',
  c: true
}
console.log('0', bLengthPlusA(c))
console.log('1', bLengthPlusA(c))
c = { ...c, c: false }
console.log('2', bLengthPlusA(c))
c = { ...c }
console.log('3', bLengthPlusA(c))
c = { ...c, a: 5 }
console.log('4', bLengthPlusA(c))
c = { ...c, b: 'Hark' }
console.log('5', bLengthPlusA(c))