From c60f17d1013f9509c055cd207c9a318fddd808c6 Mon Sep 17 00:00:00 2001
From: Jacky Zhao <j.zhao2k19@gmail.com>
Date: Mon, 24 Jul 2023 00:04:01 -0700
Subject: [PATCH] fix watch-mode batching

---
 content/features/upcoming features.md         |  1 -
 quartz/build.ts                               | 73 +++++++++++++------
 quartz/ctx.ts                                 |  2 +
 quartz/plugins/emitters/componentResources.ts |  8 +-
 quartz/plugins/filters/draft.ts               |  2 +-
 quartz/plugins/filters/explicit.ts            |  2 +-
 quartz/plugins/index.ts                       |  9 +--
 quartz/plugins/transformers/links.ts          |  7 +-
 quartz/plugins/transformers/ofm.ts            |  2 +-
 quartz/plugins/types.ts                       | 13 ++--
 quartz/processors/emit.ts                     |  2 +-
 quartz/processors/filter.ts                   |  9 +--
 quartz/processors/parse.ts                    | 30 ++++----
 quartz/worker.ts                              |  8 +-
 14 files changed, 91 insertions(+), 77 deletions(-)

diff --git a/content/features/upcoming features.md b/content/features/upcoming features.md
index 671e6a1..d7acd2a 100644
--- a/content/features/upcoming features.md	
+++ b/content/features/upcoming features.md	
@@ -4,7 +4,6 @@ draft: true
 
 ## high priority
 
-- back button doesn't work sometimes
 - images in same folder are broken on shortest path mode
 - https://help.obsidian.md/Editing+and+formatting/Tags#Nested+tags nested tags?? and big tag listing
 - watch mode for config/source code
diff --git a/quartz/build.ts b/quartz/build.ts
index 26baa1b..b96c462 100644
--- a/quartz/build.ts
+++ b/quartz/build.ts
@@ -10,7 +10,7 @@ import { parseMarkdown } from "./processors/parse"
 import { filterContent } from "./processors/filter"
 import { emitContent } from "./processors/emit"
 import cfg from "../quartz.config"
-import { FilePath } from "./path"
+import { FilePath, slugifyFilePath } from "./path"
 import chokidar from "chokidar"
 import { ProcessedContent } from "./plugins/vfile"
 import WebSocket, { WebSocketServer } from "ws"
