Add support for StyleSheet visitor function

Closes parcel-bundler#573

Author: devongovett

node/index.d.ts

@@ -1,4 +1,4 @@
-import type { Angle, CssColor, Rule, CustomProperty, EnvironmentVariable, Function, Image, LengthValue, MediaQuery, Declaration, Ratio, Resolution, Selector, SupportsCondition, Time, Token, TokenOrValue, UnknownAtRule, Url, Variable, StyleRule, DeclarationBlock, ParsedComponent, Multiplier } from './ast';
+import type { Angle, CssColor, Rule, CustomProperty, EnvironmentVariable, Function, Image, LengthValue, MediaQuery, Declaration, Ratio, Resolution, Selector, SupportsCondition, Time, Token, TokenOrValue, UnknownAtRule, Url, Variable, StyleRule, DeclarationBlock, ParsedComponent, Multiplier, StyleSheet } from './ast';
 import { Targets, Features } from './targets';
 
 export * from './ast';
@@ -182,6 +182,7 @@ type EnvironmentVariableVisitors = {
 };
 
 export interface Visitor<C extends CustomAtRules> {
+  StyleSheet?(stylesheet: StyleSheet): StyleSheet<ReturnedDeclaration, ReturnedMediaQuery> | void;
   Rule?: RuleVisitor | RuleVisitors<C>;
   RuleExit?: RuleVisitor | RuleVisitors<C>;
   Declaration?: DeclarationVisitor | DeclarationVisitors;

node/src/transformer.rs

@@ -3,7 +3,6 @@ use std::{
   ops::{Index, IndexMut},
 };
 
-use lightningcss::traits::IntoOwned;
 use lightningcss::{
   media_query::MediaFeatureValue,
   properties::{
@@ -20,6 +19,7 @@ use lightningcss::{
   },
   visitor::{Visit, VisitTypes, Visitor},
 };
+use lightningcss::{stylesheet::StyleSheet, traits::IntoOwned};
 use napi::{Env, JsFunction, JsObject, JsUnknown, Ref, ValueType};
 use serde::{Deserialize, Serialize};
 use smallvec::SmallVec;
@@ -28,6 +28,7 @@ use crate::at_rule_parser::AtRule;
 
 pub struct JsVisitor {
   env: Env,
+  visit_stylesheet: VisitorsRef,
   visit_rule: VisitorsRef,
   rule_map: VisitorsRef,
   property_map: VisitorsRef,
@@ -143,6 +144,7 @@ impl Drop for JsVisitor {
       };
     }
 
+    drop_tuple!(visit_stylesheet);
     drop_tuple!(visit_rule);
     drop_tuple!(rule_map);
     drop_tuple!(visit_declaration);
@@ -199,6 +201,7 @@ impl JsVisitor {
 
     Self {
       env,
+      visit_stylesheet: VisitorsRef::new(get!("StyleSheet", RULES), get!("StyleSheetExit", RULES)),
       visit_rule: VisitorsRef::new(get!("Rule", RULES), get!("RuleExit", RULES)),
       rule_map: VisitorsRef::new(map!("Rule", RULES), get!("RuleExit", RULES)),
       visit_declaration: VisitorsRef::new(get!("Declaration", PROPERTIES), get!("DeclarationExit", PROPERTIES)),
@@ -253,6 +256,26 @@ impl<'i> Visitor<'i, AtRule<'i>> for JsVisitor {
     self.types
   }
 
+  fn visit_stylesheet<'o>(&mut self, stylesheet: &mut StyleSheet<'i, 'o, AtRule<'i>>) -> Result<(), Self::Error> {
+    if self.types.contains(VisitTypes::RULES) {
+      let env = self.env;
+      let visit_stylesheet = self.visit_stylesheet.get::<JsFunction>(&env);
+      if let Some(visit) = visit_stylesheet.for_stage(VisitStage::Enter) {
+        call_visitor(&env, stylesheet, visit)?
+      }
+
+      stylesheet.visit_children(self)?;
+
+      if let Some(visit) = visit_stylesheet.for_stage(VisitStage::Exit) {
+        call_visitor(&env, stylesheet, visit)?
+      }
+
+      Ok(())
+    } else {
+      stylesheet.visit_children(self)
+    }
+  }
+
   fn visit_rule_list(
     &mut self,
     rules: &mut lightningcss::rules::CssRuleList<'i, AtRule<'i>>,
@@ -585,13 +608,23 @@ fn visit<V: Serialize + Deserialize<'static>>(
     .as_ref()
     .and_then(|v| env.get_reference_value_unchecked::<JsFunction>(v).ok())
   {
-    let js_value = env.to_js_value(value)?;
-    let res = visit.call(None, &[js_value])?;
-    let new_value: Option<V> = env.from_js_value(res).map(serde_detach::detach)?;
-    match new_value {
-      Some(new_value) => *value = new_value,
-      None => {}
-    }
+    call_visitor(env, value, &visit)?;
+  }
+
+  Ok(())
+}
+
+fn call_visitor<V: Serialize + Deserialize<'static>>(
+  env: &Env,
+  value: &mut V,
+  visit: &JsFunction,
+) -> napi::Result<()> {
+  let js_value = env.to_js_value(value)?;
+  let res = visit.call(None, &[js_value])?;
+  let new_value: Option<V> = env.from_js_value(res).map(serde_detach::detach)?;
+  match new_value {
+    Some(new_value) => *value = new_value,
+    None => {}
   }
 
   Ok(())

node/test/visitor.test.mjs

@@ -1080,4 +1080,28 @@ test('media query raw', () => {
   assert.equal(res.code.toString(), '.m-1{margin:10px}@media (width>=500px){.sm\\:m-1{margin:10px}}');
 });
 
+test('visit stylesheet', () => {
+  let res = transform({
+    filename: 'test.css',
+    minify: true,
+    code: Buffer.from(`
+      .foo {
+        width: 32px;
+      }
+
+      .bar {
+        width: 80px;
+      }
+    `),
+    visitor: {
+      StyleSheetExit(stylesheet) {
+        stylesheet.rules.sort((a, b) => a.value.selectors[0][0].name.localeCompare(b.value.selectors[0][0].name));
+        return stylesheet;
+      }
+    }
+  });
+
+  assert.equal(res.code.toString(), '.bar{width:80px}.foo{width:32px}');
+});
+
 test.run();

src/stylesheet.rs

@@ -303,10 +303,14 @@ where
 impl<'i, 'o, T, V> Visit<'i, T, V> for StyleSheet<'i, 'o, T>
 where
   T: Visit<'i, T, V>,
-  V: Visitor<'i, T>,
+  V: ?Sized + Visitor<'i, T>,
 {
   const CHILD_TYPES: VisitTypes = VisitTypes::all();
 
+  fn visit(&mut self, visitor: &mut V) -> Result<(), V::Error> {
+    visitor.visit_stylesheet(self)
+  }
+
   fn visit_children(&mut self, visitor: &mut V) -> Result<(), V::Error> {
     self.rules.visit(visitor)
   }

src/visitor.rs

@@ -68,6 +68,7 @@ use crate::{
   },
   rules::{supports::SupportsCondition, CssRule, CssRuleList},
   selector::{Selector, SelectorList},
+  stylesheet::StyleSheet,
   values::{
     angle::Angle,
     color::CssColor,
@@ -150,6 +151,12 @@ pub trait Visitor<'i, T: Visit<'i, T, Self> = DefaultAtRule> {
   /// `Self::TYPES`, but this can be overridden to change the value at runtime.
   fn visit_types(&self) -> VisitTypes;
 
+  /// Visits a stylesheet.
+  #[inline]
+  fn visit_stylesheet<'o>(&mut self, stylesheet: &mut StyleSheet<'i, 'o, T>) -> Result<(), Self::Error> {
+    stylesheet.visit_children(self)
+  }
+
   /// Visits a rule list.
   #[inline]
   fn visit_rule_list(&mut self, rules: &mut CssRuleList<'i, T>) -> Result<(), Self::Error> {