From ca7bc22cf588b21338d7b5baa10caf1337a0828c Mon Sep 17 00:00:00 2001
From: atanasster <atanasster@gmail.com>
Date: Tue, 29 Dec 2020 00:32:17 -0500
Subject: [PATCH] fix: ast cache issue with hmr

---
 .vscode/launch.json                              |  2 +-
 core/instrument/src/index.ts                     |  4 ++--
 core/instrument/src/misc/ast_store.ts            | 15 ++++++++++-----
 .../__snapshots__/mdx-story-source.test.ts.snap  | 16 ++++++++--------
 4 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/.vscode/launch.json b/.vscode/launch.json
index f24edbb20..ee1f22dee 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -112,7 +112,7 @@
       "name": "jest instrument",
       "program": "${workspaceFolder}/node_modules/.bin/jest",
       "cwd": "${workspaceFolder}/core/instrument",
-      "args": ["extract-component"],
+      "args": ["mdx-story-source"],
       "console": "integratedTerminal",
       "internalConsoleOptions": "neverOpen",
       "disableOptimisticBPs": true,
diff --git a/core/instrument/src/index.ts b/core/instrument/src/index.ts
index 949b2c4ba..b5daf64fb 100644
--- a/core/instrument/src/index.ts
+++ b/core/instrument/src/index.ts
@@ -61,7 +61,7 @@ const parseSource = async (
   options: Required<InstrumentOptions>,
 ): Promise<ParseStorieReturnType | undefined> => {
   const source = await prettify(code, options.prettier, filePath);
-  const { ast } = parseFile(source, options.parser, source);
+  const { ast } = parseFile(filePath, options.parser, source);
 
   const store = traverseFn(ast, options, { source, filePath });
   if (!store) {
@@ -161,7 +161,7 @@ export const parseStories = async (
       filepath: filePath,
       ...otherMDXOptions,
     });
-    const { ast } = parseFile(filePath, mergedOptions.parser, code);
+    const { ast } = parseFile(filePath, mergedOptions.parser, code, false);
 
     if (transformMDX) {
       //second pass transform - inject any necessary attributes
diff --git a/core/instrument/src/misc/ast_store.ts b/core/instrument/src/misc/ast_store.ts
index af0116b7e..34b374858 100644
--- a/core/instrument/src/misc/ast_store.ts
+++ b/core/instrument/src/misc/ast_store.ts
@@ -18,14 +18,19 @@ export const parseFile = (
   filePath: string,
   options?: parser.ParserOptions,
   sourceCode?: string,
+  cache: boolean = true,
 ): CacheProps => {
-  if (astCache[filePath]) {
-    return astCache[filePath];
-  }
   const source = sourceCode || fs.readFileSync(filePath, 'utf8');
+  if (!cache) {
+    return { ast: parser.parse(source, options), source };
+  }
+  const filekey = `${filePath}-${fs.statSync(filePath).mtime.toString()}`;
+  if (astCache[filekey]) {
+    return astCache[filekey];
+  }
   const ast = parser.parse(source, options);
-  astCache[filePath] = { ast, source };
-  return astCache[filePath];
+  astCache[filekey] = { ast, source };
+  return astCache[filekey];
 };
 
 export const parseImports = (
diff --git a/core/instrument/test/__snapshots__/mdx-story-source.test.ts.snap b/core/instrument/test/__snapshots__/mdx-story-source.test.ts.snap
index 1212c3651..cdcebe9a8 100644
--- a/core/instrument/test/__snapshots__/mdx-story-source.test.ts.snap
+++ b/core/instrument/test/__snapshots__/mdx-story-source.test.ts.snap
@@ -5,8 +5,8 @@ Object {
   "components": Object {},
   "doc": Object {
     "componentsLookup": Object {},
-    "date": 2020-10-21T01:23:26.442Z,
-    "dateModified": 2020-10-21T01:23:26.442Z,
+    "date": 2020-12-29T05:27:27.469Z,
+    "dateModified": 2020-12-29T05:27:27.469Z,
     "title": "MDX",
   },
   "exports": Object {
@@ -92,8 +92,8 @@ Object {
   "components": Object {},
   "doc": Object {
     "componentsLookup": Object {},
-    "date": 2020-10-21T01:23:26.442Z,
-    "dateModified": 2020-10-21T01:23:26.442Z,
+    "date": 2020-12-29T05:27:27.468Z,
+    "dateModified": 2020-12-29T05:27:27.469Z,
     "title": "MDX",
   },
   "exports": Object {
@@ -231,8 +231,8 @@ Object {
   "components": Object {},
   "doc": Object {
     "componentsLookup": Object {},
-    "date": 2020-10-21T01:23:26.442Z,
-    "dateModified": 2020-10-21T01:23:26.442Z,
+    "date": 2020-12-29T05:27:27.469Z,
+    "dateModified": 2020-12-29T05:27:27.469Z,
     "title": "MDX",
   },
   "exports": Object {
@@ -448,8 +448,8 @@ Object {
   "components": Object {},
   "doc": Object {
     "componentsLookup": Object {},
-    "date": 2020-10-21T01:23:26.442Z,
-    "dateModified": 2020-10-21T01:23:26.442Z,
+    "date": 2020-12-29T05:27:27.469Z,
+    "dateModified": 2020-12-29T05:27:27.470Z,
     "title": "MDX",
   },
   "exports": Object {