/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::collections::HashMap;

use juno_ast::node_cast;
use juno_ast::GCLock;
use juno_ast::Node;
use juno_support::atom_table::Atom;
use juno_support::declare_opaque_id;
use juno_support::source_manager::SourceId;

use crate::ast::NodeRc;

declare_opaque_id!(DeclId);
declare_opaque_id!(LexicalScopeId);
declare_opaque_id!(FunctionInfoId);
declare_opaque_id!(LabelId);

impl LexicalScopeId {
    pub const GLOBAL_SCOPE_ID: LexicalScopeId = unsafe { LexicalScopeId::new_unchecked(0) };
    pub fn is_global(self) -> bool {
        self == Self::GLOBAL_SCOPE_ID
    }
}

impl FunctionInfoId {
    pub const GLOBAL_FUNCTION_ID: FunctionInfoId = unsafe { FunctionInfoId::new_unchecked(0) };
    pub fn is_global(self) -> bool {
        self == Self::GLOBAL_FUNCTION_ID
    }
}

macro_rules! declare_opaque_list {
    ($name:ident, $storage:ident, $id:ident) => {
        #[derive(Default)]
        pub struct $name(Vec<$storage>);
        impl $name {
            pub fn as_slice(&self) -> &[$storage] {
                self.0.as_slice()
            }
            pub fn get(&self, id: $id) -> &$storage {
                &self.0[id.as_usize()]
            }
            pub(super) fn get_mut(&mut self, id: $id) -> &mut $storage {
                &mut self.0[id.as_usize()]
            }
        }
    };
}

declare_opaque_list!(DeclList, Decl, DeclId);
declare_opaque_list!(LexicalScopeList, LexicalScope, LexicalScopeId);
declare_opaque_list!(FunctionInfoList, FunctionInfo, FunctionInfoId);

#[derive(Debug, PartialOrd, PartialEq, Eq, Copy, Clone)]
pub enum DeclKind {
    // ==== Let-like declarations ===
    Let,
    Const,
    Class,
    Import,
    /// A single catch variable declared like this "catch (e)", see
    /// ES10 B.3.5 VariableStatements in Catch Blocks
    ES5Catch,

    // ==== other declarations ===
    FunctionExprName,
    /// Function declaration visible only in its lexical scope.
    ScopedFunction,

    // ==== Var-like declarations ===
    /// "var" in function scope.
    Var,
    Parameter,
    /// "var" in global scope.
    GlobalProperty,
    UndeclaredGlobalProperty,
}

