From a3f8d77db2cc20a191c7c9f096ac0a281e6e781a Mon Sep 17 00:00:00 2001
From: Kevin Ingersoll <kingersoll@gmail.com>
Date: Fri, 10 May 2024 12:22:05 +0100
Subject: [PATCH] feat(store-sync): add lru cache to recs encode/decode entity
 (#2808)

---
 packages/common/src/LruMap.ts                | 22 ++++++++++
 packages/common/src/index.ts                 |  1 +
 packages/store-sync/src/recs/decodeEntity.ts | 43 ++++++++++++++++++--
 packages/store-sync/src/recs/encodeEntity.ts | 35 ++++++++++++++--
 4 files changed, 94 insertions(+), 7 deletions(-)
 create mode 100644 packages/common/src/LruMap.ts

diff --git a/packages/common/src/LruMap.ts b/packages/common/src/LruMap.ts
new file mode 100644
index 0000000000..7d0d3b5f30
--- /dev/null
+++ b/packages/common/src/LruMap.ts
@@ -0,0 +1,22 @@
+/**
+ * Map with a LRU (least recently used) policy.
+ *
+ * @link https://en.wikipedia.org/wiki/Cache_replacement_policies#LRU
+ * @link https://github.com/wevm/viem/blob/0fa08e113a890e6672fdc64fa7a2206a840611ab/src/utils/lru.ts
+ */
+export class LruMap<key, value> extends Map<key, value> {
+  maxSize: number;
+
+  constructor(size: number) {
+    super();
+    this.maxSize = size;
+  }
+
+  override set(key: key, value: value): this {
+    super.set(key, value);
+    if (this.maxSize && this.size > this.maxSize) {
+      this.delete(this.keys().next().value);
+    }
+    return this;
+  }
+}
diff --git a/packages/common/src/index.ts b/packages/common/src/index.ts
index 0da5bb7794..51248832f4 100644
--- a/packages/common/src/index.ts
+++ b/packages/common/src/index.ts
@@ -7,6 +7,7 @@ export * from "./getBurnerPrivateKey";
 export * from "./getNonceManager";
 export * from "./getNonceManagerId";
 export * from "./hexToResource";
+export * from "./LruMap";
 export * from "./readHex";
 export * from "./resourceToLabel";
 export * from "./resourceToHex";
diff --git a/packages/store-sync/src/recs/decodeEntity.ts b/packages/store-sync/src/recs/decodeEntity.ts
index 63c1ff8768..973d0b66f0 100644
--- a/packages/store-sync/src/recs/decodeEntity.ts
+++ b/packages/store-sync/src/recs/decodeEntity.ts
@@ -2,11 +2,23 @@ import { Entity } from "@latticexyz/recs";
 import { Hex, decodeAbiParameters } from "viem";
 import { entityToHexKeyTuple } from "./entityToHexKeyTuple";
 import { KeySchema, SchemaToPrimitives } from "@latticexyz/protocol-parser/internal";
+import { LruMap } from "@latticexyz/common";
 
-export function decodeEntity<TKeySchema extends KeySchema>(
-  keySchema: TKeySchema,
+const caches = new Map<KeySchema, LruMap<Entity, SchemaToPrimitives<KeySchema>>>();
+
+function getCache<keySchema extends KeySchema>(keySchema: keySchema): LruMap<Entity, SchemaToPrimitives<keySchema>> {
+  const cache = caches.get(keySchema);
+  if (cache != null) return cache as never;
+
+  const map = new LruMap<Entity, SchemaToPrimitives<keySchema>>(8096);
+  caches.set(keySchema, map);
+  return map;
+}
+
+export function _decodeEntity<keySchema extends KeySchema>(
+  keySchema: keySchema,
   entity: Entity,
-): SchemaToPrimitives<TKeySchema> {
+): SchemaToPrimitives<keySchema> {
   const hexKeyTuple = entityToHexKeyTuple(entity);
   if (hexKeyTuple.length !== Object.keys(keySchema).length) {
     throw new Error(
@@ -18,5 +30,28 @@ export function decodeEntity<TKeySchema extends KeySchema>(
       key,
       decodeAbiParameters([{ type }], hexKeyTuple[index] as Hex)[0],
     ]),
-  ) as SchemaToPrimitives<TKeySchema>;
+  ) as never;
+}
+
+// decoding can get expensive if we have thousands of entities, so we use a cache to ease this
+export function decodeEntity<keySchema extends KeySchema>(
+  keySchema: keySchema,
+  entity: Entity,
+): SchemaToPrimitives<keySchema> {
+  const cache = getCache(keySchema);
+
+  const cached = cache.get(entity);
+  if (cached != null) {
+    return cached as never;
+  }
+
+  const hexKeyTuple = entityToHexKeyTuple(entity);
+  if (hexKeyTuple.length !== Object.keys(keySchema).length) {
+    throw new Error(
+      `entity key tuple length ${hexKeyTuple.length} does not match key schema length ${Object.keys(keySchema).length}`,
+    );
+  }
+  const decoded = _decodeEntity(keySchema, entity);
+  cache.set(entity, decoded);
+  return decoded;
 }
diff --git a/packages/store-sync/src/recs/encodeEntity.ts b/packages/store-sync/src/recs/encodeEntity.ts
index a4cd2f2f2c..eaa8336ad8 100644
--- a/packages/store-sync/src/recs/encodeEntity.ts
+++ b/packages/store-sync/src/recs/encodeEntity.ts
@@ -2,10 +2,22 @@ import { Entity } from "@latticexyz/recs";
 import { encodeAbiParameters } from "viem";
 import { hexKeyTupleToEntity } from "./hexKeyTupleToEntity";
 import { KeySchema, SchemaToPrimitives } from "@latticexyz/protocol-parser/internal";
+import { LruMap } from "@latticexyz/common";
 
-export function encodeEntity<TKeySchema extends KeySchema>(
-  keySchema: TKeySchema,
-  key: SchemaToPrimitives<TKeySchema>,
+const caches = new Map<KeySchema, LruMap<SchemaToPrimitives<KeySchema>, Entity>>();
+
+function getCache<keySchema extends KeySchema>(keySchema: keySchema): LruMap<SchemaToPrimitives<keySchema>, Entity> {
+  const cache = caches.get(keySchema);
+  if (cache != null) return cache as never;
+
+  const map = new LruMap<SchemaToPrimitives<keySchema>, Entity>(8096);
+  caches.set(keySchema, map);
+  return map;
+}
+
+export function _encodeEntity<keySchema extends KeySchema>(
+  keySchema: keySchema,
+  key: SchemaToPrimitives<keySchema>,
 ): Entity {
   if (Object.keys(keySchema).length !== Object.keys(key).length) {
     throw new Error(
@@ -16,3 +28,20 @@ export function encodeEntity<TKeySchema extends KeySchema>(
     Object.entries(keySchema).map(([keyName, type]) => encodeAbiParameters([{ type }], [key[keyName]])),
   );
 }
+
+// encoding can get expensive if we have thousands of entities, so we use a cache to ease this
+export function encodeEntity<keySchema extends KeySchema>(
+  keySchema: keySchema,
+  key: SchemaToPrimitives<keySchema>,
+): Entity {
+  const cache = getCache(keySchema);
+
+  const cached = cache.get(key);
+  if (cached != null) {
+    return cached as never;
+  }
+
+  const encoded = _encodeEntity(keySchema, key);
+  cache.set(key, encoded);
+  return encoded;
+}