feat(core): pdf preview (#8569)

Co-authored-by: forehalo <forehalo@gmail.com>

packages/frontend/core/src/modules/pdf/entities/pdf-page.ts

@@ -0,0 +1,39 @@
+import { DebugLogger } from '@affine/debug';
+import {
+  catchErrorInto,
+  effect,
+  Entity,
+  LiveData,
+  mapInto,
+} from '@toeverything/infra';
+import { map, switchMap } from 'rxjs';
+
+import type { RenderPageOpts } from '../renderer';
+import type { PDF } from './pdf';
+
+const logger = new DebugLogger('affine:pdf:page:render');
+
+export class PDFPage extends Entity<{ pdf: PDF; pageNum: number }> {
+  readonly pageNum: number = this.props.pageNum;
+  bitmap$ = new LiveData<ImageBitmap | null>(null);
+  error$ = new LiveData<any>(null);
+
+  render = effect(
+    switchMap((opts: Omit<RenderPageOpts, 'pageNum'>) =>
+      this.props.pdf.renderer.ob$('render', {
+        ...opts,
+        pageNum: this.pageNum,
+      })
+    ),
+    map(data => data.bitmap),
+    mapInto(this.bitmap$),
+    catchErrorInto(this.error$, error => {
+      logger.error('Failed to render page', error);
+    })
+  );
+
+  constructor() {
+    super();
+    this.disposables.push(() => this.render.unsubscribe);
+  }
+}

packages/frontend/core/src/modules/pdf/entities/pdf.ts

@@ -0,0 +1,74 @@
+import type { AttachmentBlockModel } from '@blocksuite/affine/blocks';
+import { Entity, LiveData, ObjectPool } from '@toeverything/infra';
+import { catchError, from, map, of, startWith, switchMap } from 'rxjs';
+
+import type { PDFMeta } from '../renderer';
+import { downloadBlobToBuffer, PDFRenderer } from '../renderer';
+import { PDFPage } from './pdf-page';
+
+export enum PDFStatus {
+  IDLE = 0,
+  Opening,
+  Opened,
+  Error,
+}
+
+export type PDFRendererState =
+  | {
+      status: PDFStatus.IDLE | PDFStatus.Opening;
+    }
+  | {
+      status: PDFStatus.Opened;
+      meta: PDFMeta;
+    }
+  | {
+      status: PDFStatus.Error;
+      error: Error;
+    };
+
+export class PDF extends Entity<AttachmentBlockModel> {
+  public readonly id: string = this.props.id;
+  readonly renderer = new PDFRenderer();
+  readonly pages = new ObjectPool<string, PDFPage>({
+    onDelete: page => page.dispose(),
+  });
+
+  readonly state$ = LiveData.from<PDFRendererState>(
+    // @ts-expect-error type alias
+    from(downloadBlobToBuffer(this.props)).pipe(
+      switchMap(buffer => {
+        return this.renderer.ob$('open', { data: buffer });
+      }),
+      map(meta => ({ status: PDFStatus.Opened, meta })),
+      // @ts-expect-error type alias
+      startWith({ status: PDFStatus.Opening }),
+      catchError((error: Error) => of({ status: PDFStatus.Error, error }))
+    ),
+    { status: PDFStatus.IDLE }
+  );
+
+  constructor() {
+    super();
+    this.renderer.listen();
+    this.disposables.push(() => this.pages.clear());
+  }
+
+  page(pageNum: number, size: string) {
+    const key = `${pageNum}:${size}`;
+    let rc = this.pages.get(key);
+
+    if (!rc) {
+      rc = this.pages.put(
+        key,
+        this.framework.createEntity(PDFPage, { pdf: this, pageNum })
+      );
+    }
+
+    return { page: rc.obj, release: rc.release };
+  }
+
+  override dispose() {
+    this.renderer.destroy();
+    super.dispose();
+  }
+}

packages/frontend/core/src/modules/pdf/index.ts

@@ -0,0 +1,19 @@
+import type { Framework } from '@toeverything/infra';
+import { WorkspaceScope } from '@toeverything/infra';
+
+import { PDF } from './entities/pdf';
+import { PDFPage } from './entities/pdf-page';
+import { PDFService } from './services/pdf';
+
+export function configurePDFModule(framework: Framework) {
+  framework
+    .scope(WorkspaceScope)
+    .service(PDFService)
+    .entity(PDF)
+    .entity(PDFPage);
+}
+
+export { PDF, type PDFRendererState, PDFStatus } from './entities/pdf';
+export { PDFPage } from './entities/pdf-page';
+export { PDFRenderer } from './renderer';
+export { PDFService } from './services/pdf';

packages/frontend/core/src/modules/pdf/renderer/index.ts

@@ -0,0 +1,3 @@
+export { PDFRenderer } from './renderer';
+export type { PDFMeta, RenderedPage, RenderPageOpts } from './types';
+export { downloadBlobToBuffer } from './utils';

packages/frontend/core/src/modules/pdf/renderer/ops.ts

@@ -0,0 +1,8 @@
+import type { OpSchema } from '@toeverything/infra/op';
+
+import type { PDFMeta, RenderedPage, RenderPageOpts } from './types';
+
+export interface ClientOps extends OpSchema {
+  open: [{ data: ArrayBuffer }, PDFMeta];
+  render: [RenderPageOpts, RenderedPage];
+}

packages/frontend/core/src/modules/pdf/renderer/renderer.ts

@@ -0,0 +1,28 @@
+import { OpClient } from '@toeverything/infra/op';
+
+import type { ClientOps } from './ops';
+
+export class PDFRenderer extends OpClient<ClientOps> {
+  private readonly worker: Worker;
+
+  constructor() {
+    const worker = new Worker(
+      /* webpackChunkName: "pdf.worker" */ new URL(
+        './worker.ts',
+        import.meta.url
+      )
+    );
+    super(worker);
+
+    this.worker = worker;
+  }
+
+  override destroy() {
+    super.destroy();
+    this.worker.terminate();
+  }
+
+  [Symbol.dispose]() {
+    this.destroy();
+  }
+}

packages/frontend/core/src/modules/pdf/renderer/types.ts

@@ -0,0 +1,16 @@
+export type PDFMeta = {
+  pageCount: number;
+  width: number;
+  height: number;
+};
+
+export type RenderPageOpts = {
+  pageNum: number;
+  width: number;
+  height: number;
+  scale?: number;
+};
+
+export type RenderedPage = RenderPageOpts & {
+  bitmap: ImageBitmap;
+};

packages/frontend/core/src/modules/pdf/renderer/utils.ts

@@ -0,0 +1,15 @@
+import type { AttachmentBlockModel } from '@blocksuite/affine/blocks';
+
+export async function downloadBlobToBuffer(model: AttachmentBlockModel) {
+  const sourceId = model.sourceId;
+  if (!sourceId) {
+    throw new Error('Attachment not found');
+  }
+
+  const blob = await model.doc.blobSync.get(sourceId);
+  if (!blob) {
+    throw new Error('Attachment not found');
+  }
+
+  return await blob.arrayBuffer();
+}

packages/frontend/core/src/modules/pdf/renderer/worker.ts

@@ -0,0 +1,140 @@
+import { OpConsumer, transfer } from '@toeverything/infra/op';
+import type { Document } from '@toeverything/pdf-viewer';
+import {
+  createPDFium,
+  PageRenderingflags,
+  Runtime,
+  Viewer,
+} from '@toeverything/pdf-viewer';
+import {
+  BehaviorSubject,
+  combineLatestWith,
+  filter,
+  from,
+  map,
+  Observable,
+  ReplaySubject,
+  share,
+  switchMap,
+} from 'rxjs';
+
+import type { ClientOps } from './ops';
+import type { PDFMeta, RenderPageOpts } from './types';
+
+class PDFRendererBackend extends OpConsumer<ClientOps> {
+  private readonly viewer$: Observable<Viewer> = from(
+    createPDFium().then(pdfium => {
+      return new Viewer(new Runtime(pdfium));
+    })
+  );
+
+  private readonly binary$ = new BehaviorSubject<Uint8Array | null>(null);
+
+  private readonly doc$ = this.binary$.pipe(
+    filter(Boolean),
+    combineLatestWith(this.viewer$),
+    switchMap(([buffer, viewer]) => {
+      return new Observable<Document | undefined>(observer => {
+        const doc = viewer.open(buffer);
+
+        if (!doc) {
+          observer.error(new Error('Document not opened'));
+          return;
+        }
+
+        observer.next(doc);
+
+        return () => {
+          doc.close();
+        };
+      });
+    }),
+    share({
+      connector: () => new ReplaySubject(1),
+    })
+  );
+
+  private readonly docInfo$: Observable<PDFMeta> = this.doc$.pipe(
+    map(doc => {
+      if (!doc) {
+        throw new Error('Document not opened');
+      }
+
+      const firstPage = doc.page(0);
+      if (!firstPage) {
+        throw new Error('Document has no pages');
+      }
+
+      return {
+        pageCount: doc.pageCount(),
+        width: firstPage.width(),
+        height: firstPage.height(),
+      };
+    })
+  );
+
+  open({ data }: { data: ArrayBuffer }) {
+    this.binary$.next(new Uint8Array(data));
+    return this.docInfo$;
+  }
+
+  render(opts: RenderPageOpts) {
+    return this.doc$.pipe(
+      combineLatestWith(this.viewer$),
+      switchMap(([doc, viewer]) => {
+        if (!doc) {
+          throw new Error('Document not opened');
+        }
+
+        return from(this.renderPage(viewer, doc, opts));
+      }),
+      map(bitmap => {
+        if (!bitmap) {
+          throw new Error('Failed to render page');
+        }
+
+        return transfer({ ...opts, bitmap }, [bitmap]);
+      })
+    );
+  }
+
+  async renderPage(viewer: Viewer, doc: Document, opts: RenderPageOpts) {
+    const page = doc.page(opts.pageNum);
+
+    if (!page) return;
+
+    const width = Math.ceil(opts.width * (opts.scale ?? 1));
+    const height = Math.ceil(opts.height * (opts.scale ?? 1));
+
+    const bitmap = viewer.createBitmap(width, height, 0);
+    bitmap.fill(0, 0, width, height);
+    page.render(
+      bitmap,
+      0,
+      0,
+      width,
+      height,
+      0,
+      PageRenderingflags.REVERSE_BYTE_ORDER | PageRenderingflags.ANNOT
+    );
+
+    const data = new Uint8ClampedArray(bitmap.toUint8Array());
+    const imageBitmap = await createImageBitmap(
+      new ImageData(data, width, height)
+    );
+
+    bitmap.close();
+    page.close();
+
+    return imageBitmap;
+  }
+
+  override listen(): void {
+    this.register('open', this.open.bind(this));
+    this.register('render', this.render.bind(this));
+    super.listen();
+  }
+}
+
+// @ts-expect-error how could we get correct postMessage signature for worker, exclude `window.postMessage`
+new PDFRendererBackend(self).listen();

packages/frontend/core/src/modules/pdf/services/pdf.ts

@@ -0,0 +1,31 @@
+import type { AttachmentBlockModel } from '@blocksuite/affine/blocks';
+import { ObjectPool, Service } from '@toeverything/infra';
+
+import { PDF } from '../entities/pdf';
+
+// One PDF document one worker.
+
+export class PDFService extends Service {
+  PDFs = new ObjectPool<string, PDF>({
+    onDelete: pdf => {
+      pdf.dispose();
+    },
+  });
+
+  constructor() {
+    super();
+    this.disposables.push(() => {
+      this.PDFs.clear();
+    });
+  }
+
+  get(model: AttachmentBlockModel) {
+    let rc = this.PDFs.get(model.id);
+
+    if (!rc) {
+      rc = this.PDFs.put(model.id, this.framework.createEntity(PDF, model));
+    }
+
+    return { pdf: rc.obj, release: rc.release };
+  }
+}

packages/frontend/core/src/modules/pdf/views/components.tsx

@@ -0,0 +1,185 @@
+import { Scrollable } from '@affine/component';
+import clsx from 'clsx';
+import { type CSSProperties, forwardRef, memo } from 'react';
+import type { VirtuosoProps } from 'react-virtuoso';
+
+import * as styles from './styles.css';
+
+export type PDFVirtuosoContext = {
+  width: number;
+  height: number;
+  pageClassName?: string;
+  onPageSelect?: (index: number) => void;
+};
+
+export type PDFVirtuosoProps = VirtuosoProps<unknown, PDFVirtuosoContext>;
+
+export const Scroller = forwardRef<HTMLDivElement, PDFVirtuosoProps>(
+  ({ context: _, ...props }, ref) => {
+    return (
+      <Scrollable.Root>
+        <Scrollable.Viewport ref={ref} {...props} />
+        <Scrollable.Scrollbar />
+      </Scrollable.Root>
+    );
+  }
+);
+
+Scroller.displayName = 'pdf-virtuoso-scroller';
+
+export const List = forwardRef<HTMLDivElement, PDFVirtuosoProps>(
+  ({ context: _, className, ...props }, ref) => {
+    return (
+      <div
+        className={clsx([styles.virtuosoList, className])}
+        ref={ref}
+        {...props}
+      />
+    );
+  }
+);
+
+List.displayName = 'pdf-virtuoso-list';
+
+export const ListWithSmallGap = forwardRef<HTMLDivElement, PDFVirtuosoProps>(
+  ({ context: _, className, ...props }, ref) => {
+    return (
+      <List className={clsx([className, 'small-gap'])} ref={ref} {...props} />
+    );
+  }
+);
+
+ListWithSmallGap.displayName = 'pdf-virtuoso-small-gap-list';
+
+export const Item = forwardRef<HTMLDivElement, PDFVirtuosoProps>(
+  ({ context: _, ...props }, ref) => {
+    return <div className={styles.virtuosoItem} ref={ref} {...props} />;
+  }
+);
+
+Item.displayName = 'pdf-virtuoso-item';
+
+export const ListPadding = () => (
+  <div style={{ width: '100%', height: '20px' }} />
+);
+
+export const LoadingSvg = memo(function LoadingSvg({
+  style,
+  className,
+}: {
+  style?: CSSProperties;
+  className?: string;
+}) {
+  return (
+    <svg
+      className={clsx([styles.pdfLoading, className])}
+      style={style}
+      width="16"
+      height="24"
+      viewBox="0 0 537 759"
+      fill="none"
+      xmlns="http://www.w3.org/2000/svg"
+    >
+      <rect width="537" height="759" fill="white" />
+      <rect
+        x="32"
+        y="82"
+        width="361"
+        height="30"
+        rx="4"
+        fill="black"
+        fillOpacity="0.07"
+      />
+      <rect
+        x="32"
+        y="142"
+        width="444"
+        height="30"
+        rx="4"
+        fill="black"
+        fillOpacity="0.07"
+      />
+      <rect
+        x="32"
+        y="202"
+        width="387"
+        height="30"
+        rx="4"
+        fill="black"
+        fillOpacity="0.07"
+      />
+      <rect
+        x="32"
+        y="262"
+        width="461"
+        height="30"
+        rx="4"
+        fill="black"
+        fillOpacity="0.07"
+      />
+      <rect
+        x="32"
+        y="322"
+        width="282"
+        height="30"
+        rx="4"
+        fill="black"
+        fillOpacity="0.07"
+      />
+      <rect
+        x="32"
+        y="382"
+        width="361"
+        height="30"
+        rx="4"
+        fill="black"
+        fillOpacity="0.07"
+      />
+      <rect
+        x="32"
+        y="442"
+        width="444"
+        height="30"
+        rx="4"
+        fill="black"
+        fillOpacity="0.07"
+      />
+      <rect
+        x="32"
+        y="502"
+        width="240"
+        height="30"
+        rx="4"
+        fill="black"
+        fillOpacity="0.07"
+      />
+      <rect
+        x="32"
+        y="562"
+        width="201"
+        height="30"
+        rx="4"
+        fill="black"
+        fillOpacity="0.07"
+      />
+      <rect
+        x="32"
+        y="622"
+        width="224"
+        height="30"
+        rx="4"
+        fill="black"
+        fillOpacity="0.07"
+      />
+      <rect
+        x="314"
+        y="502"
+        width="191"
+        height="166"
+        rx="4"
+        fill="black"
+        fillOpacity="0.07"
+      />
+    </svg>
+  );
+});

packages/frontend/core/src/modules/pdf/views/index.ts

@@ -0,0 +1,11 @@
+export {
+  Item,
+  List,
+  ListPadding,
+  ListWithSmallGap,
+  LoadingSvg,
+  type PDFVirtuosoContext,
+  type PDFVirtuosoProps,
+  Scroller,
+} from './components';
+export { PDFPageRenderer } from './page-renderer';

packages/frontend/core/src/modules/pdf/views/page-renderer.tsx

@@ -0,0 +1,84 @@
+import { useI18n } from '@affine/i18n';
+import { useLiveData } from '@toeverything/infra';
+import { useEffect, useRef, useState } from 'react';
+
+import type { PDF } from '../entities/pdf';
+import type { PDFPage } from '../entities/pdf-page';
+import { LoadingSvg } from './components';
+import * as styles from './styles.css';
+
+interface PDFPageProps {
+  pdf: PDF;
+  width: number;
+  height: number;
+  pageNum: number;
+  scale?: number;
+  className?: string;
+  onSelect?: (pageNum: number) => void;
+}
+
+export const PDFPageRenderer = ({
+  pdf,
+  width,
+  height,
+  pageNum,
+  className,
+  onSelect,
+  scale = window.devicePixelRatio,
+}: PDFPageProps) => {
+  const t = useI18n();
+  const [pdfPage, setPdfPage] = useState<PDFPage | null>(null);
+  const canvasRef = useRef<HTMLCanvasElement>(null);
+  const img = useLiveData(pdfPage?.bitmap$ ?? null);
+  const error = useLiveData(pdfPage?.error$ ?? null);
+  const style = { width, aspectRatio: `${width} / ${height}` };
+
+  useEffect(() => {
+    const { page, release } = pdf.page(pageNum, `${width}:${height}:${scale}`);
+    setPdfPage(page);
+
+    return release;
+  }, [pdf, width, height, pageNum, scale]);
+
+  useEffect(() => {
+    pdfPage?.render({ width, height, scale });
+
+    return pdfPage?.render.unsubscribe;
+  }, [pdfPage, width, height, scale]);
+
+  useEffect(() => {
+    const canvas = canvasRef.current;
+    if (!canvas) return;
+    if (!img) return;
+    const ctx = canvas.getContext('2d');
+    if (!ctx) return;
+
+    canvas.width = width * scale;
+    canvas.height = height * scale;
+    ctx.drawImage(img, 0, 0);
+  }, [img, width, height, scale]);
+
+  if (error) {
+    return (
+      <div className={className} style={style}>
+        <p className={styles.pdfPageError}>
+          {t['com.affine.pdf.page.render.error']()}
+        </p>
+      </div>
+    );
+  }
+
+  return (
+    <div
+      className={className}
+      style={style}
+      onClick={() => onSelect?.(pageNum)}
+    >
+      {img === null ? (
+        <LoadingSvg />
+      ) : (
+        <canvas className={styles.pdfPageCanvas} ref={canvasRef} />
+      )}
+    </div>
+  );
+};

packages/frontend/core/src/modules/pdf/views/styles.css.ts

@@ -0,0 +1,64 @@
+import { cssVarV2 } from '@toeverything/theme/v2';
+import { style } from '@vanilla-extract/css';
+
+export const virtuoso = style({
+  width: '100%',
+});
+
+export const virtuosoList = style({
+  display: 'flex',
+  flexDirection: 'column',
+  alignItems: 'center',
+  justifyContent: 'center',
+  gap: '20px',
+  selectors: {
+    '&.small-gap': {
+      gap: '12px',
+    },
+  },
+});
+
+export const virtuosoItem = style({
+  width: '100%',
+  display: 'flex',
+  alignItems: 'center',
+  justifyContent: 'center',
+});
+
+export const pdfPage = style({
+  overflow: 'hidden',
+  maxWidth: 'calc(100% - 40px)',
+  background: cssVarV2('layer/white'),
+  boxSizing: 'border-box',
+  borderWidth: '1px',
+  borderStyle: 'solid',
+  borderColor: cssVarV2('layer/insideBorder/border'),
+  boxShadow:
+    '0px 4px 20px 0px var(--transparent-black-200, rgba(0, 0, 0, 0.10))',
+});
+
+export const pdfPageError = style({
+  display: 'flex',
+  alignSelf: 'center',
+  justifyContent: 'center',
+  overflow: 'hidden',
+  textWrap: 'wrap',
+  width: '100%',
+  wordBreak: 'break-word',
+  fontSize: 14,
+  lineHeight: '22px',
+  fontWeight: 400,
+  color: cssVarV2('text/primary'),
+});
+
+export const pdfPageCanvas = style({
+  width: '100%',
+});
+
+export const pdfLoading = style({
+  display: 'flex',
+  alignSelf: 'center',
+  width: '100%',
+  height: '100%',
+  maxWidth: '537px',
+});