feat: use Program instead of Module more

src/cjs_parse.rs

@@ -12,6 +12,7 @@ use crate::swc::visit::noop_visit_type;
 use crate::swc::visit::Visit;
 use crate::swc::visit::VisitWith;
 use crate::ParsedSource;
+use crate::ProgramRef;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CjsAnalysis {
@@ -30,7 +31,14 @@ impl ParsedSource {
     }
 
     let mut visitor = CjsVisitor::default();
+<<<<<<< Updated upstream
     visitor.visit_script(self.script());
+=======
+    match self.program_ref() {
+      ProgramRef::Module(n) => visitor.visit_module(n),
+      ProgramRef::Script(n) => visitor.visit_script(n),
+    };
+>>>>>>> Stashed changes
     visitor.take_result()
   }
 }

src/dep.rs

@@ -3,7 +3,6 @@
 use std::collections::HashMap;
 
 use ast::BinExpr;
-use ast::Module;
 use serde::Deserialize;
 use serde::Serialize;
 
@@ -16,31 +15,30 @@ use crate::swc::visit::Visit;
 use crate::swc::visit::VisitWith;
 use crate::MultiThreadedComments;
 use crate::ParsedSource;
+use crate::ProgramRef;
 use crate::SourcePos;
 use crate::SourceRange;
 use crate::SourceRangedForSpanned;
 
 impl ParsedSource {
   /// Analyzes the module for a list of its imports and exports.
   pub fn analyze_dependencies(&self) -> Vec<DependencyDescriptor> {
-    match self.program_ref() {
-      ast::Program::Module(module) => {
-        analyze_module_dependencies(module, self.comments())
-      }
-      ast::Program::Script(_) => vec![],
-    }
+    analyze_program_dependencies(self.program_ref(), self.comments())
   }
 }
 
-pub fn analyze_module_dependencies(
-  module: &Module,
+pub fn analyze_program_dependencies(
+  program: ProgramRef,
   comments: &MultiThreadedComments,
 ) -> Vec<DependencyDescriptor> {
   let mut v = DependencyCollector {
     comments,
     items: vec![],
   };
-  module.visit_with(&mut v);
+  match program {
+    ProgramRef::Module(n) => n.visit_with(&mut v),
+    ProgramRef::Script(n) => n.visit_with(&mut v),
+  }
   v.items
 }
 

src/emit.rs

@@ -5,11 +5,11 @@ use std::string::FromUtf8Error;
 use base64::Engine;
 use thiserror::Error;
 
-use crate::swc::ast::Program;
 use crate::swc::codegen::text_writer::JsWriter;
 use crate::swc::codegen::Node;
 use crate::swc::common::FileName;
 use crate::ModuleSpecifier;
+use crate::ProgramRef;
 use crate::SourceMap;
 
 #[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -92,7 +92,7 @@ pub enum EmitError {
 /// Emits the program as a string of JavaScript code, possibly with the passed
 /// comments, and optionally also a source map.
 pub fn emit(
-  program: &Program,
+  program: ProgramRef,
   comments: &dyn crate::swc::common::comments::Comments,
   source_map: &SourceMap,
   emit_options: &EmitOptions,
@@ -119,9 +119,14 @@ pub fn emit(
       cm: source_map.clone(),
       wr: writer,
     };
-    program
-      .emit_with(&mut emitter)
-      .map_err(EmitError::SwcEmit)?;
+    match program {
+      ProgramRef::Module(n) => {
+        n.emit_with(&mut emitter).map_err(EmitError::SwcEmit)?;
+      }
+      ProgramRef::Script(n) => {
+        n.emit_with(&mut emitter).map_err(EmitError::SwcEmit)?;
+      }
+    }
   }
 
   let mut map: Option<Vec<u8>> = None;

src/parsed_source.rs

@@ -9,6 +9,9 @@ use dprint_swc_ext::common::SourceTextInfo;
 use dprint_swc_ext::common::SourceTextProvider;
 use dprint_swc_ext::common::StartSourcePos;
 use swc_common::Mark;
+use swc_ecma_ast::ModuleDecl;
+use swc_ecma_ast::ModuleItem;
+use swc_ecma_ast::Stmt;
 
 use crate::comments::MultiThreadedComments;
 use crate::scope_analysis_transform;
@@ -59,6 +62,92 @@ impl Globals {
   }
 }
 
