Enhanced AST Static Analysis with Typescript Language Server

Hi everyone, I'm Artur, tech lead at Apps Platform team at London-based company called Zoho. Today I'd like to talk about enhancing static code analysis using TypeScript language server. But first, let's discuss quickly what abstract syntax tree is and how can we use it for static analysis.

Imagine we want to implement some automatic code transformation during compilation to remove all CONSOLE.LOC calls from the output. In theory, we could just use a regular expression to find all CONSOLE.LOC in the code. But remember, solve the problem with a regular expression and now you have one more problem. And actually, it'll be quite tricky to write such a regular expression to handle all the different cases. Instead, we can make code transformations using abstract syntax tree.

Let's take a look at code compilation process. Compilation usually contains 4 stages – lexical analysis, syntax analysis, semantic analysis and code generation. Today we'll focus on semantic analysis with abstract syntax tree. And we'll implement our code transformation using AST. Most of the tools of our ecosystem are using AST for analysis and code changes.

Let's take a look at our code again. We're going to use a tool called AST explorer. And on the right side you can see how our code is represented by abstract syntax tree. The function call, which is console.log in our case, is represented by call expression node in the tree, which contains other nodes like callee or arguments. This call expression is a part of the block statement body, let's just remove it. Once we remove it, block statement wouldn't contain any other expressions in the body. So in the end, generating a code will get just an empty Hello World function, in this particular case. Going back to diagram, so what we did, we just transformed our abstract syntax tree, removing call expression node, and then we generated code from the new AST. Of course, we don't want to make these transformations manually, so let's implement a simple Babel plugin to remove all console.log calls. In AST Explore, you can select Built-in Babel API to implement a Babel plugin and you can see the four parts on the screen, like source code, AST, plugin implementation and output code. Original source code contains three console.log call expressions, which are represented by abstract syntax tree. You can see these three expression statements as part of the block statement body. Let's expand each of them and you can find one of the call expressions with callee and argument nodes. Now moving to the plugin implementation, the implementation is quite straightforward. All we need to do is traverse over call expressions, check if callee is a member expression and if that's a console.log, remove the whole expression. In the end, we'll get an empty function, hello world, without console.log again because they were automatically removed. You can use AST and static code analysis to implement build-time optimizations, minification, formatting, transpiling, for example, from the new language features to the old one to support older browsers, for example, and many more.

But how can TypeScript help with static analysis? Let's imagine now that we need to remove console.log calls as before, but keep logs of arguments of number type. So, for example, we want to keep 1, 2, 3 or some constant if that's a number, but we want to remove Hello World strings or Now strings with in-place values like 1, 2, 3, that's pretty straightforward because you can get this value just from the syntax. But it gets quite complicated with function calls or external variables like some constant here. We have to have some kind of type system in place or implement type inference ourselves and TypeScript can provide such types.

In this example, TypeScript can automatically info the correct type based on the function signatures. And to make TypeScript not only check types with the command-line tool but work with different code editors, TypeScript provides a language server which is a separate process like a backend and tools can connect to this server to get the information about types or use code refactoring tools provided by TypeScript. And the beauty of the language server is that it can be used not only for code editing or type checking. Following the language server protocol we can interact with it programmatically and benefit from TypeScript knowledge about the project and types at any level including we can use it during our semantic analysis.

TSMorph is an amazing library which wraps TypeScript compiler API and provides a simple interface to interact with. There is an example how you can configure TSMorph in your program. Now we can use it during abstract tree analysis. Let's see how we can integrate TSMorph into the actual plugin. So there is quite a lot of code here, but let's focus on the most important parts. So we want to have a getTypeAtPost function which accepts file name, code, start and end of the position we are interested in and then we can create a virtual source file based on the file name and code, and read the type for this provided position using TSMorph. The idea is the same as, for example, in IDE. If we need to see the type in our editor, we select an expression we are interested in and editor shows types for this specific cursor position. In this case, instead of cursor, we just have programmable interface but the idea is mostly the same. And now let's integrate TSMorph with our Babel plugin. We need to initialize the defined previously TypedClick class for TS processor. Then we can define a helper called GetTSType with file name, source code, and path to the node in the abstract syntax tree as arguments. Each path in the abstract syntax tree contains nodes with information where the node is positioned from start to end. This way we can check the type for the specific node. And in the end, we just need to return the type representation provided by TSMorph. Having these simple helpers defined, let's see how our original Babel plugin can be enhanced with TypeScript support. To remind, this was the original plugin implementation that's integrated with TypeScript's language server. First of all, again, let's initialize our TSProcessor instance and define GetTSType helper. Next, we can get code and file name from the Babel plugin state. At this state, we already did all the checks, like we check that this is console.log expression, so we just need to go over console.log arguments and check that each argument... We need to get the type for each argument and check that if that argument type is number or not. If that's not a number, we can just remove this argument completely, and if it is number, we can just leave it.

And then, we just need to check that there are any arguments left or not. If there are any, then we can keep this expression because it means that there are number arguments for this console.log call. If not, we can remove the whole expression by path.remove.

So yeah, this example is quite unrealistic. I don't think there is a real use case for console.log and types of console.log parameters. But there is a bit more practical example I've been working on with Compile Time CSS and JS library.

I've been working on library called Tidy, and the idea is that it can provide a simple functional interface with CSS function, which accepts properties like inline styles. But the idea is that instead of generating these styles in the runtime, we can compile them during the build time and extract them into the separate style sheet.

In this case, that's pretty straightforward because blue value and yellow value are static. So we can get them just from the syntax tree. But what if our code is a bit more complicated? And for example, our styles depend on the runtime. In this case, color and font-weight depend on the variant property and weight property, which are not known during the build time.

But with TypeScript, we can get the actual type of each of the properties. Like, for example, for color, we can see that it can be either violet or purple. And for font-weight, we can see that that's a union of lighter, normal, and bold values. Which makes it possible for us to go over these values and pregenerate classes for each of the value.

Thanks to the atomic representation of styles, we can compose these items in the end, just as class name composition. And all we need in the runtime is just to identify the right hash for the class name for this specific value. There is no styles generation. And again, we can benefit from styles compilation during the build time and extract it into the separate stylesheet.

If that sounds interesting to you, please check the tidy.dev website and tidy repository. It's still in a quite experimental state, but any feedback or ideas will be very much appreciated.

And in the end, just some notes. Use TSMorph for programmable interface to interact with TypeScript compiler. And you can get TypeScript type hints in Babel or any other AST traverser just from the position in the code, which is quite useful. And TypeScript integration might slow down static analysis because that's quite a lot of overhead, and it might be a good idea to apply some performance optimizations, for example, isolated into the separate worker.

There are dozens of different language server protocols. It's not just about TypeScript. So you can build a lot of different integrations as well.

And in the end, last but not least, types are not 100% trusted and they can't be because there are types like any or there are commands like TSignore. So the more attention you pay to the type safety of your codebase, the more you can benefit from it, especially during the compilation and build time. There are some useful links.