Find AstNode/node declaration by ReferenceId from Traverse

[Hookyns]

Hello guys 👋,

Is there currently a way to access AstNodes from TraverseCtx/Traverse, or otherwise retrieve the node declaration corresponding to a ReferenceId?

I'm trying to implement a transformation based on semantic information. For example, I want to find all call expressions where the callee is a specific function (imported from a specific module) and modify its arguments based on type arguments.

Is there a way to achieve this?

One idea I had was to store AstNodes in my visitor struct, but due to Rust's borrowing rules, that doesn't seem viable. The compiler complains about borrowing the Program mutably while it is still borrowed immutably via Semantic. Maybe there's a Rust solution I'm missing? 🤔

Here’s a simplified version of the code:

struct MyVisitor<'a> {
    nodes: AstNodes<'a>,
}

impl<'a> Traverse<'a> for MyVisitor<'a> {
    fn enter_call_expression(&mut self, node: &mut CallExpression<'a>, ctx: &mut TraverseCtx<'a>) {
        match &node.callee {
            Expression::Identifier(identifier) => {
                if identifier.name == "getType" {
                    if let Some(reference_id) = identifier.reference_id.get() {
                        let symbol_table = ctx.scoping.scoping();
                        let symbol_id = symbol_table.get_reference(reference_id).symbol_id();
                        
                        // TODO: Get AstNode
                        // semantic.nodes().get_node(symbol_table.symbol_declaration(symbol_id))
                    }
                }
            }
            _ => {}
        }
    }
}

pub fn test(
    file_path: String,
    source_code: String,
    transformer_context: TransformerContext,
) {
    let source_code = Arc::new(source_code);
    let path = Path::new(&file_path);
    let source_type = SourceType::from_path(path).unwrap();
    
    let allocator = Allocator::default();
    let parser_ret = Parser::new(&allocator, &source_code, source_type).parse();
    let mut program = parser_ret.program;
    let semantic = SemanticBuilder::new().with_check_syntax_error(true).build(&program);

    let (scoping, nodes) = semantic.semantic.into_scoping_and_nodes();

    let mut my_visitor = MyVisitor {
        nodes
    };

    traverse_mut(&mut my_visitor, &allocator, &mut program, scoping);
    //                                          ^^^^
    // Rust error here: cannot borrow `program` as mutable because it is also borrowed as immutable
}

The key issue is that Semantic::build borrows program immutably, and that borrow lives on in nodes, making it impossible to later mutate program during traversal.

Any guidance or workarounds would be greatly appreciated!

---

It would be nice to have AstNodes in TraverseCtx.

[overlookmotel]

I'm afraid it's not possible to have AstNodes in traverse. The 2 are fundamentally incompatible, because AstNodes contains & references to all AST nodes, and traverse provides &mut references to nodes (so you can mutate the AST).

Traverse pulls some tricks to allow reading back up the tree by carefully ensuring you can never obtain a & and &mut ref to same node at the same time. But it's not possible to generalize to an ability to get hold of any node at any time, the way AstNodes allows. This would be a violation of Rust's aliasing rules.

So I'm sorry to say, this isn't going to be possible.

Out of interest, can I ask what is it that you're trying to do? What's the purpose for which you need to get hold of the symbol's declaration node?

[Hookyns]

Thanks for reply. 🙂

Well, I'm creating runtime reflection for TS. See rttist.org, formal tst-reflect.

I have to find all call expressions of getType fnc, e.g. getType<SomeType>(). I have to know it is getType from rttist package and also what is SomeType - where it is declared, if it is imported then from which module? I have to build unique identifier for each type and it must be deterministic so based on path(/scopes) in the module.

I'll post more details tomorrow; I'm writing from phone from bed. 😅

[Hookyns]

I'll probably have to do two passes. First using the Visit and then Traverse.

I have to transform code like this:

index.ts

import { getType } from 'rttist';

class Foo {
    private foo: string;

    static Bar = class {
        baz: string;
    }
}

function func<TParam>(param: string) {
    getType<TParam>();
}

console.log(getType<Foo>());
console.log(getType<typeof Foo.Bar>());
console.log(getType<typeof func>());

const foo = new Foo();
console.log(getType(foo));

to:

import { getType } from 'rttist';

class Foo {
    private foo: string;

    static Bar = class {
        baz: string;

        // Added static block <--------------
        static {
            this.prototype['[[type]]'] = '@@this/index::Foo::Bar';
        }
    }

    // Added static block <--------------
    static {
        Foo.prototype['[[type]]'] = '@@this/index::Foo';
    }
}

function func<TParam>(param: string) {
    // Added access to callsite <--------------
    const [__tp$TParam] = Rttist.resFnCs$(func);
    {
        Rttist.cs$(getType, undefined, [__tp$TParam]);
    }
}
func.prototype['[[type]]'] = '@@this/index::func'; // Added <--------------

// getType call expressions transformed  <--------------
console.log(Rttist.cs$(getType, undefined, ['@@this/index::Foo']));
console.log(Rttist.cs$(getType, undefined, ['@@this/index::Foo::Bar']));
console.log(Rttist.cs$(getType, undefined, ['@@this/index::func']));

const foo = new Foo();
console.log(getType(foo));

So I have to transform those calls, create callsite wrappers for those calls, that will pass type arguments into the function. I also have to identify the type argument used in the call expression getType<Foo>() - from this call, from this line, I have to receive string identifier @@this/index::Foo. I have to find in which module it is declared a what the path to the type in the module is.

Here is a StackBlitz demo with old TypeScript version of the system.

[overlookmotel]

For the code you generate inside the classes, you could do it in a single Traverse pass, by building up a stack when entering a class, and popping off it when exiting the class, so then you can query "what class am I in?" at any time.

But for resolving info about variables which could be anywhere in the AST, yes, I think you'd need a 2-pass approach - either Visit, or possibly iterating through all Symbols and building a Vec of info about those symbols.

It's unfortunate it can't be done all in one pass, but sadly this is one of the restrictions of Rust, as compared to "anything goes" JS. No aliasing allowed! It's the price you pay for the superior performance that the no-aliasing invariant unlocks.

Traverse could be modified to allow you to choose between holding a &mut reference to the current node, or jumping to any other node in the AST (but not both at same time, so as not to violate aliasing rules). But this would be a major undertaking, require fairly deep understanding of Rust and undefined behavior... and we'd have to make sure that change didn't make Traverse extremely unergonomic for the common use case where you don't need this ability.

But if you really want to get into the weeds, and are willing to spend a bunch of time to implement it, I'd be happy to give some guidance ("I don't have time for that!" is a perfectly legitimate answer).

[Hookyns]

I'm a Rust newbie, so I'm not sure I'm able to do that. As you said, "this would be a major undertaking, require fairly deep understanding of Rust." I have some idea, but I don't even know if it's possible in Rust. So I'll continue with my current implementation using Traverse, learn a bit more Rust as I go, and maybe I'll come up with a clever solution later. If I come up with something, I'll definitely let you know. 🙂

[overlookmotel]

OK cool, totally understand. To be honest, I don't think messing with Traverse is a job for a newbie - it already skirts right on the edge of legality, and contains some subtleties which are pretty... subtle. But if you do come up with any bright ideas, please do let me know.

In the meantime, would you be able to mark this question as "answered"?

Good luck with your runtime reflection endeavor!