+/// A reference to a Program.
+///
+/// It is generally preferable for functions to accept this over `&Program`
+/// because it doesn't require cloning when only owning a `Module` or `Script`.
+#[derive(Debug, Clone, Copy)]
+pub enum ProgramRef<'a> {
+  Module(&'a Module),
+  Script(&'a Script),
+}
+
+impl<'a> From<&'a Program> for ProgramRef<'a> {
+  fn from(program: &'a Program) -> Self {
+    match program {
+      Program::Module(module) => ProgramRef::Module(module),
+      Program::Script(script) => ProgramRef::Script(script),
+    }
+  }
+}
+
+impl<'a> ProgramRef<'a> {
+  pub fn shebang(&self) -> &Option<swc_atoms::Atom> {
+    match self {
+      ProgramRef::Module(m) => &m.shebang,
+      ProgramRef::Script(s) => &s.shebang,
+    }
+  }
+
+  pub fn body(&self) -> impl Iterator<Item = ModuleItemRef<'a>> {
+    match self {
+      ProgramRef::Module(m) => Box::new(m.body.iter().map(|n| n.into()))
+        as Box<dyn Iterator<Item = ModuleItemRef<'a>>>,
+      ProgramRef::Script(s) => Box::new(s.body.iter().map(ModuleItemRef::Stmt)),
+    }
+  }
+
+  pub fn to_owned(&self) -> Program {
+    match self {
+      ProgramRef::Module(m) => Program::Module((*m).clone()),
+      ProgramRef::Script(s) => Program::Script((*s).clone()),
+    }
+  }
+}
+
+impl swc_common::Spanned for ProgramRef<'_> {
+  // ok because we're implementing Spanned
+  #[allow(clippy::disallowed_methods)]
+  #[allow(clippy::disallowed_types)]
+  fn span(&self) -> swc_common::Span {
+    match self {
+      Self::Module(m) => m.span,
+      Self::Script(s) => s.span,
+    }
+  }
+}
+
+/// Reference to a ModuleDecl or Stmt in a Program.
+///
+/// This is used to allow using the same API for the top level
+/// statements when working with a ProgramRef.
+#[derive(Debug, Clone, Copy)]
+pub enum ModuleItemRef<'a> {
+  ModuleDecl(&'a ModuleDecl),
+  Stmt(&'a Stmt),
+}
+
+impl swc_common::Spanned for ModuleItemRef<'_> {
+  // ok because we're implementing Spanned
+  #[allow(clippy::disallowed_methods)]
+  #[allow(clippy::disallowed_types)]
+  fn span(&self) -> swc_common::Span {
+    match self {
+      Self::ModuleDecl(n) => n.span(),
+      Self::Stmt(n) => n.span(),
+    }
+  }
+}
+
+impl<'a> From<&'a ModuleItem> for ModuleItemRef<'a> {
+  fn from(item: &'a ModuleItem) -> Self {
+    match item {
+      ModuleItem::ModuleDecl(n) => ModuleItemRef::ModuleDecl(n),
+      ModuleItem::Stmt(n) => ModuleItemRef::Stmt(n),
+    }
+  }
+}
+
 #[derive(Clone)]
 pub(crate) struct SyntaxContexts {
   pub unresolved: SyntaxContext,
@@ -126,17 +215,20 @@ impl ParsedSource {
   }
 
   /// Gets the parsed program as a reference.
-  pub fn program_ref(&self) -> &Program {
-    &self.0.program
+  pub fn program_ref(&self) -> ProgramRef<'_> {
+    match self.0.program.as_ref() {
+      Program::Module(module) => ProgramRef::Module(module),
+      Program::Script(script) => ProgramRef::Script(script),
+    }
   }
 
   /// Gets the parsed module.
   ///
   /// This will panic if the source is not a module.
   pub fn module(&self) -> &Module {
     match self.program_ref() {
-      Program::Module(module) => module,
-      Program::Script(_) => panic!("Cannot get a module when the source was a script. Use `.program()` instead."),
+      ProgramRef::Module(module) => module,
+      ProgramRef::Script(_) => panic!("Cannot get a module when the source was a script. Use `.program()` instead."),
     }
   }
 
@@ -145,8 +237,8 @@ impl ParsedSource {
   /// This will panic if the source is not a script.
   pub fn script(&self) -> &Script {
     match self.program_ref() {
-      Program::Script(script) => script,
-      Program::Module(_) => panic!("Cannot get a script when the source was a module. Use `.program()` instead."),
+      ProgramRef::Script(script) => script,
+      ProgramRef::Module(_) => panic!("Cannot get a script when the source was a module. Use `.program()` instead."),
     }
   }
 
@@ -245,12 +337,12 @@ impl ParsedSource {
 
   /// Gets if this source is a module.
   pub fn is_module(&self) -> bool {
-    matches!(self.program_ref(), Program::Module(_))
+    matches!(self.program_ref(), ProgramRef::Module(_))
   }
 
   /// Gets if this source is a script.
   pub fn is_script(&self) -> bool {
-    matches!(self.program_ref(), Program::Script(_))
+    matches!(self.program_ref(), ProgramRef::Script(_))
   }
 }
 
@@ -298,8 +390,8 @@ impl ParsedSource {
   ) -> T {
     let program_info = crate::view::ProgramInfo {
       program: match self.program_ref() {
-        Program::Module(module) => crate::view::ProgramRef::Module(module),
-        Program::Script(script) => crate::view::ProgramRef::Script(script),
+        ProgramRef::Module(module) => crate::view::ProgramRef::Module(module),
+        ProgramRef::Script(script) => crate::view::ProgramRef::Script(script),
       },
       text_info: Some(self.text_info_lazy()),
       tokens: self.0.tokens.as_ref().map(|t| t as &[TokenAndSpan]),

src/transpiling/jsx_precompile.rs

@@ -1650,6 +1650,7 @@ mod tests {
   use crate::swc::visit::FoldWith;
   use crate::EmitOptions;
   use crate::ModuleSpecifier;
+  use crate::ProgramRef;
   use crate::SourceMap;
   use pretty_assertions::assert_eq;
   use swc_common::comments::SingleThreadedComments;
@@ -3018,7 +3019,7 @@ const a = _jsxTemplate($$_tpl_1, _jsxAttr("class", "foo"), _jsxAttr("className",
   ) {
     let (source_map, module) = parse(src);
     let mut transform_folder = as_folder(transform);
-    let output = print(&source_map, module.fold_with(&mut transform_folder));
+    let output = print(&source_map, &module.fold_with(&mut transform_folder));
     assert_eq!(output, format!("{}\n", expected_output));
   }
 
@@ -3037,9 +3038,9 @@ const a = _jsxTemplate($$_tpl_1, _jsxAttr("class", "foo"), _jsxAttr("className",
     (source_map, parser.parse_module().unwrap())
   }
 
-  fn print(source_map: &SourceMap, module: Module) -> String {
+  fn print(source_map: &SourceMap, module: &Module) -> String {
     crate::emit::emit(
-      &Program::Module(module),
+      ProgramRef::Module(module),
       &SingleThreadedComments::default(),
       source_map,
       &EmitOptions {

src/transpiling/mod.rs

@@ -348,7 +348,12 @@ fn transpile(
     )
   })?;
 
-  Ok(emit(&program, &comments, &source_map, emit_options)?)
+  Ok(emit(
+    (&program).into(),
+    &comments,
+    &source_map,
+    emit_options,
+  )?)
 }
 
 #[derive(Default, Clone)]