1import { WorkerPoolContext, useWorkerPool } from "@pierre/diffs/react";
2import { WorkerPoolManager } from "@pierre/diffs/worker";
3import DiffsWorker from "@pierre/diffs/worker/worker.js?worker";
4import * as Schema from "effect/Schema";
5import {
6 useCallback,
7 useEffect,
8 useMemo,
9 useState,
10 useSyncExternalStore,
11 type ReactNode,
12} from "react";
13import { useTheme } from "../hooks/useTheme";
14import { resolveDiffThemeName, type DiffThemeName } from "../lib/diffRendering";
15import { PREFERRED_HIGHLIGHTER } from "../lib/syntaxHighlighting";
16
17export class DiffWorkerError extends Schema.TaggedError<DiffWorkerError>()("DiffWorkerError", {
18 operation: Schema.Literals(["create-worker", "get-render-options", "set-render-options"]),
19 themeName: Schema.Literals(["pierre-light", "pierre-dark"]),
20 cause: Schema.Defect(),
21}) {
22 override get message(): string {
23 return `Diff worker operation ${this.operation} failed for theme ${this.themeName}.`;
24 }
25}
26
27const DIFF_WORKER_IDLE_TTL_MS = 30_000;
28let sharedWorkerPool:
29 | {
30 readonly pool: WorkerPoolManager;
31 consumers: number;
32 idleTimer: ReturnType<typeof setTimeout> | undefined;
33 }
34 | undefined;
35
36/** Create workers after a viewer commits, then reuse them across short panel closures. */
37function acquireDiffWorkerPool(themeName: DiffThemeName, poolSize: number) {
38 const entry = (sharedWorkerPool ??= {
39 pool: new WorkerPoolManager(
40 {
41 workerFactory: () => {
42 try {
43 return new DiffsWorker();
44 } catch (cause) {
45 throw new DiffWorkerError({ operation: "create-worker", themeName, cause });
46 }
47 },
48 poolSize,
49 totalASTLRUCacheSize: 240,
50 },
51 {
52 theme: themeName,
53 preferredHighlighter: PREFERRED_HIGHLIGHTER,
54 tokenizeMaxLineLength: 1_000,
55 useTokenTransformer: true,
56 },
57 ),
58 consumers: 0,
59 idleTimer: undefined,
60 });
61 clearTimeout(entry.idleTimer);
62 entry.idleTimer = undefined;
63 entry.consumers += 1;
64 return entry;
65}
66
67function DiffWorkerThemeSync({ themeName }: { themeName: DiffThemeName }) {
68 const workerPool = useWorkerPool();
69
70 useEffect(() => {
71 if (!workerPool) {
72 return;
73 }
74
75 let operation: DiffWorkerError["operation"] = "get-render-options";
76 void (async () => {
77 try {
78 const current = workerPool.getDiffRenderOptions();
79 if (current.theme === themeName) {
80 return;
81 }
82
83 operation = "set-render-options";
84 await workerPool.setRenderOptions({
85 ...current,
86 theme: themeName,
87 });
88 } catch (cause) {
89 console.error(new DiffWorkerError({ operation, themeName, cause }));
90 }
91 })();
92 }, [themeName, workerPool]);
93
94 return null;
95}
96
97// Plain-text views do not queue a highlight task that could retry a blank first render.
98function DiffWorkerReady({ children }: { children?: ReactNode }) {
99 const workerPool = useWorkerPool();
100 const [readyPool, setReadyPool] = useState<WorkerPoolManager>();
101 const ready = workerPool
102 ? readyPool === workerPool || workerPool.isInitialized() || !workerPool.isWorkingPool()
103 : typeof window === "undefined";
104
105 useEffect(() => {
106 if (ready || !workerPool) return;
107
108 let mounted = true;
109 const finish = () => {
110 if (mounted) setReadyPool(workerPool);
111 };
112 // Failed pools use Pierre's existing main-thread highlighter.
113 void workerPool.initialize().then(finish, finish);
114 return () => {
115 mounted = false;
116 };
117 }, [ready, workerPool]);
118
119 return ready ? (
120 children
121 ) : (
122 <div
123 role="status"
124 className="flex min-h-0 flex-1 items-center justify-center p-4 text-xs text-muted-foreground"
125 >
126 Loading code...
127 </div>
128 );
129}
130
131export function DiffWorkerPoolProvider({ children }: { children?: ReactNode }) {
132 const { resolvedTheme } = useTheme();
133 const diffThemeName = resolveDiffThemeName(resolvedTheme);
134 const workerPoolSize = useMemo(() => {
135 const cores =
136 typeof navigator === "undefined" ? 4 : Math.max(1, navigator.hardwareConcurrency || 4);
137 return Math.max(2, Math.min(6, Math.floor(cores / 2)));
138 }, []);
139 const workerPool = useSyncExternalStore(
140 useCallback(
141 (onStoreChange) => {
142 if (typeof window === "undefined") return () => {};
143 const entry = acquireDiffWorkerPool(diffThemeName, workerPoolSize);
144 onStoreChange();
145 return () => {
146 entry.consumers -= 1;
147 if (entry.consumers !== 0) return;
148 entry.idleTimer = setTimeout(() => {
149 entry.idleTimer = undefined;
150 entry.pool.terminate();
151 if (sharedWorkerPool === entry) sharedWorkerPool = undefined;
152 }, DIFF_WORKER_IDLE_TTL_MS);
153 };
154 },
155 [diffThemeName, workerPoolSize],
156 ),
157 () => sharedWorkerPool?.pool,
158 () => undefined,
159 );
160
161 return (
162 <WorkerPoolContext value={workerPool}>
163 <DiffWorkerThemeSync themeName={diffThemeName} />
164 <DiffWorkerReady>{children}</DiffWorkerReady>
165 </WorkerPoolContext>
166 );
167}