impl DeclKind {
    /// Return true if this kind of declaration is function scope (and can be
    /// re-declared).
    pub fn is_var_like(self) -> bool {
        self >= Self::Var
    }
    pub fn is_var_like_or_scoped_function(self) -> bool {
        self >= Self::ScopedFunction
    }
    /// Return true if this kind of declaration is lexically scoped (and cannot
    /// be re-declared).
    pub fn is_let_like(self) -> bool {
        self <= Self::ES5Catch
    }
    /// Return true if this kind of declaration is a global property.
    pub fn is_global(self) -> bool {
        self >= Self::GlobalProperty
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Special {
    NotSpecial,
    Arguments,
    Eval,
}

pub struct Decl {
    /// The declared identifier.
    pub name: Atom,
    /// What kind of declaration it is.
    pub kind: DeclKind,
    /// If this is a special declaration, identify which one.
    pub special: Special,
    /// Set to true if this is a function initialized in this scope.
    pub function_in_scope: bool,
    /// The lexical scope of the declaration. Could be nullptr for special
    /// declarations, since they are technically unscoped.
    pub scope: LexicalScopeId,
    /// Whether the variable can be renamed.
    /// False when, e.g., it may be read/written by a local `eval` call.
    pub can_rename: bool,
}

impl Decl {
    fn dump(&self, lock: &GCLock, id: DeclId, indent: usize) {
        println!(
            "{:indent$} Decl#{id} '{name}' {kind:?} {special:?}{function_in_scope}",
            "",
            indent = indent,
            id = id.as_usize(),
            name = lock.str(self.name),
            kind = self.kind,
            special = self.special,
            function_in_scope = if self.function_in_scope {
                " functionInScope"
            } else {
                ""
            },
        );
    }
}

pub struct LexicalScope {
    /// The global depth of this scope, where 0 is the root scope.
    pub depth: u32,
    /// The function owning this lexical scope.
    pub parent_function: FunctionInfoId,
    /// The enclosing lexical scope (it could be in another function).
    pub parent_scope: Option<LexicalScopeId>,
    /// All declarations made in this scope.
    pub decls: Vec<DeclId>,
    /// A list of functions that need to be hoisted and materialized before we
    /// can generate the rest of the scope.
    pub hoisted_functions: Vec<NodeRc>,
    /// True if this scope or any descendent scopes have a local eval call.
    /// If any descendent uses local eval,
    /// it's impossible to know whether local variables are modified.
    pub local_eval: bool,
}

impl LexicalScope {
    fn dump(&self, sem: &SemContext, lock: &GCLock, id: LexicalScopeId, indent: usize) {
        println!(
            "{:indent$} Scope#{id}",
            "",
            indent = indent,
            id = id.as_usize()
        );
        for &decl in &self.decls {
            sem.decl(decl).dump(lock, decl, indent + 1);
        }
        for func in &self.hoisted_functions {
            let node = func.node(lock);
            println!(
                "{:indent$} hoistedFunction {name}",
                "",
                indent = indent + 1,
                name = lock.str(
                    node_cast!(
                        Node::Identifier,
                        node_cast!(Node::FunctionDeclaration, node).id.unwrap()
                    )
                    .name
                ),
            );
        }
    }
}

#[derive(Debug)]
pub struct FunctionInfo {
    /// The function surrounding this function.
    pub parent_function: Option<FunctionInfoId>,
    /// The enclosing lexical scope.
    pub parent_scope: Option<LexicalScopeId>,
    /// True if the function is strict mode.
    pub strict: bool,
    /// All lexical scopes in this function. The first one is the function scope.
    pub scopes: Vec<LexicalScopeId>,
    /// True if this function is an arrow function.
    pub arrow: bool,
    /// The implicitly declared "arguments" object. It is declared only if it is used.
    pub arguments_decl: Option<DeclId>,
    /// How many labels have been allocated in this function so far.
    num_labels: usize,
}

impl FunctionInfo {
    /// Allocate a new label and return its ID.
    pub fn new_label(&mut self) -> LabelId {
        let result = LabelId::new(self.num_labels);
        self.num_labels += 1;
        result
    }

    fn dump(&self, sem: &SemContext, lock: &GCLock, id: FunctionInfoId, indent: usize) {
        println!(
            "{:indent$} Func#{id}",
            "",
            indent = indent,
            id = id.as_usize()
        );
        let mut children = HashMap::<LexicalScopeId, Vec<LexicalScopeId>>::new();

        for &scope_id in &self.scopes {
            let scope = sem.scope(scope_id);
            if let Some(parent) = &scope.parent_scope {
                match children.get_mut(parent) {
                    Some(v) => v.push(scope_id),
                    None => {
                        children.insert(*parent, vec![scope_id]);
                    }
                };
            }
        }

        self.dump_scope(sem, lock, &children, self.scopes[0], indent + 1);
    }

    fn dump_scope(
        &self,
        sem: &SemContext,
        lock: &GCLock,
        children: &HashMap<LexicalScopeId, Vec<LexicalScopeId>>,
        id: LexicalScopeId,
        level: usize,
    ) {
        sem.scope(id).dump(sem, lock, id, level);
        if let Some(next) = children.get(&id) {
            for c in next {
                self.dump_scope(sem, lock, children, *c, level + 1);
            }
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub enum Resolution {
    /// Unable to resolve due to presence of `eval` or `with`.
    Unresolvable,
    /// Resolved to a proper declaration.
    Decl(DeclId),
}

#[derive(Default)]
pub struct SemContext {
    pub decls: DeclList,
    pub scopes: LexicalScopeList,
    pub funcs: FunctionInfoList,
    /// Declarations associated with Identifier AST nodes or marked [`Resolution::Unresolvable`].
    ident_decls: HashMap<NodeRc, Resolution>,
    /// Lexical scopes associated with AST nodes. Usually BlockStatement, but
    /// occasionally others.
    node_scopes: HashMap<NodeRc, LexicalScopeId>,
    /// Resolved `require` calls.
    requires: HashMap<NodeRc, SourceId>,
}

impl SemContext {
    pub(super) fn new_function(
        &mut self,
        parent_function: Option<FunctionInfoId>,
        parent_scope: Option<LexicalScopeId>,
        strict: bool,
        arrow: bool,
    ) -> (FunctionInfoId, &FunctionInfo) {
        self.funcs.0.push(FunctionInfo {
            parent_function,
            parent_scope,
            strict,
            scopes: Default::default(),
            arrow,
            arguments_decl: Default::default(),
            num_labels: 0,
        });
        (
            FunctionInfoId::new(self.funcs.0.len() - 1),
            self.funcs.0.last().unwrap(),
        )
    }

    pub(super) fn new_scope(
        &mut self,
        in_function: FunctionInfoId,
        parent_scope: Option<LexicalScopeId>,
        depth: u32,
    ) -> (LexicalScopeId, &LexicalScope) {
        self.scopes.0.push(LexicalScope {
            depth,
            parent_function: in_function,
            parent_scope,
            decls: Default::default(),
            hoisted_functions: Default::default(),
            local_eval: false,
        });
        let scope_id = LexicalScopeId::new(self.scopes.0.len() - 1);

        self.funcs.get_mut(in_function).scopes.push(scope_id);

        (scope_id, self.scopes.0.last().unwrap())
    }

    pub(super) fn new_decl_special(
        &mut self,
        scope: LexicalScopeId,
        name: Atom,
        kind: DeclKind,
        special: Special,
    ) -> DeclId {
        self.decls.0.push(Decl {
            name,
            kind,
            special,
            function_in_scope: false,
            scope,
            can_rename: false,
        });
        let decl_id = DeclId::new(self.decls.0.len() - 1);
        self.scopes.get_mut(scope).decls.push(decl_id);
        decl_id
    }
    pub(super) fn new_decl(&mut self, scope: LexicalScopeId, name: Atom, kind: DeclKind) -> DeclId {
        self.new_decl_special(scope, name, kind, Special::NotSpecial)
    }

    pub(super) fn new_global(&mut self, name: Atom, kind: DeclKind) -> DeclId {
        self.new_decl(
            self.global_scope_id()
                .expect("global scope has not been created"),
            name,
            kind,
        )
    }

    pub fn all_ident_decls(&self) -> &HashMap<NodeRc, Resolution> {
        &self.ident_decls
    }
    pub(super) fn set_ident_decl(&mut self, node: NodeRc, decl: DeclId) {
        self.ident_decls.insert(node, Resolution::Decl(decl));
    }
    pub(super) fn set_ident_unresolvable(&mut self, node: NodeRc) {
        self.ident_decls.insert(node, Resolution::Unresolvable);
    }
    pub fn ident_decl(&self, node: &NodeRc) -> Option<Resolution> {
        self.ident_decls.get(node).copied()
    }

    pub fn all_node_scopes(&self) -> &HashMap<NodeRc, LexicalScopeId> {
        &self.node_scopes
    }
    pub(super) fn set_node_scope(&mut self, node: NodeRc, scope: LexicalScopeId) {
        self.node_scopes.insert(node, scope);
    }
    pub fn node_scope(&self, node: NodeRc) -> Option<LexicalScopeId> {
        self.node_scopes.get(&node).copied()
    }

    pub fn all_decls(&self) -> &[Decl] {
        self.decls.as_slice()
    }
    pub fn decl(&self, id: DeclId) -> &Decl {
        self.decls.get(id)
    }
    pub(super) fn decl_mut(&mut self, id: DeclId) -> &mut Decl {
        self.decls.get_mut(id)
    }

    pub fn all_scopes(&self) -> &[LexicalScope] {
        self.scopes.as_slice()
    }
    pub fn scope(&self, id: LexicalScopeId) -> &LexicalScope {
        self.scopes.get(id)
    }
    pub(super) fn scope_mut(&mut self, id: LexicalScopeId) -> &mut LexicalScope {
        self.scopes.get_mut(id)
    }

    pub fn all_functions(&self) -> &[FunctionInfo] {
        self.funcs.as_slice()
    }
    pub fn function(&self, id: FunctionInfoId) -> &FunctionInfo {
        self.funcs.get(id)
    }
    pub(super) fn function_mut(&mut self, id: FunctionInfoId) -> &mut FunctionInfo {
        self.funcs.get_mut(id)
    }

    pub fn all_requires(&self) -> &HashMap<NodeRc, SourceId> {
        &self.requires
    }
    pub fn add_require(&mut self, call: NodeRc, file_id: SourceId) {
        self.requires.insert(call, file_id);
    }

    /// Return the id of the global scope in the context. This may seem
    /// redundant, since the ID is constant. The idea here that the global scope
    /// may not have been created yet.
    pub fn global_scope_id(&self) -> Option<LexicalScopeId> {
        if self.scopes.0.is_empty() {
            None
        } else {
            Some(LexicalScopeId::GLOBAL_SCOPE_ID)
        }
    }

    /// Return the id of the global function in the context. This may seem
    /// redundant, since the ID is constant. The idea here that the global
    /// function may not have been created yet.
    pub fn global_function_id(&self) -> Option<FunctionInfoId> {
        if self.funcs.0.is_empty() {
            None
        } else {
            Some(FunctionInfoId::GLOBAL_FUNCTION_ID)
        }
    }

    /// Return the special arguments declaration in the specified function.
    /// It may be None if "arguments" hasn't been used.
    pub fn func_arguments_opt(&self, func: FunctionInfoId) -> Option<DeclId> {
        self.function(func).arguments_decl
    }
    /// Return or create the special arguments declaration in the specified
    /// function. If `func` is an arrow function, find the closest ancestor that is not an arrow
    /// function and use that function's `arguments`.
    /// - `name`: the object doesn't have access to the atom table, so it has
    ///   to be passed in.
    pub(super) fn func_arguments_decl(&mut self, func: FunctionInfoId, name: Atom) -> DeclId {
        let mut arguments_func = func;
        while self.function(arguments_func).arrow {
            match self.function(arguments_func).parent_function {
                None => break,
                Some(parent) => arguments_func = parent,
            };
        }

        if let Some(d) = self.function(arguments_func).arguments_decl {
            return d;
        }

        let decl = if Some(arguments_func) == self.global_function_id() {
            // `arguments` must simply be treated as a global property in top level contexts.
            self.new_decl(
                *self
                    .function(arguments_func)
                    .scopes
                    .first()
                    .expect("Function must contain a scope"),
                name,
                DeclKind::UndeclaredGlobalProperty,
            )
        } else {
            self.new_decl_special(
                *self
                    .function(arguments_func)
                    .scopes
                    .first()
                    .expect("Function must contain a scope"),
                name,
                DeclKind::Var,
                Special::Arguments,
            )
        };
        self.function_mut(arguments_func).arguments_decl = Some(decl);

        decl
    }

    /// Dumps information about the semantic context directly to stdout.
    /// Intended for testing and debugging.
    pub fn dump(&self, lock: &GCLock) {
        println!("SemContext");
        println!("{:7} functions", self.all_functions().len());
        println!("{:7} lexical scopes", self.all_scopes().len());
        println!("{:7} declarations", self.all_decls().len());
        println!("{:7} ident resolutions", self.all_ident_decls().len());
        if !self.all_requires().is_empty() {
            println!("{:7} require resolutions", self.all_requires().len());
        }
        println!();

        let mut children = HashMap::<FunctionInfoId, Vec<FunctionInfoId>>::new();
        for i in 0..self.funcs.0.len() {
            let f = FunctionInfoId::new(i);
            if let Some(parent) = &self.function(f).parent_function {
                match children.get_mut(parent) {
                    Some(v) => v.push(f),
                    None => {
                        children.insert(*parent, vec![f]);
                    }
                };
            }
        }

        for id in 0..self.decls.0.len() {
            let id = DeclId::new(id);
            let decl = self.decl(id);
            if decl.kind == DeclKind::GlobalProperty {
                decl.dump(lock, id, 1);
            }
        }

        self.dump_function(lock, &children, FunctionInfoId::GLOBAL_FUNCTION_ID, 0);
    }

    fn dump_function(
        &self,
        lock: &GCLock,
        children: &HashMap<FunctionInfoId, Vec<FunctionInfoId>>,
        id: FunctionInfoId,
        level: usize,
    ) {
        self.function(id).dump(self, lock, id, level);
        if let Some(next) = children.get(&id) {
            for c in next {
                self.dump_function(lock, children, *c, level + 1);
            }
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#[macro_use]
mod ast;

mod gen_js;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use juno::ast::node_cast;
use juno::ast::*;
use juno::hparser;

mod validate;

#[test]
#[should_panic]
fn test_node_outlives_context() {
    let ast;
    {
        let mut ctx = Context::new();
        ast = {
            let gc = GCLock::new(&mut ctx);
            NodeRc::from_node(
                &gc,
                builder::NumericLiteral::build_template(
                    &gc,
                    template::NumericLiteral {
                        metadata: Default::default(),
                        value: 1.0f64,
                    },
                ),
            )
        };
        // Forget `ast` in order to prevent the `Drop` impl from being called on panic.
        std::mem::forget(ast);
    }
}

#[test]
#[should_panic]
#[allow(clippy::redundant_clone)]
fn test_node_clone_outlives_context() {
    let ast;
    {
        let mut ctx = Context::new();
        {
            let ast_orig = {
                let gc = GCLock::new(&mut ctx);
                NodeRc::from_node(
                    &gc,
                    builder::NumericLiteral::build_template(
                        &gc,
                        template::NumericLiteral {
                            metadata: Default::default(),
                            value: 1.0f64,
                        },
                    ),
                )
            };
            ast = ast_orig.clone();
            // Forget `ast` in order to prevent the `Drop` impl from being called on panic.
            std::mem::forget(ast);
        }
    }
}

#[test]
fn test_list_empty() {
    let mut ctx = Context::new();
    let gc = GCLock::new(&mut ctx);
    let ast = builder::Program::build_template(
        &gc,
        template::Program {
            metadata: Default::default(),
            body: NodeList::new(&gc),
        },
    );
    match ast {
        Node::Program(Program { body, .. }) => {
            assert!(body.is_empty());
            assert_eq!(body.len(), 0);
        }
        _ => unreachable!(),
    }
}

#[test]
fn test_list_elements() {
    let mut ctx = Context::new();
    let gc = GCLock::new(&mut ctx);
    let ast = builder::Program::build_template(
        &gc,
        template::Program {
            metadata: Default::default(),
            body: NodeList::from_iter(
                &gc,
                [
                    builder::NullLiteral::build_template(
                        &gc,
                        template::NullLiteral {
                            metadata: Default::default(),
                        },
                    ),
                    builder::BooleanLiteral::build_template(
                        &gc,
                        template::BooleanLiteral {
                            metadata: Default::default(),
                            value: false,
                        },
                    ),
                ],
            ),
        },
    );
    match ast {
        Node::Program(Program { body, .. }) => {
            assert!(!body.is_empty());
            assert_eq!(body.len(), 2);
            let mut it = body.iter();
            assert!(matches!(it.next(), Some(Node::NullLiteral(_))));
            assert!(matches!(
                it.next(),
                Some(Node::BooleanLiteral(BooleanLiteral { value: false, .. }))
            ));
        }
        _ => unreachable!(),
    }
}

/// Tests to ensure that list elements allocated off the free list are valid.
#[test]
fn test_list_elements_from_freelist() {
    let mut ctx = Context::new();
    // First, allocate a two-element list.
    {
        let gc = GCLock::new(&mut ctx);
        let ast = builder::Program::build_template(
            &gc,
            template::Program {
                metadata: Default::default(),
                body: NodeList::from_iter(
                    &gc,
                    [
                        builder::NullLiteral::build_template(
                            &gc,
                            template::NullLiteral {
                                metadata: Default::default(),
                            },
                        ),
                        builder::BooleanLiteral::build_template(
                            &gc,
                            template::BooleanLiteral {
                                metadata: Default::default(),
                                value: false,
                            },
                        ),
                    ],
                ),
            },
        );
        assert_eq!(node_cast!(Node::Program, ast).body.len(), 2);
    }
    // Free both elements.
    ctx.gc();
    // Allocate two one-element lists to ensure they both have one element
    // and the old connection between the two elements was not retained.
    {
        let gc = GCLock::new(&mut ctx);
        let ast = builder::Program::build_template(
            &gc,
            template::Program {
                metadata: Default::default(),
                body: NodeList::from_iter(
                    &gc,
                    [builder::NullLiteral::build_template(
                        &gc,
                        template::NullLiteral {
                            metadata: Default::default(),
                        },
                    )],
                ),
            },
        );
        assert_eq!(node_cast!(Node::Program, ast).body.len(), 1);
        let ast = builder::Program::build_template(
            &gc,
            template::Program {
                metadata: Default::default(),
                body: NodeList::from_iter(
                    &gc,
                    [builder::NullLiteral::build_template(
                        &gc,
                        template::NullLiteral {
                            metadata: Default::default(),
                        },
                    )],
                ),
            },
        );
        assert_eq!(node_cast!(Node::Program, ast).body.len(), 1);
    }
}

#[test]
#[allow(clippy::float_cmp)]
fn test_visit() {
    let mut ctx = Context::new();
    let ast = {
        let gc = GCLock::new(&mut ctx);
        NodeRc::from_node(
            &gc,
            template::BlockStatement {
                metadata: Default::default(),
                body: NodeList::from_iter(
                    &gc,
                    [
                        template::ExpressionStatement {
                            metadata: Default::default(),
                            expression: template::NumericLiteral {
                                metadata: Default::default(),
                                value: 1.0,
                            }
                            .build(&gc),
                            directive: None,
                        }
                        .build(&gc),
                        template::ExpressionStatement {
                            metadata: Default::default(),
                            expression: template::NumericLiteral {
                                metadata: Default::default(),
                                value: 2.0,
                            }
                            .build(&gc),
                            directive: None,
                        }
                        .build(&gc),
                    ],
                ),
            }
            .build(&gc),
        )
    };

    // Accumulates the numbers found in the AST.
    struct NumberFinder {
        acc: Vec<f64>,
    }

    impl<'gc> Visitor<'gc> for NumberFinder {
        /// Visit the Node `node` with the given `parent`.
        fn call(&mut self, ctx: &'gc GCLock, node: &'gc Node<'gc>, path: Option<Path<'gc>>) {
            if let Node::NumericLiteral(NumericLiteral { value, .. }) = node {
                assert!(matches!(
                    path.unwrap().parent,
                    Node::ExpressionStatement(ExpressionStatement { .. })
                ));
                self.acc.push(*value);
            }
            node.visit_children(ctx, self);
        }
    }

    let mut visitor = NumberFinder { acc: vec![] };
    let gc = GCLock::new(&mut ctx);
    ast.node(&gc).visit(&gc, &mut visitor, None);
    assert_eq!(visitor.acc, [1.0, 2.0]);
}

#[test]
fn test_visit_mut() {
    let mut ctx = Context::new();

    let (x, y) = (1.0, 2.0);
    let ast = {
        let gc = GCLock::new(&mut ctx);
        NodeRc::from_node(
            &gc,
            builder::ExpressionStatement::build_template(
                &gc,
                template::ExpressionStatement {
                    metadata: Default::default(),
                    directive: None,
                    expression: builder::BinaryExpression::build_template(
                        &gc,
                        template::BinaryExpression {
                            metadata: Default::default(),
                            operator: BinaryExpressionOperator::Plus,
                            left: builder::NumericLiteral::build_template(
                                &gc,
                                template::NumericLiteral {
                                    metadata: Default::default(),
                                    value: x,
                                },
                            ),
                            right: builder::UnaryExpression::build_template(
                                &gc,
                                template::UnaryExpression {
                                    metadata: Default::default(),
                                    prefix: true,
                                    operator: UnaryExpressionOperator::Minus,
                                    argument: builder::NumericLiteral::build_template(
                                        &gc,
                                        template::NumericLiteral {
                                            metadata: Default::default(),
                                            value: y,
                                        },
                                    ),
                                },
                            ),
                        },
                    ),
                },
            ),
        )
    };

    struct Pass {}

    impl<'gc> VisitorMut<'gc> for Pass {
        fn call(
            &mut self,
            ctx: &'gc GCLock,
            node: &'gc Node<'gc>,
            _path: Option<Path<'gc>>,
        ) -> TransformResult<&'gc Node<'gc>> {
            if let Node::BinaryExpression(
                e1 @ BinaryExpression {
                    operator: BinaryExpressionOperator::Plus,
                    right:
                        Node::UnaryExpression(UnaryExpression {
                            operator: UnaryExpressionOperator::Minus,
                            argument: e2,
                            ..
                        }),
                    ..
                },
            ) = node
            {
                let mut builder = builder::BinaryExpression::from_node(e1);
                builder.operator(BinaryExpressionOperator::Minus);
                builder.right(e2);
                return node.replace_with_new(builder.into(), ctx, self);
            }
            node.visit_children_mut(ctx, self)
        }
    }
    let mut pass = Pass {};

    let transformed = {
        let gc = GCLock::new(&mut ctx);
        NodeRc::from_node(&gc, ast.node(&gc).visit_mut(&gc, &mut pass, None).unwrap())
    };

    {
        let gc = GCLock::new(&mut ctx);
        match transformed.node(&gc) {
            Node::ExpressionStatement(ExpressionStatement {
                expression:
                    Node::BinaryExpression(BinaryExpression {
                        operator: BinaryExpressionOperator::Minus,
                        left:
                            Node::NumericLiteral(NumericLiteral {
                                value: val_left, ..
                            }),
                        right:
                            Node::NumericLiteral(NumericLiteral {
                                value: val_right, ..
                            }),
                        ..
                    }),
                ..
            }) => {
                assert_eq!(
                    (*val_left, *val_right),
                    (x, y),
                    "Transformation failed: {:#?}",
                    transformed.node(&gc),
                );
            }
            _ => panic!("Transformation failed: {:#?}", transformed.node(&gc)),
        };
    }

    {
        ctx.gc();
    }
}

#[test]
fn test_replace_var_decls() {
    let mut ctx = Context::new();
    let ast = hparser::parse_with_flags(Default::default(), "var x, y;", &mut ctx).unwrap();

    {
        let gc = GCLock::new(&mut ctx);
        match ast.node(&gc) {
            Node::Program(Program { body, .. }) => {
                assert_eq!(body.len(), 1, "Program is {:#?}", ast.node(&gc));
            }
            _ => panic!("Parse failed: {:#?}", ast.node(&gc)),
        };
    }

    struct Pass {}
    impl<'gc> VisitorMut<'gc> for Pass {
        fn call(
            &mut self,
            lock: &'gc GCLock,
            node: &'gc Node<'gc>,
            path: Option<Path<'gc>>,
        ) -> TransformResult<&'gc Node<'gc>> {
            match node {
                Node::VariableDeclaration(VariableDeclaration {
                    metadata: _,
                    kind,
                    declarations,
                }) if declarations.len() > 1 => {
                    assert_eq!(path.unwrap().field, NodeField::body);
                    let mut result: Vec<builder::Builder> = Vec::new();
                    for decl in *declarations {
                        result.push(builder::Builder::VariableDeclaration(
                            builder::VariableDeclaration::from_template(
                                template::VariableDeclaration {
                                    metadata: (*decl.range()).into(),
                                    kind: *kind,
                                    declarations: NodeList::from_iter(lock, [decl]),
                                },
                            ),
                        ));
                    }
                    node.replace_with_multiple(result, lock, self)
                }
                _ => node.visit_children_mut(lock, self),
            }
        }
    }
    let mut pass = Pass {};

    let transformed = {
        let gc = GCLock::new(&mut ctx);
        NodeRc::from_node(&gc, ast.node(&gc).visit_mut(&gc, &mut pass, None).unwrap())
    };

    {
        let gc = GCLock::new(&mut ctx);
        match transformed.node(&gc) {
            Node::Program(Program { body, .. }) => {
                assert_eq!(body.len(), 2, "Program is {:#?}", transformed.node(&gc));
                let x_decl = node_cast!(Node::VariableDeclaration, body.head().unwrap())
                    .declarations
                    .head()
                    .unwrap();
                assert_eq!(
                    gc.ctx().str(
                        node_cast!(
                            Node::Identifier,
                            node_cast!(Node::VariableDeclarator, x_decl).id
                        )
                        .name
                    ),
                    "x"
                );
                assert_eq!(x_decl.range().start.line, 1);
                assert_eq!(x_decl.range().start.col, 5);
                assert_eq!(
                    gc.ctx().str(
                        node_cast!(
                            Node::Identifier,
                            node_cast!(
                                Node::VariableDeclarator,
                                node_cast!(Node::VariableDeclaration, body.iter().nth(1).unwrap())
                                    .declarations
                                    .head()
                                    .unwrap()
                            )
                            .id
                        )
                        .name
                    ),
                    "y"
                );
            }
            _ => panic!("Transformation failed: {:#?}", transformed.node(&gc)),
        };
    }
}

#[test]
fn test_many_nodes() {
    let mut ctx = Context::new();
    let mut cached = None;
    let mut val = 0f64;
    for _ in 0..10 {
        {
            let gc = GCLock::new(&mut ctx);
            for i in 0..10_000 {
                cached = Some(NodeRc::from_node(
                    &gc,
                    builder::NumericLiteral::build_template(
                        &gc,
                        template::NumericLiteral {
                            metadata: Default::default(),
                            value: i as f64,
                        },
                    ),
                ));
                val = i as f64;
            }
        }
        ctx.gc();
    }

    let gc = GCLock::new(&mut ctx);
    match cached.unwrap().node(&gc) {
        Node::NumericLiteral(NumericLiteral { value, .. }) => {
            assert!(
                (*value - val).abs() < f64::EPSILON,
                "Incorrect cached value: {:#?}",
                *value
            );
        }
        n => {
            panic!("Incorrect cached value: {:#?}", n);
        }
    };
}

#[test]
fn test_store_node() {
    let mut ctx = Context::new();

    struct Foo<'gc> {
        n: Option<&'gc Node<'gc>>,
    }

    impl<'gc> Foo<'gc> {
        fn set_n(&mut self, node: &'gc Node<'gc>) {
            self.n = Some(node);
        }
    }

    impl<'gc> Visitor<'gc> for Foo<'gc> {
        fn call(&mut self, gc: &'gc GCLock, node: &'gc Node<'gc>, _path: Option<Path<'gc>>) {
            self.set_n(node);
            node.visit_children(gc, self)
        }
    }

    {
        let gc = GCLock::new(&mut ctx);
        let mut pass = Foo { n: None };
        let ast = NodeRc::from_node(
            &gc,
            builder::NumericLiteral::build_template(
                &gc,
                template::NumericLiteral {
                    metadata: Default::default(),
                    value: 1.0f64,
                },
            ),
        );
        ast.node(&gc).visit(&gc, &mut pass, None);
        assert!(pass.n.is_some());
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use juno::ast::*;
use juno::hparser;
use juno::hparser::ParserDialect;
use juno::hparser::ParserFlags;

fn validate_src_with_flags(
    flags: hparser::ParserFlags,
    src: &str,
) -> Result<(), TreeValidationError> {
    let mut ctx = Context::new();
    let ast = hparser::parse_with_flags(flags, src, &mut ctx).unwrap();
    validate_tree(&mut ctx, &ast)
}

fn validate_src(src: &str) -> Result<(), TreeValidationError> {
    validate_src_with_flags(Default::default(), src)
}

fn validate_src_jsx(src: &str) -> Result<(), TreeValidationError> {
    validate_src_with_flags(
        ParserFlags {
            enable_jsx: true,
            ..Default::default()
        },
        src,
    )
}

fn validate_src_flow(src: &str) -> Result<(), TreeValidationError> {
    validate_src_with_flags(
        ParserFlags {
            dialect: ParserDialect::Flow,
            ..Default::default()
        },
        src,
    )
}

#[test]
fn test_valid() {
    let mut ctx = Context::new();
    let return_stmt = {
        let gc = GCLock::new(&mut ctx);
        NodeRc::from_node(
            &gc,
            builder::ReturnStatement::build_template(
                &gc,
                template::ReturnStatement {
                    metadata: Default::default(),
                    argument: None,
                },
            ),
        )
    };
    assert!(validate_tree_pure(&mut ctx, &return_stmt).is_ok());

    let return_stmt = {
        let gc = GCLock::new(&mut ctx);
        NodeRc::from_node(
            &gc,
            builder::ReturnStatement::build_template(
                &gc,
                template::ReturnStatement {
                    metadata: Default::default(),
                    argument: Some(builder::NumericLiteral::build_template(
                        &gc,
                        template::NumericLiteral {
                            metadata: Default::default(),
                            value: 1.0,
                        },
                    )),
                },
            ),
        )
    };
    assert!(validate_tree_pure(&mut ctx, &return_stmt).is_ok());

    let return_stmt = {
        let gc = GCLock::new(&mut ctx);
        NodeRc::from_node(
            &gc,
            builder::ReturnStatement::build_template(
                &gc,
                template::ReturnStatement {
                    metadata: Default::default(),
                    argument: Some(builder::ReturnStatement::build_template(
                        &gc,
                        template::ReturnStatement {
                            metadata: Default::default(),
                            argument: None,
                        },
                    )),
                },
            ),
        )
    };
    assert!(validate_tree_pure(&mut ctx, &return_stmt).is_err());
}

#[test]
fn test_error() {
    let mut ctx = Context::new();

    let ast: NodeRc = {
        let gc = GCLock::new(&mut ctx);
        NodeRc::from_node(
            &gc,
            builder::BlockStatement::build_template(
                &gc,
                template::BlockStatement {
                    metadata: Default::default(),
                    body: NodeList::from_iter(
                        &gc,
                        [builder::ReturnStatement::build_template(
                            &gc,
                            template::ReturnStatement {
                                metadata: Default::default(),
                                argument: Some(builder::ReturnStatement::build_template(
                                    &gc,
                                    template::ReturnStatement {
                                        metadata: Default::default(),
                                        argument: None,
                                    },
                                )),
                            },
                        )],
                    ),
                },
            ),
        )
    };

    let bad_ret: NodeRc = {
        let gc = GCLock::new(&mut ctx);
        match ast.node(&gc) {
            Node::BlockStatement(BlockStatement { body, .. }) => {
                NodeRc::from_node(&gc, body.head().unwrap())
            }
            _ => {
                unreachable!("bad match");
            }
        }
    };
    match validate_tree_pure(&mut ctx, &ast) {
        Ok(()) => {
            panic!("Must be error");
        }
        Err(e) => {
            assert_eq!(e.len(), 1);
            assert_eq!(e[0].node, bad_ret);
        }
    }
}

#[test]
fn test_literals() {
    validate_src("({});").unwrap();
    validate_src("({x: y});").unwrap();
    validate_src("({x: y, ...z});").unwrap();
    validate_src("[]").unwrap();
    validate_src("[x,y,,z]").unwrap();
    validate_src("[x, y, ...z]").unwrap();
}

#[test]
fn test_patterns() {
    validate_src("[x] = y;").unwrap();
    validate_src("[x.y , , ...z] = y;").unwrap();
    validate_src("({x} = y)").unwrap();
    validate_src("({x: y} = y)").unwrap();
}

#[test]
fn test_calls() {
    validate_src("foo()").unwrap();
    validate_src("foo(1,2)").unwrap();
    validate_src("foo(1,2,...bar)").unwrap();
}

#[test]
fn test_import_export() {
    validate_src("import * as foo from 'foo'").unwrap();
    validate_src("import {x} from 'foo'").unwrap();
    validate_src("import {x as y} from 'foo'").unwrap();
    validate_src("export var x;").unwrap();
    validate_src("export * as x from 'foo';").unwrap();
}

#[test]
fn test_jsx() {
    validate_src_jsx("<foo />").unwrap();
    validate_src_jsx("<foo bar={1} />").unwrap();
    validate_src_jsx("<foo bar='abc' />").unwrap();
    validate_src_jsx("<foo><bar></bar></foo>").unwrap();
    validate_src_jsx("<foo><bar>abc</bar></foo>").unwrap();
}

#[test]
fn test_flow() {
    validate_src_flow("function foo(a: number): number { return a; }").unwrap();
    validate_src_flow("function foo(a: number): number %checks { return a; }").unwrap();
    validate_src_flow("function foo(a: number): number %checks(a[1]) { return a; }").unwrap();
    validate_src_flow(
        "class A<T> extends B<T> {
            foo: number;
            +bar: string;
    }",
    )
    .unwrap();
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use juno::ast::*;
use juno::gen_js;
use juno::hparser;
use juno::sourcemap::merge_sourcemaps;

fn do_gen<'ast>(ctx: &mut Context<'ast>, node: &NodeRc, pretty: gen_js::Pretty) -> String {
    use juno::gen_js::*;
    let mut out: Vec<u8> = vec![];
    generate(
        &mut out,
        ctx,
        node,
        gen_js::Opt {
            pretty,
            ..Default::default()
        },
    )
    .unwrap();
    String::from_utf8(out).expect("Invalid UTF-8 output in test")
}

fn test_roundtrip_with_flags(flags: hparser::ParserFlags, src1: &str) {
    use juno::ast::*;

    let mut ctx = Context::new();

    for pretty in &[gen_js::Pretty::Yes, gen_js::Pretty::No] {
        let ast1 = hparser::parse_with_flags(flags, src1, &mut ctx).unwrap();
        let mut dump: Vec<u8> = vec![];
        dump_json(&mut dump, &mut ctx, &ast1, juno::ast::Pretty::Yes).unwrap();
        let ast1_json = String::from_utf8(dump).expect("Invalid UTF-8 output in test");

        let src2 = do_gen(&mut ctx, &ast1, *pretty);
        let ast2 = hparser::parse_with_flags(flags, &src2, &mut ctx).unwrap_or_else(|_| {
            panic!(
                "Invalid JS generated: Pretty={:?}\nOriginal Source:\n{}\nGenerated Source:\n{}",
                pretty, &src1, &src2,
            )
        });
        let mut dump: Vec<u8> = vec![];
        dump_json(&mut dump, &mut ctx, &ast2, juno::ast::Pretty::Yes).unwrap();
        let ast2_json = String::from_utf8(dump).expect("Invalid UTF-8 output in test");

        assert_eq!(
            ast1_json, ast2_json,
            "AST mismatch: Pretty={:?}\nOriginal Source:\n{}\nGenerated Source:\n{}",
            pretty, &src1, &src2
        );
    }
}

fn test_roundtrip(src1: &str) {
    test_roundtrip_with_flags(Default::default(), src1)
}

fn test_roundtrip_flow(src1: &str) {
    test_roundtrip_with_flags(
        hparser::ParserFlags {
            strict_mode: false,
            enable_jsx: false,
            dialect: hparser::ParserDialect::Flow,
            store_doc_block: false,
        },
        src1,
    );
}

fn test_roundtrip_jsx(src1: &str) {
    test_roundtrip_with_flags(
        hparser::ParserFlags {
            strict_mode: false,
            enable_jsx: true,
            dialect: hparser::ParserDialect::JavaScript,
            store_doc_block: false,
        },
        src1,
    )
}

#[test]
fn test_literals() {
    let mut ctx = Context::new();
    let string = {
        let gc = ctx.lock();
        NodeRc::from_node(
            &gc,
            builder::StringLiteral::build_template(
                &gc,
                template::StringLiteral {
                    metadata: Default::default(),
                    value: gc.atom_u16(vec!['A' as u16, 0x1234u16, '\t' as u16]),
                },
            ),
        )
    };
    assert_eq!(
        do_gen(&mut ctx, &string, gen_js::Pretty::Yes).trim(),
        r#"'A\u1234\t'"#,
    );

    test_roundtrip("1");
    test_roundtrip("1n");
    test_roundtrip("11298379123162378326187361n");
    test_roundtrip("\"abc\"");
    test_roundtrip(r#" "\ud800" "#);
    test_roundtrip(r#" "\ud83d\udcd5" "#);
    test_roundtrip(r#" "\u060b" "#);
    test_roundtrip("true");
    test_roundtrip("false");
    test_roundtrip("null");
    test_roundtrip("undefined");

    test_roundtrip("/abc/");
    test_roundtrip("/abc/gi");
    test_roundtrip("/abc/gi");
    test_roundtrip(r#"/ab\/cd/gi"#);
    test_roundtrip("/ðŸ˜¹/");

    test_roundtrip(r#" `abc` "#);
    test_roundtrip(r#" `abc\ndef` "#);
    test_roundtrip(
        r#" `abc
        def` "#,
    );
    test_roundtrip(r#" `abc \ud800 def` "#);
    test_roundtrip(r#" `abc \ud800 def` "#);
    test_roundtrip(r#" `\ud83d\udcd5` "#);
    test_roundtrip(r#" `escape backtick: \` should work` "#);
    test_roundtrip(r#" `ðŸ˜¹` "#);
}

#[test]
fn test_identifier() {
    test_roundtrip("foo");
    test_roundtrip("class C { #foo() {} }");
}

#[test]
fn test_binop() {
    test_roundtrip("null + null");
    test_roundtrip("1 + 1");
    test_roundtrip("1 * 2 + (3 + 4)");
    test_roundtrip("1 ** 2 ** 3 ** 4");
    test_roundtrip("1 in 2 + (2 - 4) / 3");
    test_roundtrip("1 instanceof 2 + (2 - 4) / 3");
}

#[test]
fn test_conditional() {
    test_roundtrip("a ? b : c");
    test_roundtrip("a ? b : c ? d : e");
    test_roundtrip("(a ? b : c) ? d : e");
    test_roundtrip("a ? b : (c ? d : e)");
    test_roundtrip("a?.3:.4");
}

#[test]
fn test_vars() {
    test_roundtrip("var x=3;");
    test_roundtrip("var x=3, y=4;");
}

#[test]
fn test_functions() {
    test_roundtrip("function foo() {}");
    test_roundtrip("function foo(x, y) {}");
    test_roundtrip("function foo(x, y=3) {}");
    test_roundtrip("function foo([x, y], {z}) {}");
    test_roundtrip("function foo([x, y] = [1,2], {z:q}, {w = 1}) {}");
    test_roundtrip("function foo() { return this; }");
    test_roundtrip("function *foo() {}");
    test_roundtrip("function *foo() { yield 1; }");
    test_roundtrip("function *foo() { yield* f(); }");
    test_roundtrip("async function foo() {}");
    test_roundtrip("async function foo() { await f(); }");
    test_roundtrip("async function *foo() {}");
    test_roundtrip("async function *foo() { await f(); yield 1; }");
    test_roundtrip("x => 3");
    test_roundtrip("(x) => 3");
    test_roundtrip("(x,y) => 3");
    test_roundtrip("x => {3}");
    test_roundtrip("x => ({y: 10})");
    test_roundtrip("x => ({y: 10}[z])");
    test_roundtrip("async x => {3}");
    test_roundtrip("async (x,y) => {3}");
    test_roundtrip("(x => 1) + (y => 1)");
    test_roundtrip("x = y => 1");
    test_roundtrip("x = (y => 1)");
    test_roundtrip("x = (({a, b}) => 1)");
    test_roundtrip_flow("var x = (): (number=>string) => 1");
    test_roundtrip(
        "function foo() {
        return (y => 1);
    }",
    );
    test_roundtrip(
        "function* foo() {
        yield y => 1;
    }",
    );
}

#[test]
fn test_calls() {
    test_roundtrip("f();");
    test_roundtrip("f(1);");
    test_roundtrip("f(1, 2);");
    test_roundtrip("f(1, (2,3), 4);");
    test_roundtrip("(f?.(1, 2))(3);");
    test_roundtrip("f?.(1, 2)?.(3)(5);");
    test_roundtrip("f(...x)");
    test_roundtrip("new f();");
    test_roundtrip("new f(1);");
    test_roundtrip("new f(...x)");
    test_roundtrip("new(a.b);");
    test_roundtrip("new(a.b());");
    test_roundtrip("new(a.b())();");
    test_roundtrip("new(a.b())(c);");
    test_roundtrip("new(a?.b())(c);");
    test_roundtrip("new(1 + 2);");
    test_roundtrip("new(fn(foo)[bar])()");
    test_roundtrip("new(fn(foo)[bar])(c)");
    test_roundtrip("new(fn(foo).bar)()");
    test_roundtrip("new(fn(foo).bar)(c)");
    test_roundtrip("import('foo')");
}

#[test]
fn test_statements() {
    test_roundtrip("while (1) {}");
    test_roundtrip("while (1) { fn(); }");
    test_roundtrip("while (1) fn();");
    test_roundtrip("while (1) fn()");
    test_roundtrip("for (;;) { fn(); }");
    test_roundtrip("for (;;) fn();");
    test_roundtrip("for (x;;) { fn(); }");
    test_roundtrip("for (;x;) { fn(); }");
    test_roundtrip("for (;;x) { fn(); }");
    test_roundtrip("for (var x=1;x<10;++x) { fn(); }");
    test_roundtrip("for (x in y) { fn(); }");
    test_roundtrip("for (var x of y) { fn(); }");
    test_roundtrip(
        "async () => {
            for await (x of y) { fn(); }
        }",
    );
    test_roundtrip("do {fn();} while (1)");
    test_roundtrip("do fn(); while (1)");
    test_roundtrip("do x, y, z; while (1)");
    test_roundtrip("do if (x) y; while (1)");
    test_roundtrip("debugger");
    test_roundtrip("{fn(); fn();}");
    test_roundtrip("for (;;) { break; }");
    test_roundtrip("for (;;) { continue; }");
    test_roundtrip("function f() { return; }");
    test_roundtrip("function f() { return 3; }");
    test_roundtrip(
        "switch(x) {
            case 1:
                break;
            case 2:
            case 3:
                break;
            default:
                break;
        }",
    );
    test_roundtrip("a: var x = 3;");
    test_roundtrip(
        "try {
            fn();
        } catch {
            fn();
        }",
    );
    test_roundtrip(
        "try {
            fn();
        } catch (e) {
            fn();
        }",
    );
    test_roundtrip(
        "try {
            fn();
        } catch (e) {
            fn();
        } finally {
            fn();
        }",
    );
    test_roundtrip("if (x) {fn();}");
    test_roundtrip("if (x) {fn();} else {fn();}");
    test_roundtrip("if (x) fn(); else fn();");
    test_roundtrip(
        "if (x)
          try { } catch (e) { }
        else
          fn();",
    );
}

#[test]
fn test_logical() {
    test_roundtrip("a && b || c");
    test_roundtrip("a || b && c");
    test_roundtrip("(a || b) && c");
    test_roundtrip("(a || b) ?? c");
    test_roundtrip("(a ?? b) || c");
}

#[test]
fn test_sequences() {
    test_roundtrip("var x = (1, 2, 3);");
    test_roundtrip("foo((1, 2, 3), 4);");
}

#[test]
fn test_objects() {
    test_roundtrip("({ })");
    test_roundtrip(
        "({
            a: 1,
            [x]: 1,
            fn() {},
            b,
            ...from,
        })",
    );
    test_roundtrip_flow(
        "({
            foo<T>() {},
        })",
    );
}

#[test]
fn test_arrays() {
    test_roundtrip("([])");
    test_roundtrip("var x = [, 1, , 3]");
    test_roundtrip("var x = [1, 2, 3, ...from]");
    test_roundtrip("var x = [1, 2, 3, ...from, 4, 5, 6]");
}

#[test]
fn test_assignment() {
    test_roundtrip("x = 1");
    test_roundtrip("x = y = 1");
    test_roundtrip("x += 1");
    test_roundtrip("x -= 1");
    test_roundtrip("x *= 1");
    test_roundtrip("x /= 1");
    test_roundtrip("x **= 1");
    test_roundtrip("x |= 1");
    test_roundtrip("x &= 1");
    test_roundtrip("x ||= 1");
    test_roundtrip("x &&= 1");
    test_roundtrip("x ??= 1");
    test_roundtrip("foo()[1] = 1");
    test_roundtrip("a = b && c");
    test_roundtrip("(a = b) && c");
    test_roundtrip("a && b = c");
    test_roundtrip("a && (b = c)");
    test_roundtrip("var {x: {y: [{z}]}} = foo;");
    test_roundtrip("({x: {y: [{z}]}} = foo);");
    test_roundtrip("var [x, y] = foo;");
    test_roundtrip("([x, y] = foo);");
}

#[test]
fn test_unary() {
    test_roundtrip("+x");
    test_roundtrip("-x");
    test_roundtrip("!x");
    test_roundtrip("~x");
    test_roundtrip("-(-x)");
    test_roundtrip("-(-5)");
    test_roundtrip("--x");
    test_roundtrip("x--");
    test_roundtrip("++x");
    test_roundtrip("x++");
    test_roundtrip("+!-x");
    test_roundtrip("delete x");
    test_roundtrip("typeof x");
}

#[test]
fn test_update() {
    test_roundtrip("x++");
    test_roundtrip("x--");
    test_roundtrip("++x");
    test_roundtrip("--x");
    test_roundtrip("--(-x)");
    test_roundtrip("+x++");
}

#[test]
fn test_members() {
    test_roundtrip("a.b");
    test_roundtrip("a.b.c");
    test_roundtrip("a?.b");
    test_roundtrip("a?.[b]");
    test_roundtrip("(a?.b).c");
    test_roundtrip("a?.b().c");
    test_roundtrip("(a?.b()).c");
    test_roundtrip("a?.().b");
    test_roundtrip("a?.().b");
    test_roundtrip("a?.b?.c?.()");
    test_roundtrip("(a?.b?.c?.()).d");
    test_roundtrip("(a?.b?.c?.())?.d");
    test_roundtrip("(a?.b?.c?.())(d)");
    test_roundtrip("(a?.b?.c?.())?.(d)");
    test_roundtrip("class C { constructor() { new.target; } }");
    test_roundtrip("50..toString()");
    test_roundtrip("1.5.toString()");
    test_roundtrip("1e100.toString()");
    test_roundtrip("-1e100.toString()");
    test_roundtrip("0x10293.toString()");
}

#[test]
fn test_classes() {
    test_roundtrip("class C {}");
    test_roundtrip("class C extends D {}");
    test_roundtrip(
        "class C extends D {
            prop1;
            #prop2;
            constructor() {}
            a() {}
            #b() {}
            c(x, y) {}
            static d() {}
        }",
    );
    test_roundtrip_flow(
        "class C<T> extends D<T> {
            prop1: ?number = null;
            +prop2: number;
            -prop3;
            declare prop4;
            #prop5;
            #prop5: ?number = null;
            declare +prop6;
            static +prop7;
            static +[prop8];
            declare static +prop9;
            foo<T>() {}
        }",
    );
    test_roundtrip(
        "var cls = (class C extends D {
            prop1;
            #prop2;
            constructor() {}
            a() {}
            #b() {}
            c(x, y) {}
            static d() {}
            get e() {}
            set e(v) {}
            ;
        })",
    );
}

