From 15ae0e7c11d1863328a20b3874cd27073347fedc Mon Sep 17 00:00:00 2001
From: overlookmotel <theoverlookmotel@gmail.com>
Date: Mon, 29 Apr 2024 20:03:27 +0100
Subject: [PATCH] Add `transform` function

---
 crates/oxc_ast/src/ast/js.rs             | 16 +++++++-
 crates/oxc_ast/src/traverse/transform.rs | 48 ++++++++++++++++++++++--
 2 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/crates/oxc_ast/src/ast/js.rs b/crates/oxc_ast/src/ast/js.rs
index 3c30704c058e3..d50ba5d4593aa 100644
--- a/crates/oxc_ast/src/ast/js.rs
+++ b/crates/oxc_ast/src/ast/js.rs
@@ -18,8 +18,8 @@ use serde::Serialize;
 #[cfg(feature = "serialize")]
 use tsify::Tsify;
 
-use crate::traverse::{SharedBox, SharedVec};
-use crate::{dummy, traverse::ast::*};
+use crate::dummy;
+use crate::traverse::{ast::*, SharedBox, SharedVec, Token, Traverse, TraverseCtx};
 
 use super::inherit_variants;
 use super::{jsx::*, literal::*, ts::*};
@@ -41,6 +41,18 @@ export interface FormalParameterRest extends Span {
 }
 "#;
 
+/// Traverse AST.
+#[allow(unused_variables)] // TODO: Remove this attr once function is implemented
+pub fn traverse<'a, T: Traverse<'a>>(
+    traverser: &mut T,
+    program: SharedBox<'a, TraversableProgram<'a>>,
+    ctx: &mut TraverseCtx<'a>,
+    tk: &mut Token,
+) {
+    // TODO
+    // walk_traversable_program(traverser, program, ctx, tk);
+}
+
 #[ast_node]
 #[derive(Debug, Hash)]
 #[cfg_attr(feature = "serialize", derive(Serialize, Tsify))]
diff --git a/crates/oxc_ast/src/traverse/transform.rs b/crates/oxc_ast/src/traverse/transform.rs
index f3345cc50f1a8..75ef0299c1ca0 100644
--- a/crates/oxc_ast/src/traverse/transform.rs
+++ b/crates/oxc_ast/src/traverse/transform.rs
@@ -2,10 +2,52 @@
 
 use oxc_allocator::Allocator;
 
-use super::{Ancestor, AstBuilder, SharedBox};
+use super::{ast::TraversableProgram, Ancestor, AstBuilder, GCell, SharedBox, Token, Traverse};
+use crate::ast::{traverse, Program as StandardProgram};
 
-pub fn transform() {
-    // TODO
+/// Run transform visitor on AST.
+///
+/// The provided transformer must implement `Traverse` and will be run on a version of the AST
+/// with interior mutability, allowing traversal in any direction (up or down).
+/// Once the transform is finished, caller can continue to use the standard version of the AST
+/// in the usual way, without interior mutability.
+#[allow(unsafe_code)]
+pub fn transform<'a, T: Traverse<'a>>(
+    transformer: &mut T,
+    program: &mut StandardProgram<'a>,
+    allocator: &'a Allocator,
+) {
+    // Generate `Token` which transformer uses to access the AST.
+    // SAFETY: We only create one token, and it never leaves this function.
+    let mut token = unsafe { Token::new_unchecked() };
+
+    // Create `TraverseCtx` which transformer uses to read ancestry
+    let mut ctx = TraverseCtx::new(allocator);
+
+    // Convert AST to traversable version.
+    //
+    // SAFETY: All standard and traversable AST types are mirrors of each other, with identical layouts.
+    // This is ensured by `#[repr(C)]` on all types. Therefore one can safely be transmuted to the other.
+    // As we hold a `&mut` reference to the AST, it's guaranteed there are no other live references.
+    // We extend the lifetime of ref to `TraversableProgram` to `&'a TraversableProgram`.
+    // This is safe because the node is in the arena, and doesn't move, so the ref is valid for `'a`.
+    // `transformer` could smuggle refs out, but could not use them without a token which is only
+    // available in this function.
+    //
+    // TODO: Refs could be made invalid if the allocator is reset. Hopefully this is impossible
+    // because `Allocator::reset` takes a `&mut` ref to the allocator, so you can't hold any immut refs
+    // to data in the arena at that time. But make sure.
+    #[allow(clippy::ptr_as_ptr, clippy::undocumented_unsafe_blocks)]
+    let program = GCell::from_mut(unsafe { &mut *(program as *mut _ as *mut TraversableProgram) });
+
+    // Run transformer on the traversable AST
+    traverse(transformer, program, &mut ctx, &mut token);
+
+    // The access token goes out of scope at this point, which guarantees that no references
+    // (either mutable or immutable) to the traversable AST or the token still exist.
+    // If the transformer attempts to hold on to any references to the AST, or to the token,
+    // this will produce a compile-time error.
+    // Therefore, the caller can now safely continue using the `&mut Program` that they passed in.
 }
 
 /// Traverse context.