From d51b2b9283c6d96fe66a56e8a56db9b81b7e1bf7 Mon Sep 17 00:00:00 2001
From: Kevin Ingersoll <kingersoll@gmail.com>
Date: Fri, 10 May 2024 11:52:53 +0100
Subject: [PATCH] feat(store-sync): add lru cache to decodeEntity

---
 packages/common/src/LruMap.ts                | 22 ++++++++++++
 packages/common/src/index.ts                 |  1 +
 packages/store-sync/src/recs/decodeEntity.ts | 37 +++++++++++++++++---
 3 files changed, 56 insertions(+), 4 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..3b9eb1e380 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;
+}
+
+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,22 @@ 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 decoded = _decodeEntity(keySchema, entity);
+  cache.set(entity, decoded);
+  return decoded;
 }