#[test]
fn test_import() {
    test_roundtrip("import x from 'foo'");
    test_roundtrip("import x, {y} from 'foo'");
    test_roundtrip("import * as Foo from 'foo'");
    test_roundtrip("import x, {y as z, a as b} from 'foo'");
    test_roundtrip("import {a, b, c} from 'foo'");
    test_roundtrip("import 'foo';");
    test_roundtrip("import 'foo' assert {kind: 'json'};");
    test_roundtrip(
        "
        import 'foo';
        import 'bar';
        ",
    );
}

#[test]
fn test_export() {
    test_roundtrip("export var x = 3;");
    test_roundtrip("export function foo() {}");
    test_roundtrip("export default function foo() {}");
    test_roundtrip("export {x as y};");
    test_roundtrip("export * from 'foo';");
    test_roundtrip_flow("export type Foo = number;");
    test_roundtrip_flow("export type { x as y } from 'foo';");
}

#[test]
fn test_types() {
    test_roundtrip_flow("number | boolean & string");
    test_roundtrip_flow("type A = number");
    test_roundtrip_flow("type A = ?number");
    test_roundtrip_flow("type A = ?(number | string)");
    test_roundtrip_flow("type A = string");
    test_roundtrip_flow("type A = \"foo\"");
    test_roundtrip_flow("type A = 'foo'");
    test_roundtrip_flow("type A = 3");
    test_roundtrip_flow("type A = 3n");
    test_roundtrip_flow("type A = boolean");
    test_roundtrip_flow("type A = true | false");
    test_roundtrip_flow("type A = true & false");
    test_roundtrip_flow("type A = (X | Y) & Z");
    test_roundtrip_flow("type A = X | Y & Z");
    test_roundtrip_flow("type A = X<Y, Z>");
    test_roundtrip_flow("type A = X<Y>");
    test_roundtrip_flow("type A<X: Y, Z> = T");
    test_roundtrip_flow("type A = symbol");
    test_roundtrip_flow("type A = mixed");
    test_roundtrip_flow("type A = any");
    test_roundtrip_flow("type A = void");
    test_roundtrip_flow("type A = null");
    test_roundtrip_flow("type A = number => number");
    test_roundtrip_flow("type A = X.Y");
    test_roundtrip_flow("type A = X.Y<Z>");
    test_roundtrip_flow("type A = typeof X");
    test_roundtrip_flow("type A = [number, string]");
    test_roundtrip_flow("type A = []");
    test_roundtrip_flow("type A = number[]");
    test_roundtrip_flow("type A = number[string]");
    test_roundtrip_flow("type A = number?.[string]");
    test_roundtrip_flow("type A = [number, string][]");
    test_roundtrip_flow("type A = (foo: number) => number");
    test_roundtrip_flow("type A = (foo?: number) => number");
    test_roundtrip_flow("type A = (foo?: ?number) => number");
    test_roundtrip_flow("type A = (number, string) => number");
    test_roundtrip_flow("type A = (?number) => number");
    test_roundtrip_flow("type A = ?(number, string) => number");
    test_roundtrip_flow("type A = (this: number, number, string) => number");
    test_roundtrip_flow("interface A { }");
    test_roundtrip_flow("interface A extends B { }");
    test_roundtrip_flow("interface A extends B, C, D { }");
    test_roundtrip_flow("type A = { x: number }");
    test_roundtrip_flow("type A = {| x: number |}");
    test_roundtrip_flow(
        "
        type A = {
            a?: number,
            b: ?string,
            +[c]: string,
            (d?: number): number;
            [[e]]: number,
            [[f]]?: number,
            [[g]](a: T): number,
            ...h,
            static (i?: number): number;
            +proto: number,
            ...
        };
        ",
    );
}

