From f4c736283f3ed4fba87af2d37535f74a79ca0936 Mon Sep 17 00:00:00 2001
From: Art <4998038+Alorel@users.noreply.github.com>
Date: Wed, 13 Dec 2023 12:51:01 +0000
Subject: [PATCH] feat: Typed & arged cache access

---
 src/cache.ts |  80 ++++++++++++++++------
 src/core.ts  | 101 +++++++++++++++++++--------
 src/index.ts |  12 +++-
 src/test.ts  | 190 +++++++++++++++++++++++++++++++++++++++++++--------
 4 files changed, 303 insertions(+), 80 deletions(-)

diff --git a/src/cache.ts b/src/cache.ts
index 9c19184..e7e9a4b 100644
--- a/src/cache.ts
+++ b/src/cache.ts
@@ -2,8 +2,8 @@ import type {SerialiserFn} from './core';
 
 type Fn<T, A extends any[], R> = (this: T, ...args: A) => R;
 
-/** @see {MEMOISE_CACHE} */
-export interface Cache<K = any> {
+/** @see {MEMOISE_CACHE_TYPED} */
+export interface Cache<K = any, A extends any[] = any[]> {
 
   /** Clear the cache */
   clear(): void;
@@ -14,33 +14,45 @@ export interface Cache<K = any> {
    */
   delete(key: K): boolean;
 
+  /** Like {@link Cache#delete}, but the key gets computed */
+  deleteWithArgs(...args: A): boolean;
+
   /**
    * Check if a specific cache entry exists.
    * @param key See {@link Cache#delete delete()}
    */
   has(key: K): boolean;
+
+  /** Like {@link Cache#has}, but the key gets computed */
+  hasWithArgs(...args: A): boolean;
 }
 
 /** @internal */
-export class ArglessCtx<T, R> implements Cache {
-  public _f = true;
+export class ArglessCtx<T, R> implements Cache<any, []> {
+  private readonly _ctx: T;
+
+  private _firstCall = true;
 
-  public _r!: R;
+  private readonly _orig: Fn<T, [], R>;
 
-  public constructor(private readonly _o: Fn<T, [], R>) { // eslint-disable-line no-empty-function
+  private _value?: R;
+
+  public constructor(ctx: T, origFn: Fn<T, [], R>) {
+    this._orig = origFn;
+    this._ctx = ctx;
   }
 
-  public autoGet(ctx: T): R {
-    if (this._f) {
-      this._r = this._o.call(ctx);
-      this._f = false;
+  public autoGet(): R {
+    if (this._firstCall) {
+      this._value = this._orig.call(this._ctx);
+      this._firstCall = false;
     }
 
-    return this._r;
+    return this._value!;
   }
 
   public clear(): void {
-    this._f = true;
+    this._firstCall = true;
   }
 
   public delete(): boolean {
@@ -50,29 +62,55 @@ export class ArglessCtx<T, R> implements Cache {
     return ret;
   }
 
+  public deleteWithArgs(): boolean {
+    return this.delete();
+  }
+
   public has(): boolean {
-    return !this._f;
+    return !this._firstCall;
+  }
+
+  public hasWithArgs(): boolean {
+    return this.has();
   }
 }
 
 /** @internal */
-export class ArgedCtx<T, A extends any[], R, K> extends Map<K, R> implements Cache<K> {
-  public constructor(
-    private readonly _o: Fn<T, A, R>,
-    private readonly _s: SerialiserFn<T, A, K>
-  ) {
+export class ArgedCtx<T, A extends any[], R, K> extends Map<K, R> implements Cache<K, A> {
+  private readonly _ctx: T;
+
+  private readonly _orig: Fn<T, A, R>;
+
+  private readonly _serialiser: SerialiserFn<T, A, K>;
+
+  public constructor(ctx: T, origFn: Fn<T, A, R>, serialiser: SerialiserFn<T, A, K>) {
     super();
+    this._ctx = ctx;
+    this._orig = origFn;
+    this._serialiser = serialiser;
   }
 
-  public autoGet(ctx: T, args: A | IArguments): R {
-    const key = this._s.apply(ctx, args as A);
+  public autoGet(args: A | IArguments): R {
+    const key = this.keyFor(args);
     if (this.has(key)) {
       return this.get(key)!;
     }
 
-    const value = this._o.apply(ctx, args as A);
+    const value = this._orig.apply(this._ctx, args as A);
     this.set(key, value);
 
     return value;
   }
+
+  public deleteWithArgs(...args: A): boolean {
+    return this.delete(this.keyFor(args));
+  }
+
+  public hasWithArgs(...args: A): boolean {
+    return this.has(this.keyFor(args));
+  }
+
+  private keyFor(args: A | IArguments): K {
+    return this._serialiser.apply(this._ctx, args as A);
+  }
 }
diff --git a/src/core.ts b/src/core.ts
index 3f7f6d9..6448455 100644
--- a/src/core.ts
+++ b/src/core.ts
@@ -1,5 +1,11 @@
-import {ArgedCtx, ArglessCtx} from './cache';
 import type {Cache} from './cache';
+import {ArgedCtx, ArglessCtx} from './cache';
+
+/**
+ * @deprecated Use the typed version
+ * @see MEMOISE_CACHE_TYPED
+ */
+export const MEMOISE_CACHE: unique symbol = Symbol('Memoise cache');
 
 /**
  * Cache associated with this function/method if it's been processed with one of:
@@ -9,7 +15,9 @@ import type {Cache} from './cache';
  * - {@link memoiseFunction}
  * - {@link memoiseArglessFunction}
  */
-export const MEMOISE_CACHE: unique symbol = Symbol('Memoise cache');
+export const MEMOISE_CACHE_TYPED: unique symbol = Symbol('Memoise cache typed');
+
+export type MemoiseCacheGetFn = <T, A extends any[], R>(this: Fn<T, A, R>) => Cache<T, A> | undefined;
 
 /**
  * A serialisation function for computing cache keys. The returned key can be anything that's
@@ -20,7 +28,7 @@ export type SerialiserFn<T, A extends any[], K = any> = (this: T, ...args: A) =>
 
 type Fn<T, A extends any[], R> = (this: T, ...args: A) => R;
 export type Decorator<T, A extends any[], R> = (
-  target: any,
+  target: Fn<T, A, R>,
   ctx: ClassMethodDecoratorContext<T, Fn<T, A, R>>
 ) => undefined | Fn<T, A, R>;
 
@@ -35,7 +43,9 @@ export function identitySerialiser<T>(value: T): T {
 }
 
 interface Memoised<T, A extends any[], R> extends Fn<T, A, R> {
-  [MEMOISE_CACHE]: Cache;
+  [MEMOISE_CACHE]: Cache<T, A>;
+
+  [MEMOISE_CACHE_TYPED](): Cache<T, A>;
 }
 
 /** @internal */
@@ -51,8 +61,15 @@ function namedFn<F extends Function>(name: string, fn: F): F {
 }
 
 /** @internal */
-function applyRename<F extends Function>(origFn: F, label: string, newFn: F): F {
-  return namedFn(`${label}(${origFn.name})`, newFn);
+function applyRename<F extends Function>(origFnName: PropertyKey, label: string, newFn: F): F {
+  return namedFn(`${label}(${String(origFnName)})`, newFn);
+}
+
+function setCacheGetterFn(fn: Function, cache: () => Cache | undefined): void {
+  Object.defineProperty(fn, MEMOISE_CACHE_TYPED, {
+    configurable: true,
+    value: cache
+  });
 }
 
 function setCache(fn: Function, cache: Cache): void {
@@ -62,19 +79,25 @@ function setCache(fn: Function, cache: Cache): void {
   });
 }
 
+function memoiseFunction<T, A extends [any, ...any[]], R>(fn: Fn<T, A, R>): Memoised<T, A, R>;
+function memoiseFunction<T, A extends [any, ...any[]], R, K>(
+  fn: Fn<T, A, R>,
+  serialiser: SerialiserFn<T, A, K>
+): Memoised<T, A, R>;
+
 /**
  * Memoise the function's return value based on call arguments
  * @param fn The function to memoise
  * @param serialiser Serialiser to use for generating the cache key. Defaults to {@link defaultSerialiser}.
  */
-export function memoiseFunction<T, A extends [any, ...any[]], R>(
+function memoiseFunction<T, A extends [any, ...any[]], R, K>(
   fn: Fn<T, A, R>,
-  serialiser: SerialiserFn<T, A> = defaultSerialiser
+  serialiser: SerialiserFn<T, A, K> = defaultSerialiser as SerialiserFn<T, A>
 ): Memoised<T, A, R> {
-  const ctx = new ArgedCtx(fn, serialiser);
+  const ctx = new ArgedCtx(fn as T, fn, serialiser);
 
-  const memoisedFunction: Fn<T, A, R> = applyRename(fn, 'Memoised', function (this: T): R {
-    return ctx.autoGet(this, arguments);
+  const memoisedFunction: Fn<T, A, R> = applyRename(fn.name, 'Memoised', function (this: T): R {
+    return ctx.autoGet(arguments);
   });
 
   setCache(memoisedFunction, ctx);
@@ -82,16 +105,18 @@ export function memoiseFunction<T, A extends [any, ...any[]], R>(
   return memoisedFunction as Memoised<T, A, R>;
 }
 
+export {memoiseFunction};
+
 /**
  * Memoise the function's return value disregarding call arguments,
  * effectively turning it into a lazily-evaluated value
  * @param fn The function to memoise
  */
 export function memoiseArglessFunction<T, R>(fn: Fn<T, [], R>): Memoised<T, [], R> {
-  const ctx = new ArglessCtx(fn);
+  const ctx = new ArglessCtx(fn as T, fn);
 
-  const memoisedFunction: Fn<T, [], R> = applyRename(fn, 'MemoisedArgless', function (this: T): R {
-    return ctx.autoGet(this);
+  const memoisedFunction: Fn<T, [], R> = applyRename(fn.name, 'MemoisedArgless', function (this: T): R {
+    return ctx.autoGet();
   });
 
   setCache(memoisedFunction, ctx);
@@ -124,32 +149,48 @@ export function applyDecorator<T, A extends any[], R, K = any>(
     type Ctx = CtxArged | CtxArgless;
 
     const sName = String(name);
-    const ctxFactory: () => Ctx = hasArgs
-      ? (() => new ArgedCtx<T, A, R, K>(origFn, serialiser!))
-      : (() => new ArglessCtx<T, R>(origFn));
+    const ctxFactory: (ctx: T) => Ctx = hasArgs
+      ? (ctx => new ArgedCtx<T, A, R, K>(ctx, origFn, serialiser!))
+      : (ctx => new ArglessCtx<T, R>(ctx, origFn));
 
     if (isStatic) { // private/public static
-      const ctx = ctxFactory();
+      let ctx: Ctx;
+      const outFn = applyRename(name, 'Memoised', function (): R {
+        return ctx.autoGet(arguments);
+      });
+      addInitializer(applyRename(name, 'MemoiseInit', function (this: T) {
+        ctx = ctxFactory(this);
+      }));
 
-      const outFn = applyRename(origFn, 'Memoised', function (this: T): R {
-        return ctx.autoGet(this, arguments);
+      const getCache = (): Cache<K, A> => ctx;
+      Object.defineProperty(outFn, MEMOISE_CACHE, {
+        configurable: true,
+        get: getCache
       });
-      setCache(outFn, ctx);
+      setCacheGetterFn(outFn, getCache);
 
       return outFn;
     }
+
+    // Private/public instance
     const marker: unique symbol = Symbol(`Memoised value (${sName})`);
-      type Contextual = T & { [marker]: Ctx };
+    type Contextual = T & {[marker]: Ctx};
 
-      addInitializer(applyRename(origFn, `MemoiseInit(${sName})`, function (this: T) {
-        const ctx = ctxFactory();
-        Object.defineProperty(this, marker, {value: ctx});
-        setCache(get(this), ctx);
-      }));
+    addInitializer(applyRename(name, 'MemoiseInit', function (this: T) {
+      const ctx = ctxFactory(this);
+      Object.defineProperty(this, marker, {value: ctx});
 
-      return applyRename(origFn, 'Memoised', function (this: T): R {
-        return (this as Contextual)[marker].autoGet(this, arguments);
-      });
+      const instanceFn = get(this);
+      setCache(instanceFn, ctx);
+      setCacheGetterFn(instanceFn, () => ctx);
+    }));
 
+    return applyRename(name, 'Memoised', function (this: T): R {
+      return (this as Contextual)[marker].autoGet(arguments);
+    });
   };
 }
+
+setCacheGetterFn(Function.prototype, function (this: Fn<any, any[], any>) {
+  return this?.[MEMOISE_CACHE];
+});
diff --git a/src/index.ts b/src/index.ts
index 184836a..9fecaa6 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,10 +1,11 @@
 import type {Cache} from './cache';
-import type {Decorator, SerialiserFn} from './core';
+import type {Decorator, MemoiseCacheGetFn, SerialiserFn} from './core';
 import {
   applyDecorator,
   defaultSerialiser,
   identitySerialiser,
   MEMOISE_CACHE,
+  MEMOISE_CACHE_TYPED,
   memoiseArglessFunction,
   memoiseFunction
 } from './core';
@@ -32,6 +33,7 @@ export {
   MemoiseAll,
   MemoiseIdentity,
   MEMOISE_CACHE,
+  MEMOISE_CACHE_TYPED,
   defaultSerialiser,
   memoiseArglessFunction,
   memoiseFunction
@@ -42,6 +44,12 @@ declare global {
   interface Function {
 
     /** Always defined on decorated methods and explicitly memoised functions */
-    [MEMOISE_CACHE]?: Cache;
+    readonly [MEMOISE_CACHE]?: Cache;
+  }
+
+  interface CallableFunction {
+
+    /** Always returns non-undefined on decorated methods and explicitly memoised functions */
+    readonly [MEMOISE_CACHE_TYPED]: MemoiseCacheGetFn;
   }
 }
diff --git a/src/test.ts b/src/test.ts
index 80e20cd..7dac750 100644
--- a/src/test.ts
+++ b/src/test.ts
@@ -1,7 +1,15 @@
 import {expect} from 'chai';
-import type {Cache} from './index';
-import {MemoiseIdentity} from './index';
-import {Memoise, MEMOISE_CACHE, MemoiseAll, memoiseArglessFunction, memoiseFunction} from './index';
+import type {
+  Cache} from './index';
+import {
+  Memoise,
+  MEMOISE_CACHE,
+  MEMOISE_CACHE_TYPED,
+  MemoiseAll,
+  memoiseArglessFunction,
+  memoiseFunction,
+  MemoiseIdentity
+} from './index';
 
 /* eslint-disable @typescript-eslint/no-magic-numbers,class-methods-use-this,max-lines-per-function,no-new,max-lines */
 
@@ -20,10 +28,6 @@ class Base {
 
   public static rets: any[] = [];
 
-  public get publicCache(): Cache | undefined {
-    return this.memoedInstance[MEMOISE_CACHE];
-  }
-
   public static initArgChecks(beforeCB: () => void): void {
     this.initCommon();
     before(beforeCB);
@@ -183,6 +187,10 @@ class Base {
     return {};
   }
 
+  public get publicCache(): Cache | undefined {
+    return this.memoedInstance[MEMOISE_CACHE];
+  }
+
   public argedInstancePrivateDefaultS(a: number, b: number) {
     return this.#argedInstancePrivateDefaultS(a, b);
   }
@@ -439,26 +447,26 @@ describe('Class extensions', () => {
 
   class SubDeco extends Sup {
 
-    @MemoiseAll()
-    public override iPub() {
-      return [super.iPub(), Sup.callCount++] as const;
-    }
-
     @MemoiseAll()
     public static override sPub() {
       return [super.sPub(), Sup.callCount++] as const;
     }
-  }
-
-  class SubUndeco extends Sup {
 
+    @MemoiseAll()
     public override iPub() {
       return [super.iPub(), Sup.callCount++] as const;
     }
+  }
+
+  class SubUndeco extends Sup {
 
     public static override sPub() {
       return [super.sPub(), Sup.callCount++] as const;
     }
+
+    public override iPub() {
+      return [super.iPub(), Sup.callCount++] as const;
+    }
   }
 
   afterEach(Sup.reset);
@@ -509,6 +517,10 @@ describe('Class extensions', () => {
 describe('Cache', () => {
   describe('Access', () => {
     class Source {
+      public static get sPriv(): (this: typeof Source) => number {
+        return Source.#sPriv;
+      }
+
       public static reset(): void {
         Source.sPub[MEMOISE_CACHE]!.clear();
         Source.#sPriv[MEMOISE_CACHE]!.clear();
@@ -519,27 +531,23 @@ describe('Cache', () => {
         return 1;
       }
 
-      @MemoiseAll()
-      public iPub() {
-        return 2;
-      }
-
       @MemoiseAll()
       static #sPriv() {
         return 3;
       }
 
-      @MemoiseAll()
-      #iPriv() {
-        return 4;
-      }
-
       public get iPriv(): (this: Source) => number {
         return this.#iPriv;
       }
 
-      public static get sPriv(): (this: typeof Source) => number {
-        return Source.#sPriv;
+      @MemoiseAll()
+      public iPub() {
+        return 2;
+      }
+
+      @MemoiseAll()
+      #iPriv() {
+        return 4;
       }
     }
 
@@ -684,3 +692,131 @@ describe('identity', () => {
   expect(c1).to.eq(c2, 'c1 === c2');
   expect(c3).to.eq(c4, 'c3 === c4');
 });
+
+describe('has/delete by args', () => {
+  interface Spec<T, A extends any[]> {
+    args: () => A,
+
+    label: string;
+
+    call(inst: T, args: A): any;
+
+    getCache(inst: T): Cache<T, A> | undefined;
+
+    has(cache: Cache<T, A>, args: A): boolean;
+
+    init(): T;
+
+    rm(cache: Cache<T, A>, args: A): boolean;
+  }
+
+  class Src {
+    @MemoiseAll()
+    public static s0() {
+      return {};
+    }
+
+    @Memoise()
+    public static s2(a: number, b: number) {
+      return {a, b};
+    }
+
+    @MemoiseAll()
+    public i0() {
+      return {};
+    }
+
+    @Memoise()
+    public i2(a: string, b: string) {
+      return {a, b};
+    }
+  }
+
+  const fnArgless = memoiseArglessFunction(function fnArgless() {
+    return {};
+  });
+
+  const fnArged = memoiseFunction(function fnArged(a: number, b: number) {
+    return [a, b] as const;
+  });
+
+  const specs = [
+    {
+      args: () => [],
+      call: inst => inst(),
+      getCache: inst => inst[MEMOISE_CACHE_TYPED](),
+      has: cache => cache.hasWithArgs(),
+      init: () => fnArgless,
+      label: 'Argless fn',
+      rm: cache => cache.deleteWithArgs()
+    } satisfies Spec<typeof fnArgless, []>,
+
+    {
+      args: () => [0, 1],
+      call: (inst, args) => inst(...args),
+      getCache: inst => inst[MEMOISE_CACHE_TYPED](),
+      has: (cache, args) => cache.hasWithArgs(...args),
+      init: () => fnArged,
+      label: 'Arged fn',
+      rm: (cache, args) => cache.deleteWithArgs(...args)
+    } satisfies Spec<typeof fnArged, [number, number]>,
+
+    // /////////////////////////
+    {
+      args: () => [],
+      call: inst => inst.s0(),
+      getCache: inst => inst.s0[MEMOISE_CACHE_TYPED](),
+      has: cache => cache.hasWithArgs(),
+      init: () => Src,
+      label: 'Argless static decorated',
+      rm: cache => cache.deleteWithArgs()
+    } satisfies Spec<typeof Src, []>,
+
+    {
+      args: () => [0, 1],
+      call: (inst, args) => inst.s2(...args),
+      getCache: inst => inst.s2[MEMOISE_CACHE_TYPED](),
+      has: (cache, args) => cache.hasWithArgs(...args),
+      init: () => Src,
+      label: 'Arged static decorated',
+      rm: (cache, args) => cache.deleteWithArgs(...args)
+    } satisfies Spec<typeof Src, [number, number]>,
+
+    {
+      args: () => [],
+      call: inst => inst.i0(),
+      getCache: inst => inst.i0[MEMOISE_CACHE_TYPED](),
+      has: cache => cache.hasWithArgs(),
+      init: () => new Src(),
+      label: 'Argless instance decorated',
+      rm: cache => cache.deleteWithArgs()
+    } satisfies Spec<Src, []>,
+
+    {
+      args: () => ['a', 'b'],
+      call: (inst, args) => inst.i2(...args),
+      getCache: inst => inst.i2[MEMOISE_CACHE_TYPED](),
+      has: (cache, args) => cache.hasWithArgs(...args),
+      init: () => new Src(),
+      label: 'Arged instance decorated',
+      rm: (cache, args) => cache.deleteWithArgs(...args)
+    } satisfies Spec<Src, [string, string]>
+  ];
+
+  for (const spec of specs as Array<Spec<any, any[]>>) {
+    it(spec.label, () => {
+      const state = spec.init();
+      const cache = spec.getCache(state)!;
+
+      expect(cache.hasWithArgs(...spec.args())).to.eq(false, 'has [1]');
+      expect(cache.deleteWithArgs(...spec.args())).to.eq(false, 'delete [1]');
+
+      spec.call(state, spec.args());
+      expect(cache.hasWithArgs(...spec.args())).to.eq(true, 'has [2]');
+      expect(cache.deleteWithArgs(...spec.args())).to.eq(true, 'delete [2]');
+
+      expect(cache.hasWithArgs(...spec.args())).to.eq(false, 'has [3]');
+      expect(cache.deleteWithArgs(...spec.args())).to.eq(false, 'delete [3]');
+    });
+  }
+});