From 2f2079e020ce7078ab59e78bc8f8cc54059d2a6c Mon Sep 17 00:00:00 2001
From: Mythie <me@lucasjamessmith.me>
Date: Sat, 21 Oct 2023 13:08:29 +1100
Subject: [PATCH] fix: optimise pdf viewer rerendering

---
 package-lock.json                             |  6 ++++++
 .../ee/server-only/limits/provider/client.tsx | 21 +++++++++++++++----
 packages/lib/package.json                     |  1 +
 packages/ui/primitives/pdf-viewer.tsx         | 11 +++++++---
 4 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 4f55b6418..0861bd0f3 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -16479,6 +16479,11 @@
         "url": "https://opencollective.com/unified"
       }
     },
+    "node_modules/remeda": {
+      "version": "1.27.1",
+      "resolved": "https://registry.npmjs.org/remeda/-/remeda-1.27.1.tgz",
+      "integrity": "sha512-va05uuDBz/E55O9wmpDdbVlwdWbHGJJy3oC0EAHSFn74MpWF3S81NVJDz/FW05bc/UDg769t1u6YPhBK/gmvLw=="
+    },
     "node_modules/repeat-string": {
       "version": "1.6.1",
       "resolved": "https://registry.npmjs.org/repeat-string/-/repeat-string-1.6.1.tgz",
@@ -19812,6 +19817,7 @@
         "next-auth": "4.22.3",
         "pdf-lib": "^1.17.1",
         "react": "18.2.0",
+        "remeda": "^1.27.1",
         "stripe": "^12.7.0",
         "ts-pattern": "^5.0.5",
         "zod": "^3.22.4"
diff --git a/packages/ee/server-only/limits/provider/client.tsx b/packages/ee/server-only/limits/provider/client.tsx
index 11b6bcae0..5f17678ee 100644
--- a/packages/ee/server-only/limits/provider/client.tsx
+++ b/packages/ee/server-only/limits/provider/client.tsx
@@ -1,6 +1,8 @@
 'use client';
 
-import { createContext, useContext, useEffect, useState } from 'react';
+import { createContext, useContext, useEffect, useRef, useState } from 'react';
+
+import { equals } from 'remeda';
 
 import { getLimits } from '../client';
 import { FREE_PLAN_LIMITS } from '../constants';
@@ -31,15 +33,26 @@ export const LimitsProvider = ({ initialValue, children }: LimitsProviderProps)
     remaining: FREE_PLAN_LIMITS,
   };
 
-  const [limits, setLimits] = useState(() => initialValue ?? defaultValue);
+  const $limits = useRef(initialValue ?? defaultValue);
+  const [limits, setLimits] = useState(() => $limits.current);
+
+  const refreshLimits = async () => {
+    const newLimits = await getLimits();
+
+    if (equals(newLimits, $limits.current)) {
+      return;
+    }
+
+    setLimits(newLimits);
+  };
 
   useEffect(() => {
-    void getLimits().then((limits) => setLimits(limits));
+    void refreshLimits();
   }, []);
 
   useEffect(() => {
     const onFocus = () => {
-      void getLimits().then((limits) => setLimits(limits));
+      void refreshLimits();
     };
 
     window.addEventListener('focus', onFocus);
diff --git a/packages/lib/package.json b/packages/lib/package.json
index 589145110..f200880d2 100644
--- a/packages/lib/package.json
+++ b/packages/lib/package.json
@@ -32,6 +32,7 @@
     "next-auth": "4.22.3",
     "pdf-lib": "^1.17.1",
     "react": "18.2.0",
+    "remeda": "^1.27.1",
     "stripe": "^12.7.0",
     "ts-pattern": "^5.0.5",
     "zod": "^3.22.4"
diff --git a/packages/ui/primitives/pdf-viewer.tsx b/packages/ui/primitives/pdf-viewer.tsx
index ed1733559..d057909d2 100644
--- a/packages/ui/primitives/pdf-viewer.tsx
+++ b/packages/ui/primitives/pdf-viewer.tsx
@@ -1,6 +1,6 @@
 'use client';
 
-import React, { useEffect, useRef, useState } from 'react';
+import React, { useEffect, useMemo, useRef, useState } from 'react';
 
 import { Loader } from 'lucide-react';
 import { PDFDocumentProxy } from 'pdfjs-dist';
@@ -66,6 +66,11 @@ export const PDFViewer = ({
   const [numPages, setNumPages] = useState(0);
   const [pdfError, setPdfError] = useState(false);
 
+  const memoizedData = useMemo(
+    () => ({ type: documentData.type, data: documentData.data }),
+    [documentData.data, documentData.type],
+  );
+
   const isLoading = isDocumentBytesLoading || !documentBytes;
 
   const onDocumentLoaded = (doc: LoadedPDFDocument) => {
@@ -134,7 +139,7 @@ export const PDFViewer = ({
       try {
         setIsDocumentBytesLoading(true);
 
-        const bytes = await getFile(documentData);
+        const bytes = await getFile(memoizedData);
 
         setDocumentBytes(bytes);
 
@@ -151,7 +156,7 @@ export const PDFViewer = ({
     };
 
     void fetchDocumentBytes();
-  }, [documentData, toast]);
+  }, [memoizedData, toast]);
 
   return (
     <div ref={$el} className={cn('overflow-hidden', className)} {...props}>