#[test]
fn test_declare() {
    test_roundtrip_flow("declare function foo(): number;");
    test_roundtrip_flow("declare var x : number;");
    test_roundtrip_flow("declare export var x: number;");
    test_roundtrip_flow("declare opaque type x;");
    test_roundtrip_flow("declare export opaque type x: y;");
    test_roundtrip_flow("declare type x = number;");
    test_roundtrip_flow("declare interface Foo {}");
    test_roundtrip_flow("declare class A extends B {}");
    test_roundtrip_flow("declare class A extends B mixins C, D implements E {}");
    test_roundtrip_flow("declare export class A extends B {}");
    test_roundtrip_flow("declare module A {}");
    test_roundtrip_flow("declare module.exports: number;");
    test_roundtrip_flow("declare export function foo(): number;");
}

#[test]
fn test_enum() {
    test_roundtrip_flow("enum Foo {}");
    test_roundtrip_flow("enum Foo of string {A = 'A', B = 'B'}");
    test_roundtrip_flow("enum Foo of string {A, B, C}");
    test_roundtrip_flow("enum Foo of string {A = 'A', B = 'B', ...}");
    test_roundtrip_flow("enum Foo of number {A = 1}");
    test_roundtrip_flow("enum Foo of boolean {A = true}");
}

#[test]
fn test_typecast() {
    test_roundtrip_flow("async function foo() { return (x: any); }");
    test_roundtrip_flow("var x = (y: number | number => string)");
}

