From 8436472c7874cb16caf9432660b395ca9ba68f9d Mon Sep 17 00:00:00 2001
From: Divyansh Singh <40380293+brc-dd@users.noreply.github.com>
Date: Sat, 12 Oct 2024 20:51:31 +0530
Subject: [PATCH] feat(experimental): support passing function for rewrites
 (#4274)

---
 docs/en/guide/routing.md           | 44 ++++++++++++++++++++----------
 src/node/plugins/rewritesPlugin.ts | 39 +++++++++++++++++---------
 src/node/siteConfig.ts             |  2 +-
 3 files changed, 56 insertions(+), 29 deletions(-)

diff --git a/docs/en/guide/routing.md b/docs/en/guide/routing.md
index 945e63c0fceb..55b3cd4cb26e 100644
--- a/docs/en/guide/routing.md
+++ b/docs/en/guide/routing.md
@@ -156,22 +156,24 @@ You can customize the mapping between the source directory structure and the gen
 
 ```
 .
-├─ packages
-│  ├─ pkg-a
-│  │  └─ src
-│  │      ├─ pkg-a-code.ts
-│  │      └─ pkg-a-docs.md
-│  └─ pkg-b
-│     └─ src
-│         ├─ pkg-b-code.ts
-│         └─ pkg-b-docs.md
+└─ packages
+   ├─ pkg-a
+   │  └─ src
+   │     ├─ foo.md
+   │     └─ index.md
+   └─ pkg-b
+      └─ src
+         ├─ bar.md
+         └─ index.md
 ```
 
 And you want the VitePress pages to be generated like this:
 
 ```
-packages/pkg-a/src/pkg-a-docs.md  -->  /pkg-a/index.html
-packages/pkg-b/src/pkg-b-docs.md  -->  /pkg-b/index.html
+packages/pkg-a/src/index.md  -->  /pkg-a/index.html
+packages/pkg-a/src/foo.md    -->  /pkg-a/foo.html
+packages/pkg-b/src/index.md  -->  /pkg-b/index.html
+packages/pkg-b/src/bar.md    -->  /pkg-b/bar.html
 ```
 
 You can achieve this by configuring the [`rewrites`](../reference/site-config#rewrites) option like this:
@@ -180,8 +182,10 @@ You can achieve this by configuring the [`rewrites`](../reference/site-config#re
 // .vitepress/config.js
 export default {
   rewrites: {
-    'packages/pkg-a/src/pkg-a-docs.md': 'pkg-a/index.md',
-    'packages/pkg-b/src/pkg-b-docs.md': 'pkg-b/index.md'
+    'packages/pkg-a/src/index.md': 'pkg-a/index.md',
+    'packages/pkg-a/src/foo.md': 'pkg-a/foo.md',
+    'packages/pkg-b/src/index.md': 'pkg-b/index.md',
+    'packages/pkg-b/src/bar.md': 'pkg-b/bar.md'
   }
 }
 ```
@@ -191,12 +195,22 @@ The `rewrites` option also supports dynamic route parameters. In the above examp
 ```ts
 export default {
   rewrites: {
-    'packages/:pkg/src/(.*)': ':pkg/index.md'
+    'packages/:pkg/src/:slug*': ':pkg/:slug*'
   }
 }
 ```
 
-The rewrite paths are compiled using the `path-to-regexp` package - consult [its documentation](https://github.com/pillarjs/path-to-regexp#parameters) for more advanced syntax.
+The rewrite paths are compiled using the `path-to-regexp` package - consult [its documentation](https://github.com/pillarjs/path-to-regexp/tree/6.x#parameters) for more advanced syntax.
+
+`rewrites` can also be a function that receives the original path and returns the new path:
+
+```ts
+export default {
+  rewrites(id) {
+    return id.replace(/^packages\/([^/]+)\/src\//, '$1/')
+  }
+}
+```
 
 ::: warning Relative Links with Rewrites
 
diff --git a/src/node/plugins/rewritesPlugin.ts b/src/node/plugins/rewritesPlugin.ts
index 51794a1dbd36..236973f14c1b 100644
--- a/src/node/plugins/rewritesPlugin.ts
+++ b/src/node/plugins/rewritesPlugin.ts
@@ -6,22 +6,35 @@ export function resolveRewrites(
   pages: string[],
   userRewrites: UserConfig['rewrites']
 ) {
-  const rewriteRules = Object.entries(userRewrites || {}).map(([from, to]) => ({
-    toPath: compile(`/${to}`, { validate: false }),
-    matchUrl: match(from.startsWith('^') ? new RegExp(from) : from)
-  }))
-
   const pageToRewrite: Record<string, string> = {}
   const rewriteToPage: Record<string, string> = {}
-  if (rewriteRules.length) {
+
+  if (typeof userRewrites === 'function') {
     for (const page of pages) {
-      for (const { matchUrl, toPath } of rewriteRules) {
-        const res = matchUrl(page)
-        if (res) {
-          const dest = toPath(res.params).slice(1)
-          pageToRewrite[page] = dest
-          rewriteToPage[dest] = page
-          break
+      const dest = userRewrites(page)
+      if (dest && dest !== page) {
+        pageToRewrite[page] = dest
+        rewriteToPage[dest] = page
+      }
+    }
+  } else if (typeof userRewrites === 'object') {
+    const rewriteRules = Object.entries(userRewrites || {}).map(
+      ([from, to]) => ({
+        toPath: compile(`/${to}`, { validate: false }),
+        matchUrl: match(from.startsWith('^') ? new RegExp(from) : from)
+      })
+    )
+
+    if (rewriteRules.length) {
+      for (const page of pages) {
+        for (const { matchUrl, toPath } of rewriteRules) {
+          const res = matchUrl(page)
+          if (res) {
+            const dest = toPath(res.params).slice(1)
+            pageToRewrite[page] = dest
+            rewriteToPage[dest] = page
+            break
+          }
         }
       }
     }
diff --git a/src/node/siteConfig.ts b/src/node/siteConfig.ts
index aac3545dacc2..ebd990968a79 100644
--- a/src/node/siteConfig.ts
+++ b/src/node/siteConfig.ts
@@ -162,7 +162,7 @@ export interface UserConfig<ThemeConfig = any>
    *
    * source -> destination
    */
-  rewrites?: Record<string, string>
+  rewrites?: Record<string, string> | ((id: string) => string)
 
   /**
    * @experimental