@@ -20,6 +20,7 @@ async function buildQuartz(argv: Argv, version: string) {
   const ctx: BuildCtx = {
     argv,
     cfg,
+    allSlugs: [],
   }
 
   console.log(chalk.bgGreen.black(`\n Quartz v${version} \n`))
@@ -51,6 +52,8 @@ async function buildQuartz(argv: Argv, version: string) {
   )
 
   const filePaths = fps.map((fp) => `${argv.directory}${path.sep}${fp}` as FilePath)
+  ctx.allSlugs = fps.map((fp) => slugifyFilePath(fp as FilePath))
+
   const parsedFiles = await parseMarkdown(ctx, filePaths)
   const filteredContent = filterContent(ctx, parsedFiles)
   await emitContent(ctx, filteredContent)
@@ -74,30 +77,54 @@ async function startServing(ctx: BuildCtx, initialContent: ProcessedContent[]) {
     contentMap.set(vfile.data.filePath!, content)
   }
 
-  async function rebuild(fp: string, action: "add" | "change" | "unlink") {
-    const perf = new PerfTimer()
+  let timeoutId: ReturnType<typeof setTimeout> | null = null
+  let toRebuild: Set<FilePath> = new Set()
+  let toRemove: Set<FilePath> = new Set()
+  async function rebuild(fp: string, action: "add" | "change" | "delete") {
     if (!ignored(fp)) {
-      console.log(chalk.yellow(`Detected change in ${fp}, rebuilding...`))
-      const fullPath = `${argv.directory}${path.sep}${fp}` as FilePath
-
-      try {
-        if (action === "add" || action === "change") {
-          const [parsedContent] = await parseMarkdown(ctx, [fullPath])
-          contentMap.set(fullPath, parsedContent)
-        } else if (action === "unlink") {
-          contentMap.delete(fullPath)
-        }
-
-        await rimraf(argv.output)
-        const parsedFiles = [...contentMap.values()]
-        const filteredContent = filterContent(ctx, parsedFiles)
-        await emitContent(ctx, filteredContent)
-        console.log(chalk.green(`Done rebuilding in ${perf.timeSince()}`))
-      } catch {
-        console.log(chalk.yellow(`Rebuild failed. Waiting on a change to fix the error...`))
+      const filePath = `${argv.directory}${path.sep}${fp}` as FilePath
+      if (action === "add" || action === "change") {
+        toRebuild.add(filePath)
+      } else if (action === "delete") {
+        toRemove.add(filePath)
       }
 
-      connections.forEach((conn) => conn.send("rebuild"))
+      if (timeoutId) {
+        clearTimeout(timeoutId)
+      }
+
+      timeoutId = setTimeout(async () => {
+        const perf = new PerfTimer()
+        console.log(chalk.yellow("Detected change, rebuilding..."))
+        try {
+          const filesToRebuild = [...toRebuild].filter((fp) => !toRemove.has(fp))
+
+          ctx.allSlugs = [...new Set([...contentMap.keys(), ...toRebuild])]
+            .filter((fp) => !toRemove.has(fp))
+            .map((fp) => slugifyFilePath(path.relative(argv.directory, fp) as FilePath))
+
+          const parsedContent = await parseMarkdown(ctx, filesToRebuild)
+          for (const content of parsedContent) {
+            const [_tree, vfile] = content
+            contentMap.set(vfile.data.filePath!, content)
+          }
+
+          for (const fp of toRemove) {
+            contentMap.delete(fp)
+          }
+
+          await rimraf(argv.output)
+          const parsedFiles = [...contentMap.values()]
+          const filteredContent = filterContent(ctx, parsedFiles)
+          await emitContent(ctx, filteredContent)
+          console.log(chalk.green(`Done rebuilding in ${perf.timeSince()}`))
+        } catch {
+          console.log(chalk.yellow(`Rebuild failed. Waiting on a change to fix the error...`))
+        }
+        connections.forEach((conn) => conn.send("rebuild"))
+        toRebuild.clear()
+        toRemove.clear()
+      }, 250)
     }
   }
 
@@ -110,7 +137,7 @@ async function startServing(ctx: BuildCtx, initialContent: ProcessedContent[]) {
   watcher
     .on("add", (fp) => rebuild(fp, "add"))
     .on("change", (fp) => rebuild(fp, "change"))
-    .on("unlink", (fp) => rebuild(fp, "unlink"))
+    .on("unlink", (fp) => rebuild(fp, "delete"))
 
   const server = http.createServer(async (req, res) => {
     await serveHandler(req, res, {
diff --git a/quartz/ctx.ts b/quartz/ctx.ts
index 355b4cb..8a7b803 100644
--- a/quartz/ctx.ts
+++ b/quartz/ctx.ts
@@ -1,4 +1,5 @@
 import { QuartzConfig } from "./cfg"
+import { ServerSlug } from "./path"
 
 export interface Argv {
   directory: string
@@ -11,4 +12,5 @@ export interface Argv {
 export interface BuildCtx {
   argv: Argv
   cfg: QuartzConfig
+  allSlugs: ServerSlug[]
 }
diff --git a/quartz/plugins/emitters/componentResources.ts b/quartz/plugins/emitters/componentResources.ts
index 99f9657..bc1d4ab 100644
--- a/quartz/plugins/emitters/componentResources.ts
+++ b/quartz/plugins/emitters/componentResources.ts
@@ -20,10 +20,10 @@ type ComponentResources = {
   afterDOMLoaded: string[]
 }
 
-function getComponentResources(plugins: PluginTypes): ComponentResources {
+function getComponentResources(ctx: BuildCtx): ComponentResources {
   const allComponents: Set<QuartzComponent> = new Set()
-  for (const emitter of plugins.emitters) {
-    const components = emitter.getQuartzComponents()
+  for (const emitter of ctx.cfg.plugins.emitters) {
+    const components = emitter.getQuartzComponents(ctx)
     for (const component of components) {
       allComponents.add(component)
     }
@@ -127,7 +127,7 @@ export const ComponentResources: QuartzEmitterPlugin = () => ({
   },
   async emit(ctx, _content, resources, emit): Promise<FilePath[]> {
     // component specific scripts and styles
-    const componentResources = getComponentResources(ctx.cfg.plugins)
+    const componentResources = getComponentResources(ctx)
     // important that this goes *after* component scripts
     // as the "nav" event gets triggered here and we should make sure
     // that everyone else had the chance to register a listener for it
diff --git a/quartz/plugins/filters/draft.ts b/quartz/plugins/filters/draft.ts
index e4a1f8f..65e2d6b 100644
--- a/quartz/plugins/filters/draft.ts
+++ b/quartz/plugins/filters/draft.ts
@@ -2,7 +2,7 @@ import { QuartzFilterPlugin } from "../types"
 
 export const RemoveDrafts: QuartzFilterPlugin<{}> = () => ({
   name: "RemoveDrafts",
-  shouldPublish([_tree, vfile]) {
+  shouldPublish(_ctx, [_tree, vfile]) {
     const draftFlag: boolean = vfile.data?.frontmatter?.draft ?? false
     return !draftFlag
   },
diff --git a/quartz/plugins/filters/explicit.ts b/quartz/plugins/filters/explicit.ts
index e0395a4..30f0b37 100644
--- a/quartz/plugins/filters/explicit.ts
+++ b/quartz/plugins/filters/explicit.ts
@@ -2,7 +2,7 @@ import { QuartzFilterPlugin } from "../types"
 
 export const ExplicitPublish: QuartzFilterPlugin = () => ({
   name: "ExplicitPublish",
-  shouldPublish([_tree, vfile]) {
+  shouldPublish(_ctx, [_tree, vfile]) {
     const publishFlag: boolean = vfile.data?.frontmatter?.publish ?? false
     return publishFlag
   },
diff --git a/quartz/plugins/index.ts b/quartz/plugins/index.ts
index a8208e3..23440fb 100644
--- a/quartz/plugins/index.ts
+++ b/quartz/plugins/index.ts
@@ -1,15 +1,15 @@
 import { StaticResources } from "../resources"
-import { PluginTypes } from "./types"
 import { FilePath, ServerSlug } from "../path"
+import { BuildCtx } from "../ctx"
 
-export function getStaticResourcesFromPlugins(plugins: PluginTypes) {
+export function getStaticResourcesFromPlugins(ctx: BuildCtx) {
   const staticResources: StaticResources = {
     css: [],
     js: [],
   }
 
-  for (const transformer of plugins.transformers) {
-    const res = transformer.externalResources ? transformer.externalResources() : {}
+  for (const transformer of ctx.cfg.plugins.transformers) {
+    const res = transformer.externalResources ? transformer.externalResources(ctx) : {}
     if (res?.js) {
       staticResources.js.push(...res.js)
     }
@@ -29,7 +29,6 @@ declare module "vfile" {
   // inserted in processors.ts
   interface DataMap {
     slug: ServerSlug
-    allSlugs: ServerSlug[]
     filePath: FilePath
   }
 }
diff --git a/quartz/plugins/transformers/links.ts b/quartz/plugins/transformers/links.ts
index e496171..7e8a278 100644
--- a/quartz/plugins/transformers/links.ts
+++ b/quartz/plugins/transformers/links.ts
@@ -29,7 +29,7 @@ export const CrawlLinks: QuartzTransformerPlugin<Partial<Options> | undefined> =
   const opts = { ...defaultOptions, ...userOpts }
   return {
     name: "LinkProcessing",
-    htmlPlugins() {
+    htmlPlugins(ctx) {
       return [
         () => {
           return (tree, file) => {
@@ -40,11 +40,8 @@ export const CrawlLinks: QuartzTransformerPlugin<Partial<Options> | undefined> =
               if (opts.markdownLinkResolution === "relative") {
                 return targetSlug as RelativeURL
               } else if (opts.markdownLinkResolution === "shortest") {
-                // https://forum.obsidian.md/t/settings-new-link-format-what-is-shortest-path-when-possible/6748/5
-                const allSlugs = file.data.allSlugs!
-
                 // if the file name is unique, then it's just the filename
-                const matchingFileNames = allSlugs.filter((slug) => {
+                const matchingFileNames = ctx.allSlugs.filter((slug) => {
                   const parts = slug.split(path.posix.sep)
                   const fileName = parts.at(-1)
                   return targetCanonical === fileName
diff --git a/quartz/plugins/transformers/ofm.ts b/quartz/plugins/transformers/ofm.ts
index 3f58d0f..36e79b0 100644
--- a/quartz/plugins/transformers/ofm.ts
+++ b/quartz/plugins/transformers/ofm.ts
@@ -119,7 +119,7 @@ export const ObsidianFlavoredMarkdown: QuartzTransformerPlugin<Partial<Options>
   const opts = { ...defaultOptions, ...userOpts }
   return {
     name: "ObsidianFlavoredMarkdown",
-    textTransform(src) {
+    textTransform(_ctx, src) {
       // pre-transform wikilinks (fix anchors to things that may contain illegal syntax e.g. codeblocks, latex)
       if (opts.wikilinks) {
         src = src.toString()
diff --git a/quartz/plugins/types.ts b/quartz/plugins/types.ts
index 4145e8f..2662aed 100644
--- a/quartz/plugins/types.ts
+++ b/quartz/plugins/types.ts
@@ -1,7 +1,6 @@
 import { PluggableList } from "unified"
 import { StaticResources } from "../resources"
 import { ProcessedContent } from "./vfile"
-import { GlobalConfiguration } from "../cfg"
 import { QuartzComponent } from "../components/types"
 import { FilePath, ServerSlug } from "../path"
 import { BuildCtx } from "../ctx"
@@ -18,10 +17,10 @@ export type QuartzTransformerPlugin<Options extends OptionType = undefined> = (
 ) => QuartzTransformerPluginInstance
 export type QuartzTransformerPluginInstance = {
   name: string
-  textTransform?: (src: string | Buffer) => string | Buffer
-  markdownPlugins?: () => PluggableList
-  htmlPlugins?: () => PluggableList
-  externalResources?: () => Partial<StaticResources>
+  textTransform?: (ctx: BuildCtx, src: string | Buffer) => string | Buffer
+  markdownPlugins?: (ctx: BuildCtx) => PluggableList
+  htmlPlugins?: (ctx: BuildCtx) => PluggableList
+  externalResources?: (ctx: BuildCtx) => Partial<StaticResources>
 }
 
 export type QuartzFilterPlugin<Options extends OptionType = undefined> = (
@@ -29,7 +28,7 @@ export type QuartzFilterPlugin<Options extends OptionType = undefined> = (
 ) => QuartzFilterPluginInstance
 export type QuartzFilterPluginInstance = {
   name: string
-  shouldPublish(content: ProcessedContent): boolean
+  shouldPublish(ctx: BuildCtx, content: ProcessedContent): boolean
 }
 
 export type QuartzEmitterPlugin<Options extends OptionType = undefined> = (
@@ -43,7 +42,7 @@ export type QuartzEmitterPluginInstance = {
     resources: StaticResources,
     emitCallback: EmitCallback,
   ): Promise<FilePath[]>
-  getQuartzComponents(): QuartzComponent[]
+  getQuartzComponents(ctx: BuildCtx): QuartzComponent[]
 }
 
 export interface EmitOptions {
diff --git a/quartz/processors/emit.ts b/quartz/processors/emit.ts
index 570f505..72d6085 100644
--- a/quartz/processors/emit.ts
+++ b/quartz/processors/emit.ts
@@ -24,7 +24,7 @@ export async function emitContent(ctx: BuildCtx, content: ProcessedContent[]) {
   }
 
   let emittedFiles = 0
-  const staticResources = getStaticResourcesFromPlugins(cfg.plugins)
+  const staticResources = getStaticResourcesFromPlugins(ctx)
   for (const emitter of cfg.plugins.emitters) {
     try {
       const emitted = await emitter.emit(ctx, content, staticResources, emit)
diff --git a/quartz/processors/filter.ts b/quartz/processors/filter.ts
index 12c5b48..dae6a3d 100644
--- a/quartz/processors/filter.ts
+++ b/quartz/processors/filter.ts
@@ -1,16 +1,13 @@
 import { BuildCtx } from "../ctx"
 import { PerfTimer } from "../perf"
-import { QuartzFilterPluginInstance } from "../plugins/types"
 import { ProcessedContent } from "../plugins/vfile"
 
-export function filterContent(
-  { cfg, argv }: BuildCtx,
-  content: ProcessedContent[],
-): ProcessedContent[] {
+export function filterContent(ctx: BuildCtx, content: ProcessedContent[]): ProcessedContent[] {
+  const { cfg, argv } = ctx
   const perf = new PerfTimer()
   const initialLength = content.length
   for (const plugin of cfg.plugins.filters) {
-    const updatedContent = content.filter(plugin.shouldPublish)
+    const updatedContent = content.filter((item) => plugin.shouldPublish(ctx, item))
 
     if (argv.verbose) {
       const diff = content.filter((x) => !updatedContent.includes(x))
diff --git a/quartz/processors/parse.ts b/quartz/processors/parse.ts
index aec2276..23af762 100644
--- a/quartz/processors/parse.ts
+++ b/quartz/processors/parse.ts
@@ -7,23 +7,24 @@ import { Root as HTMLRoot } from "hast"
 import { ProcessedContent } from "../plugins/vfile"
 import { PerfTimer } from "../perf"
 import { read } from "to-vfile"
-import { FilePath, QUARTZ, ServerSlug, slugifyFilePath } from "../path"
+import { FilePath, QUARTZ, slugifyFilePath } from "../path"
 import path from "path"
 import os from "os"
 import workerpool, { Promise as WorkerPromise } from "workerpool"
-import { QuartzTransformerPluginInstance } from "../plugins/types"
 import { QuartzLogger } from "../log"
 import { trace } from "../trace"
 import { BuildCtx } from "../ctx"
 
 export type QuartzProcessor = Processor<MDRoot, HTMLRoot, void>
-export function createProcessor(transformers: QuartzTransformerPluginInstance[]): QuartzProcessor {
+export function createProcessor(ctx: BuildCtx): QuartzProcessor {
+  const transformers = ctx.cfg.plugins.transformers
+
   // base Markdown -> MD AST
   let processor = unified().use(remarkParse)
 
   // MD AST -> MD AST transforms
   for (const plugin of transformers.filter((p) => p.markdownPlugins)) {
-    processor = processor.use(plugin.markdownPlugins!())
+    processor = processor.use(plugin.markdownPlugins!(ctx))
   }
 
   // MD AST -> HTML AST
@@ -31,7 +32,7 @@ export function createProcessor(transformers: QuartzTransformerPluginInstance[])
 
   // HTML AST -> HTML AST transforms
   for (const plugin of transformers.filter((p) => p.htmlPlugins)) {
-    processor = processor.use(plugin.htmlPlugins!())
+    processor = processor.use(plugin.htmlPlugins!(ctx))
   }
 
   return processor
@@ -73,7 +74,8 @@ async function transpileWorkerScript() {
   })
 }
 
-export function createFileParser({ argv, cfg }: BuildCtx, fps: FilePath[], allSlugs: ServerSlug[]) {
+export function createFileParser(ctx: BuildCtx, fps: FilePath[]) {
+  const { argv, cfg } = ctx
   return async (processor: QuartzProcessor) => {
     const res: ProcessedContent[] = []
     for (const fp of fps) {
@@ -85,12 +87,11 @@ export function createFileParser({ argv, cfg }: BuildCtx, fps: FilePath[], allSl
 
         // Text -> Text transforms
         for (const plugin of cfg.plugins.transformers.filter((p) => p.textTransform)) {
-          file.value = plugin.textTransform!(file.value)
+          file.value = plugin.textTransform!(ctx, file.value)
         }
 
         // base data properties that plugins may use
         file.data.slug = slugifyFilePath(path.relative(argv.directory, file.path) as FilePath)
-        file.data.allSlugs = allSlugs
         file.data.filePath = fp
 
         const ast = processor.parse(file)
@@ -111,24 +112,19 @@ export function createFileParser({ argv, cfg }: BuildCtx, fps: FilePath[], allSl
 }
 
 export async function parseMarkdown(ctx: BuildCtx, fps: FilePath[]): Promise<ProcessedContent[]> {
-  const { argv, cfg } = ctx
+  const { argv } = ctx
   const perf = new PerfTimer()
   const log = new QuartzLogger(argv.verbose)
 
   const CHUNK_SIZE = 128
   let concurrency = fps.length < CHUNK_SIZE ? 1 : os.availableParallelism()
 
-  // get all slugs ahead of time as each thread needs a copy
-  const allSlugs = fps.map((fp) =>
-    slugifyFilePath(path.relative(argv.directory, path.resolve(fp)) as FilePath),
-  )
-
   let res: ProcessedContent[] = []
   log.start(`Parsing input files using ${concurrency} threads`)
   if (concurrency === 1) {
     try {
-      const processor = createProcessor(cfg.plugins.transformers)
-      const parse = createFileParser(ctx, fps, allSlugs)
+      const processor = createProcessor(ctx)
+      const parse = createFileParser(ctx, fps)
       res = await parse(processor)
     } catch (error) {
       log.end()
@@ -144,7 +140,7 @@ export async function parseMarkdown(ctx: BuildCtx, fps: FilePath[]): Promise<Pro
 
     const childPromises: WorkerPromise<ProcessedContent[]>[] = []
     for (const chunk of chunks(fps, CHUNK_SIZE)) {
-      childPromises.push(pool.exec("parseFiles", [argv, chunk, allSlugs]))
+      childPromises.push(pool.exec("parseFiles", [argv, chunk, ctx.allSlugs]))
     }
 
     const results: ProcessedContent[][] = await WorkerPromise.all(childPromises)
diff --git a/quartz/worker.ts b/quartz/worker.ts
index eef4907..d97c483 100644
--- a/quartz/worker.ts
+++ b/quartz/worker.ts
@@ -3,16 +3,14 @@ import { Argv, BuildCtx } from "./ctx"
 import { FilePath, ServerSlug } from "./path"
 import { createFileParser, createProcessor } from "./processors/parse"
 
-const transformers = cfg.plugins.transformers
-const processor = createProcessor(transformers)
-
 // only called from worker thread
 export async function parseFiles(argv: Argv, fps: FilePath[], allSlugs: ServerSlug[]) {
   const ctx: BuildCtx = {
     cfg,
     argv,
+    allSlugs,
   }
-
-  const parse = createFileParser(ctx, fps, allSlugs)
+  const processor = createProcessor(ctx)
+  const parse = createFileParser(ctx, fps)
   return parse(processor)
 }