#[test]
fn test_predicate() {
    test_roundtrip_flow("function foo(): %checks {}");
    test_roundtrip_flow("function foo(): number %checks {}");
    test_roundtrip_flow("function foo(): number %checks(bar) {}");
    test_roundtrip_flow("((x): %checks => 3)");
    test_roundtrip_flow("((x): number %checks => 3)");
    test_roundtrip_flow("((x): number %checks(bar) => 3)");
}

#[test]
fn test_this_param() {
    test_roundtrip_flow("function foo(this: number): number {}");
    test_roundtrip_flow("function foo(this: number, x: number): number {}");
    test_roundtrip_flow("declare function foo(this: number): number;");
    test_roundtrip_flow("declare function foo(this: number, x: number): number;");
}

#[test]
fn test_jsx() {
    test_roundtrip_jsx("<foo />");
    test_roundtrip_jsx("<foo></foo>");
    test_roundtrip_jsx("<foo>abc</foo>");
    test_roundtrip_jsx(
        r#"
        <asdf desc="foo
            bar"
            prop2='foo """ bar'>
            body
        </asdf>
        "#,
    );
    test_roundtrip_jsx("<></>");
    test_roundtrip_jsx(
        "
        <foo>
            <bar x={1} y='3' {...z} />
            abcdef and an emoji: ðŸ˜¹
            > &x1f639;
            end of test text
            <baz.bar />
            <hello:goodbye />
        </foo>
        ",
    );
}

