From bd5c3b96be2c6c4a0b84b096c3baa3c30feb95d6 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Fri, 10 Jul 2020 16:30:58 -0400
Subject: [PATCH] feat(compiler-sfc): `<style vars>` CSS variable injection

---
 .../__snapshots__/compileScript.spec.ts.snap  | 55 ++++++++++++++
 .../__tests__/compileScript.spec.ts           | 45 +++++++++++
 .../__tests__/compileStyle.spec.ts            | 15 ++++
 packages/compiler-sfc/src/compileScript.ts    | 39 +++++++---
 packages/compiler-sfc/src/genCssVars.ts       | 76 +++++++++++++++++++
 packages/compiler-sfc/src/index.ts            |  1 +
 packages/compiler-sfc/src/rewriteDefault.ts   | 36 +++++++++
 .../compiler-sfc/src/stylePluginScoped.ts     | 32 +++++---
 8 files changed, 280 insertions(+), 19 deletions(-)
 create mode 100644 packages/compiler-sfc/src/genCssVars.ts
 create mode 100644 packages/compiler-sfc/src/rewriteDefault.ts

diff --git a/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap b/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
index 7b225dcc..e4b096b3 100644
--- a/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
+++ b/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
@@ -105,6 +105,61 @@ export default __define__({
 })"
 `;
 
+exports[`SFC compile <script setup> CSS vars injection <script> w/ default export 1`] = `
+"const __default__ = { setup() {} }
+import { useCSSVars as __useCSSVars__ } from 'vue'
+const __injectCSSVars__ = () => {
+__useCSSVars__(_ctx => ({ color: _ctx.color }))
+}
+const __setup__ = __default__.setup
+__default__.setup = __setup__
+  ? (props, ctx) => { __injectCSSVars__();return __setup__(props, ctx) }
+  : __injectCSSVars__
+export default __default__"
+`;
+
+exports[`SFC compile <script setup> CSS vars injection <script> w/ default export in strings/comments 1`] = `
+"
+          // export default {}
+          const __default__ = {}
+        
+import { useCSSVars as __useCSSVars__ } from 'vue'
+const __injectCSSVars__ = () => {
+__useCSSVars__(_ctx => ({ color: _ctx.color }))
+}
+const __setup__ = __default__.setup
+__default__.setup = __setup__
+  ? (props, ctx) => { __injectCSSVars__();return __setup__(props, ctx) }
+  : __injectCSSVars__
+export default __default__"
+`;
+
+exports[`SFC compile <script setup> CSS vars injection <script> w/ no default export 1`] = `
+"const a = 1
+const __default__ = {}
+import { useCSSVars as __useCSSVars__ } from 'vue'
+const __injectCSSVars__ = () => {
+__useCSSVars__(_ctx => ({ color: _ctx.color }))
+}
+const __setup__ = __default__.setup
+__default__.setup = __setup__
+  ? (props, ctx) => { __injectCSSVars__();return __setup__(props, ctx) }
+  : __injectCSSVars__
+export default __default__"
+`;
+
+exports[`SFC compile <script setup> CSS vars injection w/ <script setup> 1`] = `
+"import { useCSSVars as __useCSSVars__ } from 'vue'
+
+export function setup() {
+const color = 'red'
+__useCSSVars__(_ctx => ({ color }))
+return { color }
+}
+
+export default { setup }"
+`;
+
 exports[`SFC compile <script setup> errors should allow export default referencing imported binding 1`] = `
 "import { bar } from './bar'
           
diff --git a/packages/compiler-sfc/__tests__/compileScript.spec.ts b/packages/compiler-sfc/__tests__/compileScript.spec.ts
index f9ae4cd5..4be8f5e1 100644
--- a/packages/compiler-sfc/__tests__/compileScript.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript.spec.ts
@@ -49,6 +49,10 @@ describe('SFC compile <script setup>', () => {
     )
   })
 