#[test]
fn test_sourcemap() {
    use juno::gen_js::*;
    let mut ctx = Context::new();
    let ast1: NodeRc = hparser::parse(&mut ctx, "function foo() { return 1 }").unwrap();
    let mut out: Vec<u8> = vec![];
    let sourcemap = generate(&mut out, &mut ctx, &ast1, gen_js::Opt::new()).unwrap();
    let string = String::from_utf8(out).expect("Invalid UTF-8 output in test");
    assert_eq!(
        string,
        "function foo() {
  return 1;
}\n"
    );
    assert_eq!(sourcemap.get_token_count(), 4);

    assert_eq!(
        sourcemap.get_token(0).unwrap().get_raw_token(),
        sourcemap::RawToken {
            dst_line: 0,
            dst_col: 0,
            src_line: 0,
            src_col: 0,
            src_id: 0,
            name_id: !0,
        }
    );

    assert_eq!(
        sourcemap.get_token(1).unwrap().get_raw_token(),
        sourcemap::RawToken {
            dst_line: 0,
            dst_col: 9,
            src_line: 0,
            src_col: 9,
            src_id: 0,
            name_id: !0,
        }
    );

    assert_eq!(
        sourcemap.get_token(2).unwrap().get_raw_token(),
        sourcemap::RawToken {
            dst_line: 1,
            dst_col: 2,
            src_line: 0,
            src_col: 17,
            src_id: 0,
            name_id: !0,
        }
    );

    assert_eq!(
        sourcemap.get_token(3).unwrap().get_raw_token(),
        sourcemap::RawToken {
            dst_line: 1,
            dst_col: 9,
            src_line: 0,
            src_col: 24,
            src_id: 0,
            name_id: !0,
        }
    );
}

#[test]
fn test_sourcemap_merged() {
    use juno::gen_js::*;
    let mut context = Context::new();
    let ctx = &mut context;
    // original TS source:
    // type A = number;
    // function foo() { 1; }

    let input_src = r#"function foo() { 1; }"#;
    let input_map = sourcemap::SourceMap::from_slice(
        br#"{
            "version":3,
            "file":"test.js",
            "sourceRoot":"",
            "sources":["test.ts"],
            "names":[],
            "mappings":"AACA,SAAS,GAAG,KAAK,CAAC,CAAC,CAAC,CAAC"
        }"#,
    )
    .unwrap();
    let mut out: Vec<u8> = vec![];
    let node = hparser::parse_with_flags(Default::default(), input_src, ctx).unwrap();
    let output_map = generate(&mut out, ctx, &node, gen_js::Opt::new()).unwrap();
    let output = String::from_utf8(out).expect("Invalid UTF-8 output in test");
    assert_eq!(output, "function foo() {\n  1;\n}\n",);

    let merged = merge_sourcemaps(&input_map, &output_map);

    // Source maps are 0-indexed, so use those for testing purposes.
    let input_token: sourcemap::Token = merged.lookup_token(1, 2).unwrap();
    assert_eq!(input_token.get_source().unwrap(), "test.ts");
    assert_eq!(input_token.get_src(), (1, 17));
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! Garbage-collected Storage structures for AST nodes.

use std::cell::Cell;
use std::cell::UnsafeCell;
use std::hash::Hash;
use std::hash::Hasher;
use std::ops::Deref;
use std::pin::Pin;
use std::ptr::NonNull;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering;

use juno_support::atom_table::Atom;
use juno_support::atom_table::AtomTable;
use juno_support::atom_table::AtomU16;
use juno_support::Deque;
use juno_support::HeapSize;
use libc::c_void;
use memoffset::offset_of;

use crate::Node;
use crate::Path;
use crate::SourceManager;
use crate::Visitor;

/// ID which indicates a `StorageEntry` is free.
const FREE_ENTRY: u32 = 0;

/// A single entry in the heap.
#[derive(Debug)]
struct StorageEntry<'ctx> {
    /// ID of the context to which this entry belongs.
    /// Top bit is used as a mark bit, and flips meaning every time a GC happens.
    /// If this field is `0`, then this entry is free.
    ctx_id_markbit: Cell<u32>,

    /// Refcount of how many [`NodeRc`] point to this node.
    /// Entry may only be freed if this number is `0` and no other entries reference this entry
    /// directly.
    count: Cell<u32>,

    /// Actual node stored in this entry.
    inner: Node<'ctx>,
}

impl<'ctx> StorageEntry<'ctx> {
    unsafe fn from_node<'a>(node: &'a Node<'a>) -> &'a StorageEntry<'a> {
        let inner_offset = offset_of!(StorageEntry, inner) as isize;
        let inner = node as *const Node<'a>;
        &*(inner.offset(-inner_offset) as *const StorageEntry<'a>)
    }

    #[inline]
    fn set_markbit(&self, bit: bool) {
        let id = self.ctx_id_markbit.get();
        if bit {
            self.ctx_id_markbit.set(id | 1 << 31);
        } else {
            self.ctx_id_markbit.set(id & !(1 << 31));
        }
    }

    #[inline]
    fn markbit(&self) -> bool {
        (self.ctx_id_markbit.get() >> 31) != 0
    }

    fn is_free(&self) -> bool {
        self.ctx_id_markbit.get() == FREE_ENTRY
    }
}

/// A single entry in the NodeList storage.
/// These are also immutable from the user's perspective, like `Node`s,
/// but they are temporarily mutated here during construction only, in order to append elements.
#[derive(Debug)]
pub(crate) struct NodeListElement<'ctx> {
    /// ID of the context to which this entry belongs.
    /// Top bit is used as a mark bit, and flips meaning every time a GC happens.
    /// If this field is `0`, then this entry is free.
    ctx_id_markbit: Cell<u32>,

    /// Actual node stored in this entry.
    /// Must not be null, because empty lists are represented as null pointers in the [`NodeList`].
    pub inner: *const Node<'ctx>,

    /// Pointer to the next element in the NodeList.
    /// Stored in a `Cell` to allow for simple appends.
    pub next: Cell<*const NodeListElement<'ctx>>,
}

impl<'ctx> NodeListElement<'ctx> {
    #[inline]
    fn set_markbit(&self, bit: bool) {
        let id = self.ctx_id_markbit.get();
        if bit {
            self.ctx_id_markbit.set(id | 1 << 31);
        } else {
            self.ctx_id_markbit.set(id & !(1 << 31));
        }
    }

    #[inline]
    fn markbit(&self) -> bool {
        (self.ctx_id_markbit.get() >> 31) != 0
    }

    fn is_free(&self) -> bool {
        self.ctx_id_markbit.get() == FREE_ENTRY
    }
}

/// Structure pointed to by `Context` and `NodeRc` to facilitate panicking if there are
/// outstanding `NodeRc` when the `Context` is dropped.
#[derive(Debug)]
struct NodeRcCounter {
    /// ID of the context owning the counter.
    ctx_id: u32,

    /// Number of [`NodeRc`]s allocated in this `Context`.
    /// Must be `0` when `Context` is dropped.
    count: Cell<usize>,
}

/// The storage for AST nodes.
///
/// Can be used to allocate and free nodes.
/// Nodes allocated in one `Context` must not be referenced by another `Context`'s AST.
#[derive(Debug)]
pub struct Context<'ast> {
    /// Unique number used to identify this context.
    id: u32,

    /// List of all the nodes stored in this context.
    /// Each element is a "chunk" of nodes.
    /// None of the chunks are ever resized after allocation.
    nodes: UnsafeCell<Deque<StorageEntry<'ast>>>,

    /// Free list for AST nodes.
    free_nodes: UnsafeCell<Vec<NonNull<StorageEntry<'ast>>>>,

    /// Every `NodeListElement` allocated in this context.
    /// These store the links in the linked lists.
    list_elements: UnsafeCell<Deque<NodeListElement<'ast>>>,

    /// Free list for `NodeListElement`s.
    free_list_elements: UnsafeCell<Vec<NonNull<NodeListElement<'ast>>>>,

    /// `NodeRc` count stored in a `Box` to ensure that `NodeRc`s can also point to it
    /// and decrement the count on drop.
    /// Placed separately to guard against `Context` moving, though relying on that behavior is
    /// technically unsafe.
    noderc_count: Pin<Box<NodeRcCounter>>,

    /// All identifiers are kept here.
    pub atom_table: AtomTable,

    /// Source manager of this context.
    pub source_mgr: SourceManager,

    /// `true` if `1` indicates an entry is marked, `false` if `0` indicates an entry is marked.
    /// Flipped every time GC occurs.
    markbit_marked: bool,

    /// Whether strict mode has been forced.
    strict_mode: bool,

    /// Whether to warn about undefined variables in strict mode functions.
    pub warn_undefined: bool,
}

impl Default for Context<'_> {
    fn default() -> Self {
        Self::new()
    }
}

impl<'ast> Context<'ast> {
    /// Allocate a new `Context` with a new ID.
    pub fn new() -> Self {
        static NEXT_ID: AtomicU32 = AtomicU32::new(FREE_ENTRY + 1);
        let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
        Self {
            id,
            nodes: Default::default(),
            free_nodes: Default::default(),
            list_elements: Default::default(),
            free_list_elements: Default::default(),
            noderc_count: Pin::new(Box::new(NodeRcCounter {
                ctx_id: id,
                count: Cell::new(0),
            })),
            atom_table: Default::default(),
            source_mgr: Default::default(),
            markbit_marked: true,
            strict_mode: false,
            warn_undefined: false,
        }
    }

    /// Acquire a [`GCLock`] on this `Context`.
    /// This is just a more ergonomic way to call `GCLock::new`.
    pub fn lock<'ctx>(&'ctx mut self) -> GCLock<'ast, 'ctx> {
        GCLock::new(self)
    }

    /// Allocate a new `Node` in this `Context`.
    pub(crate) fn alloc<'s>(&'s self, n: Node<'_>) -> &'s Node<'s> {
        let free = unsafe { &mut *self.free_nodes.get() };
        let nodes: &mut Deque<StorageEntry<'ast>> = unsafe { &mut *self.nodes.get() };
        let node = unsafe { std::mem::transmute(n) };
        let entry: &StorageEntry<'ast> = if let Some(mut entry) = free.pop() {
            let entry: &mut StorageEntry<'ast> = unsafe { entry.as_mut() };
            debug_assert!(
                entry.ctx_id_markbit.get() == FREE_ENTRY,
                "Incorrect context ID"
            );
            debug_assert!(entry.count.get() == 0, "Freed entry has pointers to it");
            entry.ctx_id_markbit.set(self.id);
            entry.set_markbit(!self.markbit_marked);
            entry.inner = node;
            entry
        } else {
            let entry: &StorageEntry = nodes.push(StorageEntry {
                ctx_id_markbit: Cell::new(self.id),
                count: Cell::new(0),
                inner: node,
            });
            entry.set_markbit(!self.markbit_marked);
            entry
        };
        // Transmute here to handle the fact that Cell<> is invariant over its type,
        // meaning the lifetime doesn't automatically narrow from `'ast` to `'s`.
        unsafe { std::mem::transmute(&entry.inner) }
    }

    /// Allocate a list element in the context with the provided previous element if it exists.
    /// `prev` will be updated to point to `node` as its next element.
    pub(crate) fn append_list_element<'a>(
        &'a self,
        prev: Option<&'a NodeListElement<'a>>,
        node: &'a Node<'a>,
    ) -> &'a NodeListElement<'a> {
        let elements: &mut Deque<NodeListElement<'ast>> = unsafe { &mut *self.list_elements.get() };
        let free = unsafe { &mut *self.free_list_elements.get() };
        // Transmutation is safe here, because `Node`s can only be allocated through
        // this path and only one GCLock can be made available at a time per thread.
        let node: &'ast Node<'ast> = unsafe { std::mem::transmute(node) };
        let prev: Option<&'ast NodeListElement<'ast>> = unsafe { std::mem::transmute(prev) };
        let entry = if let Some(mut entry) = free.pop() {
            let entry: &mut NodeListElement<'ast> = unsafe { entry.as_mut() };
            debug_assert!(
                entry.ctx_id_markbit.get() == FREE_ENTRY,
                "Incorrect context ID"
            );
            entry.ctx_id_markbit.set(self.id);
            entry.set_markbit(!self.markbit_marked);
            entry.inner = node;
            entry.next.set(std::ptr::null());
            if let Some(prev) = prev {
                prev.next.set(entry as *const _);
            }
            entry
        } else {
            let entry = elements.push(NodeListElement {
                ctx_id_markbit: Cell::new(self.id),
                inner: node,
                next: Cell::new(std::ptr::null()),
            });
            entry.set_markbit(!self.markbit_marked);
            if let Some(prev) = prev {
                prev.next.set(entry as *const _);
            }
            entry
        };
        debug_assert!(!entry.is_free(), "Entry must not be free");
        // Transmute here to handle the fact that Cell<> is invariant over its type,
        // meaning the lifetime doesn't automatically narrow from `'ast` to `'s`.
        unsafe { std::mem::transmute(entry) }
    }

    /// Return the atom table.
    pub fn atom_table(&self) -> &AtomTable {
        &self.atom_table
    }

    /// Add a string to the identifier table.
    #[inline]
    pub fn atom<V: Into<String> + AsRef<str>>(&self, value: V) -> Atom {
        self.atom_table.atom(value)
    }

    /// Add a string to the identifier table.
    #[inline]
    pub fn atom_u16<V: Into<Vec<u16>> + AsRef<[u16]>>(&self, value: V) -> AtomU16 {
        self.atom_table.atom_u16(value)
    }

    /// Obtain the contents of an atom from the atom table.
    #[inline]
    pub fn str(&self, index: Atom) -> &str {
        self.atom_table.str(index)
    }

    /// Obtain the contents of an atom from the atom table.
    #[inline]
    pub fn str_u16(&self, index: AtomU16) -> &[u16] {
        self.atom_table.str_u16(index)
    }

    /// Return an immutable reference to SourceManager
    pub fn sm(&self) -> &SourceManager {
        &self.source_mgr
    }

    /// Return a mutable reference to SourceManager
    pub fn sm_mut(&mut self) -> &mut SourceManager {
        &mut self.source_mgr
    }

    /// Return true if strict mode has been forced globally.
    pub fn strict_mode(&self) -> bool {
        self.strict_mode
    }

    /// Enable strict mode. Note that it cannot be unset.
    pub fn enable_strict_mode(&mut self) {
        self.strict_mode = true;
    }

    pub fn gc(&mut self) {
        let nodes = unsafe { &mut *self.nodes.get() };
        let free_nodes = unsafe { &mut *self.free_nodes.get() };

        let list_elements = unsafe { &mut *self.list_elements.get() };
        let free_list_elements = unsafe { &mut *self.free_list_elements.get() };

        {
            // Begin by collecting all the roots: entries with non-zero refcount.
            let mut roots: Vec<&StorageEntry> = vec![];
            for entry in nodes.iter() {
                if entry.is_free() {
                    continue;
                }
                debug_assert!(
                    entry.markbit() != self.markbit_marked,
                    "Entry marked before start of GC: \
                        {:?}\nentry.markbit()={}\nmarkbit_marked={}",
                    &entry,
                    entry.markbit(),
                    self.markbit_marked,
                );
                if entry.count.get() > 0 {
                    // Transmuting the lifetime here because we have to store the roots from
                    // across accesses to `nodes`, meaning we must translate
                    // from `'ast` to the lifetime of this scope.
                    roots.push(unsafe { std::mem::transmute(entry) });
                }
            }

            struct Marker {
                markbit_marked: bool,
            }

            impl<'gc> Visitor<'gc> for Marker {
                fn call(
                    &mut self,
                    gc: &'gc GCLock,
                    node: &'gc Node<'gc>,
                    _path: Option<Path<'gc>>,
                ) {
                    let entry = unsafe { StorageEntry::from_node(node) };
                    if entry.markbit() == self.markbit_marked {
                        // Stop visiting early if we've already marked this part,
                        // because we must have also marked all the children.
                        return;
                    }
                    entry.set_markbit(self.markbit_marked);
                    node.mark_lists(gc, |elem| {
                        elem.set_markbit(self.markbit_marked);
                    });
                    node.visit_children(gc, self);
                }
            }

            // Use a visitor to mark every node reachable from roots.
            let mut marker = Marker {
                markbit_marked: self.markbit_marked,
            };
            {
                let gc: GCLock = GCLock::new(self);
                for root in roots {
                    root.inner.visit(&gc, &mut marker, None);
                }
            }
        }

        for entry in nodes.iter_mut() {
            if entry.is_free() {
                // Skip free entries.
                continue;
            }
            if entry.count.get() > 0 {
                // Keep referenced entries alive.
                continue;
            }
            if entry.markbit() == self.markbit_marked {
                // Keep marked entries alive.
                continue;
            }
            // Passed all checks, this entry is free.
            entry.ctx_id_markbit.set(FREE_ENTRY);
            free_nodes.push(unsafe { NonNull::new_unchecked(entry as *mut StorageEntry) });
        }

        for element in list_elements.iter_mut() {
            if element.is_free() {
                // Skip free entries.
                continue;
            }
            if element.markbit() == self.markbit_marked {
                // Keep marked entries alive.
                continue;
            }
            // Passed all checks, this element is free.
            element.ctx_id_markbit.set(FREE_ENTRY);
            free_list_elements
                .push(unsafe { NonNull::new_unchecked(element as *mut NodeListElement) });
        }

        self.markbit_marked = !self.markbit_marked;
    }

    /// Returns the number of node slots which have been allocated.
    /// Includes nodes currently in use as well as nodes in the free list.
    pub fn num_nodes(&self) -> usize {
        let nodes = unsafe { &*self.nodes.get() };
        nodes.len()
    }

    /// Returns the approximate size of just the AST storages in bytes.
    /// Includes the allocated nodes, lists, as well as free lists for both.
    pub fn storage_size(&self) -> usize {
        let nodes = unsafe { &*self.nodes.get() };
        let free_nodes = unsafe { &*self.free_nodes.get() };
        let list_elements = unsafe { &*self.list_elements.get() };
        let free_list_elements = unsafe { &*self.free_list_elements.get() };
        let mut result = 0;
        result += nodes.heap_size();
        result += free_nodes.heap_size();
        result += list_elements.heap_size();
        result += free_list_elements.heap_size();
        result
    }
}