+  test('async/await detection', () => {
+    // TODO
+  })
+
   describe('exports', () => {
     test('export const x = ...', () => {
       const { content, bindings } = compile(
@@ -288,6 +292,47 @@ describe('SFC compile <script setup>', () => {
     })
   })
 
+  describe('CSS vars injection', () => {
+    test('<script> w/ no default export', () => {
+      assertCode(
+        compile(
+          `<script>const a = 1</script>\n` +
+            `<style vars="{ color }">div{ color: var(--color); }</style>`
+        ).content
+      )
+    })
+
+    test('<script> w/ default export', () => {
+      assertCode(
+        compile(
+          `<script>export default { setup() {} }</script>\n` +
+            `<style vars="{ color }">div{ color: var(--color); }</style>`
+        ).content
+      )
+    })
+
+    test('<script> w/ default export in strings/comments', () => {
+      assertCode(
+        compile(
+          `<script>
+          // export default {}
+          export default {}
+        </script>\n` +
+            `<style vars="{ color }">div{ color: var(--color); }</style>`
+        ).content
+      )
+    })
+
+    test('w/ <script setup>', () => {
+      assertCode(
+        compile(
+          `<script setup>export const color = 'red'</script>\n` +
+            `<style vars="{ color }">div{ color: var(--color); }</style>`
+        ).content
+      )
+    })
+  })
+
   describe('errors', () => {
     test('<script> and <script setup> must have same lang', () => {
       expect(
diff --git a/packages/compiler-sfc/__tests__/compileStyle.spec.ts b/packages/compiler-sfc/__tests__/compileStyle.spec.ts
index 08d50c11..628eaffd 100644
--- a/packages/compiler-sfc/__tests__/compileStyle.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileStyle.spec.ts
@@ -237,6 +237,21 @@ describe('SFC scoped CSS', () => {
       ).toHaveBeenWarned()
     })
   })
+
+  describe('<style vars>', () => {
+    test('should rewrite CSS vars in scoped mode', () => {
+      const code = compileScoped(`.foo {
+        color: var(--color);
+        font-size: var(--global:font);
+      }`)
+      expect(code).toMatchInlineSnapshot(`
+        ".foo[test] {
+                color: var(--test-color);
+                font-size: var(--font);
+        }"
+      `)
+    })
+  })
 })
 
 describe('SFC CSS modules', () => {
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index 0efe3848..8b7711e2 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -18,6 +18,7 @@ import {
 } from '@babel/types'
 import { walk } from 'estree-walker'
 import { RawSourceMap } from 'source-map'
+import { genCssVarsCode, injectCssVarsCalls } from './genCssVars'
 
 export interface SFCScriptCompileOptions {
   /**
@@ -49,13 +50,26 @@ export function compileScript(
     )
   }
 
-  const { script, scriptSetup, source, filename } = sfc
+  const { script, scriptSetup, styles, source, filename } = sfc
+  const hasCssVars = styles.some(s => typeof s.attrs.vars === 'string')
+
+  const isTS =
+    (script && script.lang === 'ts') ||
+    (scriptSetup && scriptSetup.lang === 'ts')
+
+  const plugins: ParserPlugin[] = [
+    ...(options.babelParserPlugins || []),
+    ...babelParserDefautPlugins,
+    ...(isTS ? (['typescript'] as const) : [])
+  ]
+
   if (!scriptSetup) {
     if (!script) {
       throw new Error(`SFC contains no <script> tags.`)
     }
     return {
       ...script,
+      content: hasCssVars ? injectCssVarsCalls(sfc, plugins) : script.content,
       bindings: analyzeScriptBindings(script)
     }
   }
@@ -95,13 +109,6 @@ export function compileScript(
   const scriptStartOffset = script && script.loc.start.offset
   const scriptEndOffset = script && script.loc.end.offset
 
-  const isTS = scriptSetup.lang === 'ts'
-  const plugins: ParserPlugin[] = [
-    ...(options.babelParserPlugins || []),
-    ...babelParserDefautPlugins,
-    ...(isTS ? (['typescript'] as const) : [])
-  ]
-
   // 1. process normal <script> first if it exists
   if (script) {
     // import dedupe between <script> and <script setup>
@@ -496,7 +503,7 @@ export function compileScript(
   // 6. wrap setup code with function.
   // export the content of <script setup> as a named export, `setup`.
   // this allows `import { setup } from '*.vue'` for testing purposes.
-  s.appendLeft(startOffset, `\nexport function setup(${args}) {\n`)
+  s.prependLeft(startOffset, `\nexport function setup(${args}) {\n`)
 
   // generate return statement
   let returned = `{ ${Object.keys(setupExports).join(', ')} }`
@@ -511,6 +518,20 @@ export function compileScript(
     returned = `Object.assign(\n  ${returned}\n)`
   }
 
+  // inject `useCSSVars` calls
+  if (hasCssVars) {
+    s.prepend(`import { useCSSVars as __useCSSVars__ } from 'vue'\n`)
+    for (const style of styles) {
+      const vars = style.attrs.vars
+      if (typeof vars === 'string') {
+        s.prependRight(
+          endOffset,
+          `\n${genCssVarsCode(vars, !!style.scoped, setupExports)}`
+        )
+      }
+    }
+  }
+
   s.appendRight(endOffset, `\nreturn ${returned}\n}\n\n`)
 
   // 7. finalize default export
diff --git a/packages/compiler-sfc/src/genCssVars.ts b/packages/compiler-sfc/src/genCssVars.ts
new file mode 100644
index 00000000..0b9a7eba
--- /dev/null
+++ b/packages/compiler-sfc/src/genCssVars.ts
@@ -0,0 +1,76 @@
+import {
+  processExpression,
+  createTransformContext,
+  createSimpleExpression,
+  createRoot,
+  NodeTypes,
+  SimpleExpressionNode
+} from '@vue/compiler-dom'
+import { SFCDescriptor } from './parse'
+import { rewriteDefault } from './rewriteDefault'
+import { ParserPlugin } from '@babel/parser'
+
+export function genCssVarsCode(
+  varsExp: string,
+  scoped: boolean,
+  knownBindings?: Record<string, boolean>
+) {
+  const exp = createSimpleExpression(varsExp, false)
+  const context = createTransformContext(createRoot([]), {
+    prefixIdentifiers: true
+  })
+  if (knownBindings) {
+    // when compiling <script setup> we already know what bindings are exposed
+    // so we can avoid prefixing them from the ctx.
+    for (const key in knownBindings) {
+      context.identifiers[key] = 1
+    }
+  }
+  const transformed = processExpression(exp, context)
+  const transformedString =
+    transformed.type === NodeTypes.SIMPLE_EXPRESSION
+      ? transformed.content
+      : transformed.children
+          .map(c => {
+            return typeof c === 'string'
+              ? c
+              : (c as SimpleExpressionNode).content
+          })
+          .join('')
+
+  return `__useCSSVars__(_ctx => (${transformedString})${
+    scoped ? `, true` : ``
+  })`
+}
+
+// <script setup> already gets the calls injected as part of the transform
+// this is only for single normal <script>
+export function injectCssVarsCalls(
+  sfc: SFCDescriptor,
+  parserPlugins: ParserPlugin[]
+): string {
+  const script = rewriteDefault(
+    sfc.script!.content,
+    `__default__`,
+    parserPlugins
+  )
+
+  let calls = ``
+  for (const style of sfc.styles) {
+    const vars = style.attrs.vars
+    if (typeof vars === 'string') {
+      calls += genCssVarsCode(vars, !!style.scoped) + '\n'
+    }
+  }
+
+  return (
+    script +
+    `\nimport { useCSSVars as __useCSSVars__ } from 'vue'\n` +
+    `const __injectCSSVars__ = () => {\n${calls}}\n` +
+    `const __setup__ = __default__.setup\n` +
+    `__default__.setup = __setup__\n` +
+    `  ? (props, ctx) => { __injectCSSVars__();return __setup__(props, ctx) }\n` +
+    `  : __injectCSSVars__\n` +
+    `export default __default__`
+  )
+}
diff --git a/packages/compiler-sfc/src/index.ts b/packages/compiler-sfc/src/index.ts
index 86636a4d..dbaca968 100644
--- a/packages/compiler-sfc/src/index.ts
+++ b/packages/compiler-sfc/src/index.ts
@@ -3,6 +3,7 @@ export { parse } from './parse'
 export { compileTemplate } from './compileTemplate'
 export { compileStyle, compileStyleAsync } from './compileStyle'
 export { compileScript, analyzeScriptBindings } from './compileScript'
+export { rewriteDefault } from './rewriteDefault'
 
 // Types
 export {
diff --git a/packages/compiler-sfc/src/rewriteDefault.ts b/packages/compiler-sfc/src/rewriteDefault.ts
new file mode 100644
index 00000000..04f783d2
--- /dev/null
+++ b/packages/compiler-sfc/src/rewriteDefault.ts
@@ -0,0 +1,36 @@
+import { parse, ParserPlugin } from '@babel/parser'
+import MagicString from 'magic-string'
+
+const defaultExportRE = /((?:^|\n|;)\s*)export default/
+
+/**
+ * Utility for rewriting `export default` in a script block into a varaible
+ * declaration so that we can inject things into it
+ */
+export function rewriteDefault(
+  input: string,
+  as: string,
+  parserPlugins?: ParserPlugin[]
+): string {
+  if (!defaultExportRE.test(input)) {
+    return input + `\nconst ${as} = {}`
+  }
+
+  const replaced = input.replace(defaultExportRE, `$1const ${as} =`)
+  if (!defaultExportRE.test(replaced)) {
+    return replaced
+  }
+
+  // if the script somehow still contains `default export`, it probably has
+  // multi-line comments or template strings. fallback to a full parse.
+  const s = new MagicString(input)
+  const ast = parse(input, {
+    plugins: parserPlugins
+  }).program.body
+  ast.forEach(node => {
+    if (node.type === 'ExportDefaultDeclaration') {
+      s.overwrite(node.start!, node.declaration.start!, `const ${as} = `)
+    }
+  })
+  return s.toString()
+}
diff --git a/packages/compiler-sfc/src/stylePluginScoped.ts b/packages/compiler-sfc/src/stylePluginScoped.ts
index 5f9abfa5..f9c1d779 100644
--- a/packages/compiler-sfc/src/stylePluginScoped.ts
+++ b/packages/compiler-sfc/src/stylePluginScoped.ts
@@ -1,6 +1,10 @@
 import postcss, { Root } from 'postcss'
 import selectorParser, { Node, Selector } from 'postcss-selector-parser'
 
+const animationNameRE = /^(-\w+-)?animation-name$/
+const animationRE = /^(-\w+-)?animation$/
+const cssVarRE = /\bvar\(--(global:)?([^)]+)\)/g
+
 export default postcss.plugin('vue-scoped', (options: any) => (root: Root) => {
   const id: string = options
   const keyframes = Object.create(null)
@@ -129,21 +133,22 @@ export default postcss.plugin('vue-scoped', (options: any) => (root: Root) => {
     }).processSync(node.selector)
   })
 
-  // If keyframes are found in this <style>, find and rewrite animation names
-  // in declarations.
-  // Caveat: this only works for keyframes and animation rules in the same
-  // <style> element.
-  if (Object.keys(keyframes).length) {
-    root.walkDecls(decl => {
+  const hasKeyframes = Object.keys(keyframes).length
+  root.walkDecls(decl => {
+    // If keyframes are found in this <style>, find and rewrite animation names
+    // in declarations.
+    // Caveat: this only works for keyframes and animation rules in the same
+    // <style> element.
+    if (hasKeyframes) {
       // individual animation-name declaration
-      if (/^(-\w+-)?animation-name$/.test(decl.prop)) {
+      if (animationNameRE.test(decl.prop)) {
         decl.value = decl.value
           .split(',')
           .map(v => keyframes[v.trim()] || v.trim())
           .join(',')
       }
       // shorthand
-      if (/^(-\w+-)?animation$/.test(decl.prop)) {
+      if (animationRE.test(decl.prop)) {
         decl.value = decl.value
           .split(',')
           .map(v => {
@@ -158,8 +163,15 @@ export default postcss.plugin('vue-scoped', (options: any) => (root: Root) => {
           })
           .join(',')
       }
-    })
-  }
+    }
+
+    // rewrite CSS variables
+    if (cssVarRE.test(decl.value)) {
+      decl.value = decl.value.replace(cssVarRE, (_, $1, $2) => {
+        return $1 ? `var(--${$2})` : `var(--${id}-${$2})`
+      })
+    }
+  })
 })
 
 function isSpaceCombinator(node: Node) {