impl HeapSize for Context<'_> {
    fn heap_size(&self) -> usize {
        let nodes = unsafe { &*self.nodes.get() };
        let free_nodes = unsafe { &*self.free_nodes.get() };
        let list_elements = unsafe { &*self.list_elements.get() };
        let free_list_elements = unsafe { &*self.free_list_elements.get() };
        let mut result = 0;
        result += nodes.heap_size();
        result += free_nodes.heap_size();
        result += list_elements.heap_size();
        result += free_list_elements.heap_size();
        result += std::mem::size_of::<NodeRcCounter>();
        result += self.atom_table.heap_size();
        result += self.source_mgr.heap_size();
        result
    }
}

impl Drop for Context<'_> {
    /// Ensure that there are no outstanding `NodeRc`s into this `Context` which will be
    /// invalidated once it is dropped.
    ///
    /// # Panics
    ///
    /// Will panic if there are any `NodeRc`s stored when this `Context` is dropped.
    fn drop(&mut self) {
        if self.noderc_count.count.get() > 0 {
            #[cfg(debug_assertions)]
            {
                // In debug mode, provide more information on which node was leaked.
                let nodes = unsafe { &*self.nodes.get() };
                for entry in nodes.iter() {
                    assert!(
                        entry.count.get() == 0,
                        "NodeRc must not outlive Context: {:#?}\n",
                        &entry.inner
                    );
                }
            }
            // In release mode, just panic immediately.
            panic!("NodeRc must not outlive Context");
        }
    }
}

thread_local! {
    /// Whether there exists a `GCLock` on the current thread.
    static GCLOCK_IN_USE: Cell<bool> = Cell::new(false);
}

/// A way to view the [`Context`].
///
/// Provides the user the ability to create new nodes and dereference [`NodeRc`].
///
/// **At most one is allowed to be active in any thread at any time.**
/// This is to ensure no `&Node` can be shared between `Context`s.
pub struct GCLock<'ast, 'ctx> {
    ctx: &'ctx mut Context<'ast>,
}

impl Drop for GCLock<'_, '_> {
    fn drop(&mut self) {
        GCLOCK_IN_USE.with(|flag| {
            flag.set(false);
        });
    }
}

impl<'ast, 'ctx> GCLock<'ast, 'ctx> {
    /// # Panics
    ///
    /// Will panic if there is already an active `GCLock` on this thread.
    pub fn new(ctx: &'ctx mut Context<'ast>) -> Self {
        GCLOCK_IN_USE.with(|flag| {
            if flag.get() {
                panic!("Attempt to create multiple GCLocks in a single thread");
            }
            flag.set(true);
        });
        GCLock { ctx }
    }

    /// Allocate a node in the `ctx`.
    #[inline]
    pub(crate) fn alloc<'s>(&'s self, n: Node<'s>) -> &'s Node<'s> {
        self.ctx.alloc(n)
    }

    /// Append `node` to the `prev` element if provided, else create the element as the first
    /// element in the `NodeList`.
    #[inline]
    pub(crate) fn append_list_element<'s>(
        &'s self,
        prev: Option<&'s NodeListElement<'s>>,
        n: &'s Node<'s>,
    ) -> &'s NodeListElement<'s> {
        self.ctx.append_list_element(prev, n)
    }

    /// Return a reference to the owning Context.
    pub fn ctx(&self) -> &Context<'ast> {
        self.ctx
    }

    /// Add a string to the identifier table.
    #[inline]
    pub fn atom<V: Into<String> + AsRef<str>>(&self, value: V) -> Atom {
        self.ctx.atom(value)
    }

    /// Add a string to the identifier table.
    #[inline]
    pub fn atom_u16<V: Into<Vec<u16>> + AsRef<[u16]>>(&self, value: V) -> AtomU16 {
        self.ctx.atom_u16(value)
    }

    /// Obtain the contents of an atom from the atom table.
    #[inline]
    pub fn str(&self, index: Atom) -> &str {
        self.ctx.str(index)
    }

    /// Obtain the contents of an atom from the atom table.
    #[inline]
    pub fn str_u16(&self, index: AtomU16) -> &[u16] {
        self.ctx.str_u16(index)
    }

    /// Return an immutable reference to SourceManager.
    #[inline]
    pub fn sm(&self) -> &SourceManager {
        self.ctx.sm()
    }

    /// Return a mutable reference to SourceManager.
    #[inline]
    pub fn sm_mut(&mut self) -> &mut SourceManager {
        self.ctx.sm_mut()
    }
}

/// A wrapper around Node&, with "shallow" hashing and equality, suitable for
/// hash tables.
#[derive(Debug, Copy, Clone)]
pub struct NodePtr<'gc>(pub &'gc Node<'gc>);

impl<'gc> NodePtr<'gc> {
    pub fn from_node(node: &'gc Node<'gc>) -> Self {
        Self(node)
    }
}

impl<'gc> PartialEq for NodePtr<'gc> {
    fn eq(&self, other: &Self) -> bool {
        std::ptr::eq(self.0, other.0)
    }
}

impl Eq for NodePtr<'_> {}

impl Hash for NodePtr<'_> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        (self.0 as *const Node).hash(state)
    }
}

impl<'gc> Deref for NodePtr<'gc> {
    type Target = Node<'gc>;
    fn deref(&self) -> &'gc Self::Target {
        self.0
    }
}

impl<'gc> AsRef<Node<'gc>> for NodePtr<'gc> {
    fn as_ref(&self) -> &'gc Node<'gc> {
        self.0
    }
}

impl<'gc> From<&'gc Node<'gc>> for NodePtr<'gc> {
    fn from(node: &'gc Node<'gc>) -> Self {
        NodePtr(node)
    }
}

/// Reference counted pointer to a [`Node`] in any [`Context`].
///
/// It can be used to keep references to `Node`s outside of the lifetime of a [`GCLock`],
/// but the only way to derefence and inspect the `Node` is to use a `GCLock`.
#[derive(Debug, Eq)]
pub struct NodeRc {
    /// The `NodeRcCounter` counting for the `Context` to which this belongs.
    counter: NonNull<NodeRcCounter>,

    /// Pointer to the `StorageEntry` containing the `Node`.
    /// Stored as `c_void` to avoid specifying lifetimes, as dereferencing is checked manually.
    entry: NonNull<c_void>,
}

impl Hash for NodeRc {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.entry.hash(state)
    }
}

impl PartialEq for NodeRc {
    fn eq(&self, other: &Self) -> bool {
        self.entry == other.entry
    }
}

impl Drop for NodeRc {
    fn drop(&mut self) {
        let entry = unsafe { self.entry().as_mut() };
        let c = entry.count.get();
        debug_assert!(c > 0);
        entry.count.set(c - 1);

        let noderc_count = unsafe { self.counter.as_mut() };
        let c = noderc_count.count.get();
        debug_assert!(c > 0);
        noderc_count.count.set(c - 1);
    }
}

impl Clone for NodeRc {
    /// Cloning a `NodeRc` increments refcounts on the entry and the context.
    fn clone(&self) -> Self {
        let mut cloned = NodeRc { ..*self };

        let entry = unsafe { cloned.entry().as_mut() };
        let c = entry.count.get();
        entry.count.set(c + 1);

        let noderc_count = unsafe { cloned.counter.as_mut() };
        let c = noderc_count.count.get();
        noderc_count.count.set(c + 1);

        cloned
    }
}

impl NodeRc {
    /// Turn a node reference into a `NodeRc` for storage outside `GCLock`.
    pub fn from_node<'gc>(gc: &'gc GCLock, node: &'gc Node<'gc>) -> NodeRc {
        let inner_offset = offset_of!(StorageEntry, inner) as isize;
        let inner = node as *const Node<'gc>;
        unsafe {
            let entry: &mut StorageEntry = &mut *(inner.offset(-inner_offset) as *mut StorageEntry);
            Self::from_entry(gc, entry)
        }
    }

    /// Return the actual `Node` that `self` points to.
    ///
    /// # Panics
    ///
    /// Will panic if `gc` is not for the same context as this `NodeRc` was created in.
    pub fn node<'gc>(&'_ self, gc: &'gc GCLock<'_, '_>) -> &'gc Node {
        unsafe {
            assert_eq!(
                self.counter.as_ref().ctx_id,
                gc.ctx.id,
                "Attempt to derefence NodeRc allocated context {} in context {}",
                self.counter.as_ref().ctx_id,
                gc.ctx.id
            );
            &self.entry().as_ref().inner
        }
    }

    /// Get the pointer to the `StorageEntry`.
    unsafe fn entry(&self) -> NonNull<StorageEntry> {
        let outer = self.entry.as_ptr() as *mut StorageEntry;
        NonNull::new_unchecked(outer)
    }

    unsafe fn from_entry(gc: &GCLock, entry: &StorageEntry<'_>) -> NodeRc {
        let c = entry.count.get();
        entry.count.set(c + 1);

        let c = gc.ctx.noderc_count.count.get();
        gc.ctx.noderc_count.count.set(c + 1);

        NodeRc {
            counter: NonNull::new_unchecked(gc.ctx.noderc_count.as_ref().get_ref()
                as *const NodeRcCounter
                as *mut NodeRcCounter),
            entry: NonNull::new_unchecked(entry as *const StorageEntry as *mut c_void),
        }
    }
}