Eli Bendersky's website - Rust

Plugins case study: mdBook preprocessors

2025-12-17T18:11:00-08:00

mdBook is a tool for easily creating books out of Markdown files. It's very popular in the Rust ecosystem, where it's used (among other things) to publish the official Rust book.

mdBook has a simple yet effective plugin mechanism that can be used to modify the book output in arbitrary ways, using any programming language or tool. This post describes the mechanism and how it aligns with the fundamental concepts of plugin infrastructures.

mdBook preprocessors

mdBook's architecture is pretty simple: your contents go into a directory tree of Markdown files. mdBook then renders these into a book, with one file per chapter. The book's output is HTML by default, but mdBook supports other outputs like PDF.

The preprocessor mechanism lets us register an arbitrary program that runs on the book's source after it's loaded from Markdown files; this program can modify the book's contents in any way it wishes before it all gets sent to the renderer for generating output.

The official documentation explains this process very well.

Sample plugin

I rewrote my classical "nacrissist" plugin for mdBook; the code is available here.

In fact, there are two renditions of the same plugin there:

One in Python, to demonstrate how mdBook can invoke preprocessors written in any programming language.
One in Rust, to demonstrate how mdBook exposes an application API to plugins written in Rust (since mdBook is itself written in Rust).

Fundamental plugin concepts in this case study

Let's see how this case study of mdBook preprocessors measures against the Fundamental plugin concepts that were covered several times on this blog.

Discovery

Discovery in mdBook is very explicit. For every plugin we want mdBook to use, it has to be listed in the project's book.toml configuration file. For example, in the code sample for this post, the Python narcissist plugin is noted in book.toml as follows:

[preprocessor.narcissistpy]
command = "python3 ../preprocessor-python-narcissist/narcissist.py"

Each preprocessor is a command for mdBook to execute in a sub-process. Here it uses Python, but it can be anything else that can be validly executed.

Registration

For the purpose of registration, mdBook actually invokes the plugin command twice. The first time, it passes the arguments supports <renderer> where <renderer> is the name of the renderer (e.g. html). If the command returns 0, it means the preprocessor supports this renderer; otherwise, it doesn't.

In the second invocation, mdBook passes some metadata plus the entire book in JSON format to the preprocessor through stdin, and expects the preprocessor to return the modified book as JSON to stdout (using the same schema).

Hooks

In terms of hooks, mdBook takes a very coarse-grained approach. The preprocessor gets the entire book in a single JSON object (along with a context object that contains metadata), and is expected to emit the entire modified book in a single JSON object. It's up to the preprocessor to figure out which parts of the book to read and which parts to modify.

Given that books and other documentation typically have limited sizes, this is a reasonable design choice. Even tens of MiB of JSON-encoded data are very quick to pass between sub-processes via stdout and marshal/unmarshal. But we wouldn't be able to implement Wikipedia using this design.

Exposing an application API to plugins

This is tricky, given that the preprocessor mechanism is language-agnostic. Here, mdBook only offers additional utilities to preprocessors implemented in Rust. These get access to mdBook's API to unmarshal the JSON representing the context metadata and book's contents. mdBook offers the Preprocessor trait Rust preprocessors can implement, which makes it easier to wrangle the book's contents. See my Rust version of the narcissist preprocessor for a basic example of this.

Renderers / backends

Actually, mdBook has another plugin mechanism, but it's very similar conceptually to preprocessors. A renderer (also called a backend in some of mdBook's own doc pages) takes the same input as a preprocessor, but is free to do whatever it wants with it. The default renderer emits the HTML for the book; other renderers can do other things.

The idea is that the book can go through multiple preprocessors, but at the end a single renderer.

The data a renderer receives is exactly the same as a preprocessor - JSON encoded book contents. Due to this similarity, there's no real point getting deeper into renderers in this post.

FAAS in Go with WASM, WASI and Rust

2023-05-06T06:44:00-07:00

This post is best described as a technology demonstration; it melds together web servers, plugins, WebAssembly, Go, Rust and ABIs. Here's what it shows:

How to load WASM code with WASI in a Go environment and hook it up to a web server.
How to implement web server plugins in any language that can be compiled to WASM.
How to translate Go programs into WASM that uses WASI.
How to translate Rust programs into WASM that uses WASI.
How to write WAT (WebAssembly Text) code that uses WASI to interact with a non-JS environment.

We're going to build a simple FAAS (Function as a Service) server in Go that lets us write modules in any language that has a WASM target. Comparing to existing technologies, it's something between GCP's Cloud Functions, Cloud Run and good old CGI.

Design

Let's start with a high-level diagram describing how the system works:

The steps numbered in the diagram are:

The FAAS server receives an HTTP GET request, with a path consisting of a module name (func in the example in the diagram) and an arbitrary query string.
The FAAS server finds and loads the WASM module corresponding to the module name it was provided, and invokes it with a description of the HTTP request.
The module emits output to its stdout, which is captured by the FAAS server.
The FAAS server uses the module's stdout as the contents of an HTTP Response to the request it received.

The FAAS server

We'll start our deep dive with the FAAS server itself (full code here). The HTTP handling part is straightforward:

func httpHandler(w http.ResponseWriter, req *http.Request) {
  parts := strings.Split(strings.Trim(req.URL.Path, "/"), "/")
  if len(parts) < 1 {
    http.Error(w, "want /{modulename} prefix", http.StatusBadRequest)
    return
  }
  mod := parts[0]
  log.Printf("module %v requested with query %v", mod, req.URL.Query())

  env := map[string]string{
    "http_path":   req.URL.Path,
    "http_method": req.Method,
    "http_host":   req.Host,
    "http_query":  req.URL.Query().Encode(),
    "remote_addr": req.RemoteAddr,
  }

  modpath := fmt.Sprintf("target/%v.wasm", mod)
  log.Printf("loading module %v", modpath)
  out, err := invokeWasmModule(mod, modpath, env)
  if err != nil {
    log.Printf("error loading module %v", modpath)
    http.Error(w, "unable to find module "+modpath, http.StatusNotFound)
    return
  }

  // The module's stdout is written into the response.
  fmt.Fprint(w, out)
}

func main() {
  mux := http.NewServeMux()
  mux.HandleFunc("/", httpHandler)
  log.Fatal(http.ListenAndServe(":8080", mux))
}

This server listens on port 8080 (feel free to change this or make it more configurable), and registers a catch-all handler for the root path. The handler parses the actual request URL to find the module name. It then stores some information to pass to the loaded module in the env map.

The loaded module is found with a filesystem lookup in the target directory relative to the FAAS server binary. All of this is just for demonstration purposes and can be easily changed, of course. The handler then calls invokeWasmModule, which we'll get to shortly. This function returns the invoked module's stdout, which the handler prints out into the HTTP response.

Running WASM code in Go

Given a WASM module, how do we run it programmatically in Go? There are several high-quality WASM runtimes that work outside the browser environment, and many of them have Go bindings; for example wasmtime-go. The one I like most, however, is wazero; it's a zero-dependency, pure Go runtime that doesn't have any prerequisites except running a go get. Our FAAS server is using wazero to load and run WASM modules.

Here's invokeWasmModule:

// invokeWasmModule invokes the given WASM module (given as a file path),
// setting its env vars according to env. Returns the module's stdout.
func invokeWasmModule(modname string, wasmPath string, env map[string]string) (string, error) {
  ctx := context.Background()

  r := wazero.NewRuntime(ctx)
  defer r.Close(ctx)
  wasi_snapshot_preview1.MustInstantiate(ctx, r)

  // Instantiate the wasm runtime, setting up exported functions from the host
  // that the wasm module can use for logging purposes.
  _, err := r.NewHostModuleBuilder("env").
    NewFunctionBuilder().
    WithFunc(func(v uint32) {
      log.Printf("[%v]: %v", modname, v)
    }).
    Export("log_i32").
    NewFunctionBuilder().
    WithFunc(func(ctx context.Context, mod api.Module, ptr uint32, len uint32) {
      // Read the string from the module's exported memory.
      if bytes, ok := mod.Memory().Read(ptr, len); ok {
        log.Printf("[%v]: %v", modname, string(bytes))
      } else {
        log.Printf("[%v]: log_string: unable to read wasm memory", modname)
      }
    }).
    Export("log_string").
    Instantiate(ctx)
  if err != nil {
    return "", err
  }

  wasmObj, err := os.ReadFile(wasmPath)
  if err != nil {
    return "", err
  }

  // Set up stdout redirection and env vars for the module.
  var stdoutBuf bytes.Buffer
  config := wazero.NewModuleConfig().WithStdout(&stdoutBuf)

  for k, v := range env {
    config = config.WithEnv(k, v)
  }

  // Instantiate the module. This invokes the _start function by default.
  _, err = r.InstantiateWithConfig(ctx, wasmObj, config)
  if err != nil {
    return "", err
  }

  return stdoutBuf.String(), nil
}

Interesting things to note about this code:

wazero supports WASI, which has to be instantiated explicitly to be usable by the loaded modules.
A lot of the code deals with exporting logging functions from the host (the Go code of the FAAS server) to the WASM module.
We set up the loaded module's stdout to be redirected to a buffer, and set up its environment variables to match the env map passed in.

There are several way for host code to interact with WASM modules using only the WASI API and ABI. Here, we opt for using environment variables for input and stdout for output, but there are other options (see the Other resources section in the bottom for some pointers).

This is it - the whole FAAS server, about 100 LOC of commented Go code. Now let's move on to see some WASM modules this thing can load and run.

Writing modules in Go

We can compile Go code to WASM that uses WASI. Here's a basic Go program that emits a greeting and a listing of its environment variables to stdout:

package main

import (
  "fmt"
  "os"
)

func main() {
  fmt.Println("goenv environment:")

  for _, e := range os.Environ() {
    fmt.Println(" ", e)
  }
}

To build using the standard Go toolchain:

$ GOOS=wasip1 GOARCH=wasm gotip build -o target/goenv.wasm examples/goenv/goenv.go

(the wasip1 target name refers to "WASI Preview 1")

Sharp-eyed readers will recall that the target/ directory is precisely where the FAAS server looks for *.wasm files to load as modules. Now that we've placed a module named goenv.wasm there, we're ready to launch our server with go run . in the root directory. We can issue a HTTP request to its goenv module in a separate terminal:

$ curl "localhost:8080/goenv?foo=bar&id=1234"
goenv environment:
  http_method=GET
  http_host=localhost:8080
  http_query=foo=bar&id=1234
  remote_addr=127.0.0.1:59268
  http_path=/goenv

And looking at the terminal where the FAAS server runs we'll see some logging like:

2023/04/29 06:35:59 module goenv requested with query map[foo:[bar] id:[1234]]
2023/04/29 06:35:59 loading module target/goenv.wasm

Another option to build this is using the TinyGo compiler. In our FAAS project structure, the invocation from the root directory is:

$ tinygo build -o target/goenv.wasm -target=wasi examples/goenv/goenv.go

Before Go 1.21, TinyGo has been the only way to target WASI, but this support has been part of the standard toolchain for a few releases now.

Writing modules in Rust

Rust is another language that has good support for WASM and WASI in the build system. After adding the wasm32-wasi target with rustup, it's as simple as passing the target name to cargo:

$ cargo build --target wasm32-wasi --release

The code is straightforward, similarly to the Go version:

use std::env;

fn main() {
    println!("rustenv environment:");

    for (key, value) in env::vars() {
        println!("  {key}: {value}");
    }
}

Writing modules in WebAssembly Text (WAT)

As we've seen, compiling Go and Rust code to WASM is fairly easy; looking for a challenge, let's write a module in WAT! As I've written before, I enjoy writing directly in WAT; it's educational, and produces remarkably compact binaries.

The "educational" aspect quickly becomes apparent when thinking about our task. How exactly am I supposed to write to stdout or read environment variables using WASM? This is where WASI comes in. WASI defines both an API and ABI, both of which will be visible in our sample. The following shows some code snippets with explanations; for the full code check out the sample repository.

First, I want to show how output to stdout is done; we start by importing the fd_write WASI system call:

(import "wasi_snapshot_preview1" "fd_write" (func $fd_write (param i32 i32 i32 i32) (result i32)))

Apparently, it has four i32 parameters and returns an i32; what do all of these mean? Unfortunately, WASI documentation could use a lot of work; the resources I found useful are [1]:

With this in hand, I was able to concoct a useful println equivalent in WAT that uses fd_write under the hood:

;; println prints a string to stdout using WASI.
;; It takes the string's address and length as parameters.
(func $println (param $strptr i32) (param $len i32)
    ;; Print the string pointed to by $strptr first.
    ;;   fd=1
    ;;   data vector with the pointer and length
    (i32.store (global.get $datavec_addr) (local.get $strptr))
    (i32.store (global.get $datavec_len) (local.get $len))
    (call $fd_write
        (i32.const 1)
        (global.get $datavec_addr)
        (i32.const 1)
        (global.get $fdwrite_ret)
    )
    drop

    ;; Print out a newline.
    (i32.store (global.get $datavec_addr) (i32.const 850))
    (i32.store (global.get $datavec_len) (i32.const 1))
    (call $fd_write
        (i32.const 1)
        (global.get $datavec_addr)
        (i32.const 1)
        (global.get $fdwrite_ret)
    )
    drop
)

This uses some globals that you'll have to look up in the full code sample if you're interested. Here's another helper function that prints out a zero-terminated string to stdout:

;; show_env emits a single env var pair to stdout. envptr points to it,
;; and it's 0-terminated.
(func $show_env (param $envptr i32)
    (local $i i32)
    (local.set $i (i32.const 0))

    ;; for i = 0; envptr[i] != 0; i++
    (loop $count_loop (block $break_count_loop
        (i32.eqz (i32.load8_u (i32.add (local.get $envptr) (local.get $i))))
        br_if $break_count_loop

        (local.set $i (i32.add (local.get $i) (i32.const 1)))
        br $count_loop
    ))

    (call $println (local.get $envptr) (local.get $i))
)

The fun part about writing assembly is that there are no abstractions. Everything is out in the open. You know how strings are typically represented using either zero termination (like in C) or a (start, len) pair? In manual WAT code that uses WASI we have the pleasure of using both approaches in the same program :-)

Finally, our main function:

(func $main (export "_start")
    (local $i i32)
    (local $num_of_envs i32)
    (local $next_env_ptr i32)

    (call $log_string (i32.const 750) (i32.const 19))

    ;; Find out the number of env vars.
    (call $environ_sizes_get (global.get $env_count) (global.get $env_len))
    drop

    ;; Get the env vars themselves into memory.
    (call $environ_get (global.get $env_ptrs) (global.get $env_buf))
    drop

    ;; Print out the preamble
    (call $println (i32.const 800) (i32.const 19))

    ;; for i = 0; i != *env_count; i++
    ;;   show env var i
    (local.set $num_of_envs (i32.load (global.get $env_count)))
    (local.set $i (i32.const 0))
    (loop $envvar_loop (block $break_envvar_loop
        (i32.eq (local.get $i) (local.get $num_of_envs))
        (br_if $break_envvar_loop)

        ;; next_env_ptr <- env_ptrs[i*4]
        (local.set
            $next_env_ptr
            (i32.load (i32.add  (global.get $env_ptrs)
                                (i32.mul (local.get $i) (i32.const 4)))))

        ;; print out this env var
        (call $show_env (local.get $next_env_ptr))

        (local.set $i (i32.add (local.get $i) (i32.const 1)))
        (br $envvar_loop)
    ))
)

We can now compile this WAT code into a FAAS module and re-run the server:

$ wat2wasm examples/watenv.wat -o target/watenv.wasm
$ go run .

Let's try it:

$ curl "localhost:8080/watenv?foo=bar&id=1234"
watenv environment:
http_host=localhost:8080
http_query=foo=bar&id=1234
remote_addr=127.0.0.1:43868
http_path=/watenv
http_method=GET

WASI: API and ABI

I've mentioned the WASI API and ABI earlier; now it's a good time to explain what that means. An API is a set of functions that programs using WASI have access to; one can think of it as a standard library of sorts. Go programmers have access to the fmt package and the Println function within it. Programs targeting WASI have access to the fd_write system call in the wasi_snapshow_preview1 module, and so on. The API of fd_write also defines how this function takes parameters and what it returns. Our sample uses three WASI functions: fd_write, environ_sizes_get and environ_get.

An ABI is a little bit less familiar to most programmers; it's the run-time contract between a program and its environment. The WASI ABI is currently unstable and is described here. In our program, the ABI manifests in two ways:

The main entry point we export is the _start function. This is automatically called by a WASI-supporting host after setup.
Our WASM code exports its linear memory to the host with (memory (export "memory") 1). Since WASI APIs require passing pointers to memory, both the host and the WASM module need a shared understanding of how to access this memory.

Naturally, both the Go and Rust implementations of FAAS modules comply to the WASI API and ABI, but this is hidden by the compiler from programmers. In the Go program, for example, all we need to do is write a main function as usual and therein emit to stdout using Println. The Go compiler will properly export _start and memory:

$ wasm-objdump -x target/goenv.wasm

... snip

Export[2]:
 - func[1028] <_rt0_wasm_wasip1> -> "_start"
 - memory[0] -> "memory"

... snip

And will properly hook things up to call our code from _start, etc.

WASI and plugins

The FAAS server presented in this post is clearly an example of developing plugins using WASM and WASI. This is an emerging and exciting area in programming and lots of progress is being made on multiple fronts. Right now, WASI modules are limited to interacting with the environment via means like environment variables and stdin/stdout; while this is fine for interacting with the outside world, for host-to-module communication it's not amazing, in my experience. Therefore the WASM standards committee is working of further improvements to WASI that may include sockets and other means of passing data between hosts and modules.

In the meantime, projects are making do with what they have. For example, the sqlc Go package supports WASM plugins. The communication with plugins happens as follows: the host encodes a command into a protobuf and emits it to the plugin's stdin; it then reads the plugin's stdout for a protobuf-encoded response.

Other projects are taking more maverick approaches; for example, the Envoy proxy supports WASM plugins by defining a custom API and ABI between the host and WASM modules. I'll probably write more about this in a later post.

Other resources

Here are some additional resources on the same topic as this post:

Building a WebAssembly-powered serverless platform - the post that inspired this one. Other WASM-related articles from that blog are also recommended.
wagi - a more featured approch with the same idea, implemented in Rust.
WCGI - an even more CGI-conformant approach.

[1]	The main Go compiler supported WASM for a while, but only as a target for browsers & JS (I wrote a bit about it recently). The news in 1.21 is the `GOOS=wasip1` support, which hooks up Go's internal syscalls to WASI APIs, sets up the ABI etc.

[2]	These resources tend to disappear and move around without notice; if you notice a dead link, please drop me a note and I'll look for a more up-to-date source.

How I went about learning Rust

2022-07-15T20:23:00-07:00

Rust has been on my radar for a long time now, and about a year ago I finally began allocating some time every week to learning it.

In this post I'll provide details on the learning path I've followed for Rust, in the hope that this may prove useful for others. You'll note that this isn't exactly a "Learn X in 24 hours" kind of journey, as it stretches over a long time period and covers lots of different materials. YMMV.

This learning journey alternates consuming information (mostly books) with coding activities (some guided by books, some not). In my experience, both are critical for really learning a language or any other tool deeply, and the combination of the two is powerful.

Consuming information

These are the Rust-specific books I've read, roughly in chronological order. There are additional books in the "writing code" section, but those are not Rust-specific.

Programming Rust (link to review) is the first book I read for an initial introduction to the language.

The Rust Programming Language - I haven't read this one cover to cover, but have skimmed major parts of it, implementing the projects it lists. I've also used it a bunch as a reference since it usually comes up high in Google rankings; it's a better reference than Programming Rust, IMHO.

Rust in Action (link to review).

Rust for Rustaceans (link to review).

In addition to these books, I've watched a few of John Gjengset's videos on YouTube. These are very interesting, and I'll probably watch more to learn about specific advanced concepts of Rust, as needed.

Finally, over this year I've read a bunch of blog posts about Rust, mostly stuff that made it to the front pages of HN or r/rust from time to time.

Writing code

rustlings - small exercises for reading and writing snippets of Rust code. Nice concept, though not very comprehensive. It's still quite useful when just getting started.

Advent of Code - I love AOC, and the 2021 edition was a great opportunity to practice some Rust. I've implemented the solutions for days 1-18 in Rust, for a total of ~2 KLOC. Solving AOC problems is a terrific way to learn and practice programming languages. I'm very likely to tackle AOC 2022 in Rust again, to keep my skills sharp.

The Ray Tracer Challenge book (link to review). This book is language agnostic, and I picked Rust to implement a simple ray tracer. Another great way to practice a new language! Almost 4 KLOC for this one.

Crafting Interpreters (link to review). Reimplemented the majority of the clox compiler and VM in Rust, ~2 KLOC.

In addition to these, I've done quite a bit of experimentation, writing code to explore various areas of Rust. Some of this was mentioned in blog posts, but most wasn't. I estimate there's a total of ~6 KLOC of Rust there.

Overall, over these past months I've written about 14 KLOC of Rust in the process of learning and practicing the language.

What's next

I wouldn't call myself a Rust expert, but over this last year I believe I've developed a good understanding and intuition for the language, to the extent that I would feel confident using it for a production project.

Overall, I like Rust and I'm happy to have added it to my toolbox. It has issues, and I certainly have opinions about it, but that will have to wait for a separate post :-)

Rewriting the lexer benchmark in Rust

2022-05-30T05:53:00-07:00

I've reimplemented my lexical analyzer (lexer) for the the TableGen language in Python, JS and Go so far. The latest post in the series discussed how several years of new Go versions improved my lexer's performance by roughly 2x, and how several additional optimizations won another 37%; the final result is a lexer that churns through 1 MiB of source code in just 5.6 milliseconds.

Since I've also been playing with Rust recently, I thought it would be interesting to re-implement this lexer once again, this time in Rust. Rust is certainly the lowest-level language among those I've used so far for this task, so I expect to see some top-notch performance.

The full code for this post is available on GitHub.

Designing an API

I find that Rust's strict ownership rules makes one think carefully about API design from very early on. If we want to create a lexer and pass it a string as input - who owns the string? In Rust, the answer to this question is not an implicit contract (like it always is in C and sometimes in C++), and it cannot be deferred to the runtime either (like one would do in Python, JS or Go). The answer has to be explicitly encoded into the types of the program.

Since one of the goals of this series of posts is performance, I decided to go with a zero-copy API at first, where the user of the lexer owns the input string; the lexer, in turn, returns tokens that contain references into this input string - so the user ends up owning these too. Rust's lifetime specifiers make this fairly natural; here's the type:

pub struct Lexer<'source> {
  input: &'source str,
  iter: Peekable<CharIndices<'source>>,

  // c is the last char taken from iter, and ci is its offset in the input.
  c: char,
  ci: usize,

  // error is true iff the lexer encountered and error.
  error: bool,
}

Ignore the fields for now, focusing just on the struct definition. This is the constructor:

impl<'source> Lexer<'source> {
    pub fn new(input: &'source str) -> Self {
      // ...
    }
}

The 'source lifetime specifier is used to explicitly annotate the lifetime of the input string slice, since we want to store it in our Lexer and refer to it later. Once a Lexer is created, the user obtains new tokens by calling next_token:

pub fn next_token(&mut self) -> Token<'source> {
  // ...
}

Note that the returned Token also has the same lifetime annotation. Here's how the type is defined:

#[derive(Debug, PartialEq, Clone, Copy)]
pub struct Token<'source> {
    pub value: TokenValue<'source>,
    pub pos: usize,
}

#[derive(Debug, PartialEq, Clone, Copy)]
pub enum TokenValue<'source> {
    EOF,
    Error,

    Plus,
    Minus,
    Multiply,
    Divide,
    Period,
    Backslash,
    Colon,
    Percent,
    Pipe,
    Exclamation,
    Question,
    Pound,
    Ampersand,
    Semi,
    Comma,
    LeftParen,
    RightParen,
    LeftAng,
    RightAng,
    LeftBrace,
    RightBrace,
    LeftBracket,
    RightBracket,
    Equals,
    Comment(&'source str),
    Identifier(&'source str),
    Number(&'source str),
    Quote(&'source str),
}

Now it should be clear how these lifetimes are tied together. Some tokens hold slices into the user's input, and this is encoded in the explicit lifetimes. The signature of the next_token method says "we return tokens with a lifetime that's tied to the lifetime of the input passed into the constructor".

We can also provide a more natural iteration API for the Lexer by implementing the Iterator trait:

// Lexer is an Iterator; it returns tokens until EOF is encountered, when it
// returns None (the EOF token itself is not returned). Note that errors are
// still returned as tokens with TokenValue::Error.
impl<'source> Iterator for Lexer<'source> {
    type Item = Token<'source>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.error {
            // If an error has already been set before we invoke next_token,
            // it means we've already returned TokenValue::Error once and now
            // we should terminate the iteration.
            return None;
        }

        let tok = self.next_token();
        if tok.value == TokenValue::EOF {
            None
        } else {
            Some(tok)
        }
    }
}

The iterator implementation makes it possible to integrate the lexer with the rest of Rust very elegantly; for example, to obtain all the tokens in a given input as a vector:

pub fn tokenize_all_collect<'source>(data: &'source str) -> Vec<Token<'source>> {
    let lex = Lexer::new(&data);
    lex.collect()
}

Implementation

Generally, the implementation of the lexer in Rust follows the same approach used by my previous hand-written lexers in this series. I'd like to highlight a couple of Rust-specific aspects here.

Our lexer fully supports Unicode, so I decided to use Rust's string iterator support to obtain chars from &str. Rust provides a helpful iterator called CharIndices, which yields chars along with their position in the input - we need the position to report token locations. Furthermore, since our lexer requires a bit of look-ahead [1], the iterator is wrapped in Peekable, which provides the peek method. As we've seen above already, the final definition is:

iter: Peekable<CharIndices<'source>>,

With this in mind, the code of the lexer should be very readable, even for someone not too familiar with Rust.

Another note is how sub-string extraction happens when tokens are returned. As an example, let's look at the scan_number method which is invoked when the lexer's current character is a digit:

fn scan_number(&mut self) -> Token<'source> {
    let startpos = self.ci;
    while self.c.is_digit(10) {
        self.scan_char();
    }
    Token {
        value: TokenValue::Number(&self.input[startpos..self.ci]),
        pos: startpos,
    }
}

Recall that the Number variant of the TokenValue enum is defined as Number(&'source str) - it contains a reference to the input source string. In the code for scan_number, we see how this is actually implemented, by sub-slicing the input slice. This creates a slice with the same lifetime as the input slice (which is encoded in the lifetimes of the types). As in Go, sub-slicing is a very cheap operation in Rust (no heap allocation).

Performance

The performance of this Rust-implemented lexer is very good! I ran it on the same large TableGen file I've used for all the benchmarking, and it finishes tokenizing it in just 3.7 ms; this is about 33% faster than my fastest Go version from the previous post.

Profiling Rust is quite a bit trickier than Go, but I managed to cobble together enough magic flags and perf invocations to ascertain that next_token indeed takes the bulk of the time; this is good - the time is being spent where it's supposed to be spent.

Trying a variant with allocations

As described above, the API for this Rust lexer was designed to be zero-copy. Since my Go lexers experimented with different approaches, I've decided to see how a different API in Rust would look.

There are two aspects to ownership we can change: taking ownership of the input string, and/or returning owned Strings in the tokens.

For taking ownership of the input string, our constructor would have to look like:

pub fn new(input: String) -> Self

Does it make sense for a lexer to own its input? This isn't clear, and in reality it turned out to be tricky to implement due to lifetime issues. Rust is very unhappy when a struct field is a reference to another field of the same struct, because there is no safe way to move instances of such structs. In our case, the iterator (a struct field) needs a reference to the string (another field in the same struct), and this is doesn't pass the Rust compiler's scrutiny. I wanted to be able to implement this without actively fooling the compiler by using opaque indices or unsafe.

When the language fights you this hard, it may be a good sign that the design is wrong - it's better to leave the ownership to the code that creates a Lexer.

Ownership of returned tokens was easier to set up. The code for this is available in the owning.rs file of the repository. In this variant, the constructor doesn't change, but tokens are defined differently:

#[derive(Debug, PartialEq, Clone)]
pub struct Token {
    pub value: TokenValue,
    pub pos: usize,
}

#[derive(Debug, PartialEq, Clone)]
pub enum TokenValue {
    // .. types
    Comment(String),
    Identifier(String),
    Number(String),
    Quote(String),
}

Note that variants like Identifier now hold an owning String instead of a string reference. Therefore, lifetime annotations are no longer necessary on Token.

The scanning code now allocates a new String and reads characters into it from the iterator. We've seen previously how scan_number looks; here it is again, for the owning token variant (with some helper methods):

fn scan_number(&mut self) -> Token {
    let startpos = self.ci;
    Token {
        value: TokenValue::Number(self.scan_while_true(|c| c.is_digit(10))),
        pos: startpos,
    }
}

// Helper to scan chars while `pred(c)` returns true, into the given `s`.
fn scan_while_true_into<F>(&mut self, s: &mut String, pred: F)
where
    F: Fn(char) -> bool,
{
    while pred(self.c) {
        s.push(self.c);
        self.scan_char();
    }
}

// Helper to scan chars while `pred(c)` returns true and return all scanned
// chars in a new String.
fn scan_while_true<F>(&mut self, pred: F) -> String
where
    F: Fn(char) -> bool,
{
    let mut s = String::with_capacity(8);
    self.scan_while_true_into(&mut s, pred);
    s
}

Performance of this variant

When I first tried this variant, its performance wasn't great at all! It was about 30% slower than the string-copying version in Go. I think this has to do with the slight difference in approach - in Go, my lexer figures out the token boundaries using numerical indices and then converts a []byte into string in one fell swoop. The Rust version fetches chars one by one from an iterator and writes them into a String.

In particular, since Rust's String is dynamically allocated, this may incur reallocations (depending on how much is allocated initially).

So my solution was to create these strings with_capacity - as you can see in the previous code sample. This cut down the execution time to be roughly equal to Go's version where strings are copied.

Which API is better?

IMHO there's little doubt that the original "slice" API is better, for multiple reasons:

Performance: the results speak for themselves - the slice API is zero-copy and incurs no additional heap allocations. Like in Go, it deals in slice headers, which are just pointers to parts of a string.
The ownership story is very clear, symmetrical and explicitly documented with lifetime annotations. The 'source lifetime controls everything: it's the lifetime of the &str passed into the lexer's constructor, and it's the lifetime of the tokens. The symmetry feels important - the code that creates a lexer controls the lifetime of the input, as well as the output.
In general, there's a known good practice in API design which is not to force allocations on users, if possible. This is well articulated in this blog post for Go, but it applies equally well in Rust. In our slice API, the user can always call to_owned on the returned slice, if they so wish. But why do it for them? Why should we assume the users want us to return them an owned string? Returning a slice provides more flexibility in using the API.

Performance improvement: avoiding `Peekable`

After this post was published, the reader Utkarsh Kukreti mentioned that using Peekable is not the most optimal approach, since its next is slower than the underlying iterator's next, and calling it is on the hot path. Instead of using Peekable, we could just clone the iterator when we see a / and invoke next on the clone; iterators are small so cloning them is cheap.

Here's a patch with this change. Applying it to my lexer makes it ~11% faster on my machine, finishing the benchmark in about 3.4 ms!

To be honest, this is a little bit surprising and disappointing to me, since I was hoping that Peekable would fall into the zero abstraction promise of Rust. Perhaps this is something that can fixed in the future.

[1]	For properly tokenizing comments; when the lexer sees `/`, it needs to peek forward to see if this is the beginning of a `//`, or else the division operator. This is a common use case in lexers, especially when tokenizing multi-character operators (like `>=`).

Rust extension traits, greppability and IDEs

2022-01-29T06:16:00-08:00

Traits are a central feature of Rust, critical for its implementation of polymorphism; traits are used for both static (by serving as bounds for generic parameters) and dynamic (by having trait objects to serve as interfaces) polymorphism.

This post assumes some familiarity with traits and discusses only a specific aspect of them - how extension traits affect code readability. To learn the basics of traits in Rust, the official book is a good starting point.

Extension traits

This Rust RFC provides a good, short definition of extension traits:

Extension traits are a programming pattern that makes it possible to add methods to an existing type outside of the crate defining that type.

For example, here's a trait with a single method:

trait Magic {
    fn magic_num(&self) -> usize;
}

We can now implement the Magic trait for our types:

struct Foobar {
    name: String,
}

impl Magic for Foobar {
    fn magic_num(&self) -> usize {
        if self.name.is_empty() {
            2
        } else {
            33
        }
    }
}

Now a FooBar can be passed wherever a Magic is expected. FooBar is a custom type, but what's really interesting is that we can also implement Magic for any other type, including types that we did not define. Let's implement it for bool:

impl Magic for bool {
    fn magic_num(&self) -> usize {
        if *self {
            3
        } else {
            54
        }
    }
}

We can now write code like true.magic_num() and it will work! We've added a method to a built-in Rust type. Obviously, we can also implement this trait for types in the standard library; e.g.:

impl<T> Magic for Vec<T> {
    fn magic_num(&self) -> usize {
        if self.is_empty() {
            10
        } else {
            5
        }
    }
}

Extension traits in the wild

Extension traits aren't just a fringe feature; they are widely used in the Rust ecosystem.

One example is the popular serde crate, which includes code that serializes and deserializes data structures in multiple formats. One of the traits serde provides is serde::Serialize; once we import this trait and one of the concrete serializers serde provides, we can do stuff like [1]:

let mut serializer = serde_json::Serializer::new(std::io::stdout());
185.serialize(&mut serializer).unwrap();

Importing serde::Serialize is critical for this code to work, even though we don't refer to Serialize anywhere in our code explicitly. Rust requires traits to be explicitly imported to imbue their methods onto existing types; otherwise it's hard to avoid naming collisions in case multiple traits from different crates provide the same methods.

Another example is the byteorder crate, which helps encode numbers into buffers with explicit length and endianness. To write some numbers into a vector byte-by-byte, we have to import the relevant trait and enum first, and then we can call the newly-added methods directly on a vector:

use byteorder::{LittleEndian, WriteBytesExt};

// ...

let mut wv = vec![];
wv.write_u16::<LittleEndian>(259).unwrap();
wv.write_u16::<LittleEndian>(517).unwrap();

The write_u16 method is part of the WriteBytesExt trait, and it's implemented on a Vec by the byteorder crate. To be more precise, it's automatically implemented on any type that implements the Write trait.

Finally, let's look at rayon - a library for simplified data-parallelism. It provides magical iterators that have the same functionality as iter but compute their results in parallel, leveraging multiple CPU cores. The rayon documentation recommends to import the traits the crate injects as follows:

It is recommended that you import all of these traits at once by adding use rayon::prelude::* at the top of each module that uses Rayon methods.

Having imported it thus, we can proceed to use Rayon as follows:

let exps = vec![2, 4, 6, 12, 24];
let pows_of_two: Vec<_> = exps.par_iter().map(|n| 2_u64.pow(*n)).collect();

Note the par_iter, which replaces a regular iter. It's been magically implemented on a vector, as well as a bunch of other types that support iteration.

On greppability and code readability

All these uses of extension traits are pretty cool and useful, no doubt. But that's not the main point of my post. What I really want to discuss is how the general approach relates to code readability, which is in my mind one of the most important aspects of programming we should all be thinking about.

This Rust technique fails the greppability test; it's not a word I made up - google it! If it's not immediately apparent, greppability means the ability to explore a code base using textual search tools like grep, git grep, ripgrep, pss or what have you.

Suppose you encounter this piece of code in a project you're exploring:

let mut wv = vec![];
wv.write_u16::<LittleEndian>(259).unwrap();

"Interesting", you think, "I didn't know that Vec has a write_u16 method". You quickly check the documentation - indeed, it doesn't! So where is it coming from? You grep the project... nothing. It's nowhere in the imports. You examine the imports one by one, and notice the:

use byteorder::{LittleEndian, WriteBytesExt};

"Aha!", you say, "this imports LittleEndian, so maybe this has to do with the byteorder crate". You check the documentation of that crate and indeed, you find the write_u16 method there; phew.

With par_iter you're less lucky. Nothing in imports will catch your eye, unless you're already familiar with the rayon crate. If you're not, then use rayon::prelude::* won't ring much of a bell in relation to par_iter.

Of course, you can just google this symbol like this and you'll find it. Or maybe you don't even understand what the problem is, because your IDE is perfectly familiar with these symbols and will gladly pop up their documentation when you hover over them.

IDEs and language servers

These days we have free, powerful and fast IDEs that make all of this a non-issue (looking at Visual Studio Code, of course). Coupled with smart language servers, these IDEs are as familiar with your code as the compiler; the language servers typically run a full front-end sequence on the code, ending up with type-checked ASTs cross-referenced with symbol tables that let them understand where each symbol is coming from, its type and so on. For Rust the language server is RLS, for Go its gopls; all popular languages have them these days [2].

It's entirely possible that using a language like Rust without a sophisticated IDE is madness, and I'm somewhat stuck in the past. But I have to say, I do lament the loss of greppability. There's something very universal about being able to understand a project using only grep and the official documentation.

In fact, for some languages it's likely that this has been the case for a long while already. Who in their right mind has the courage to tackle a Java project without an IDE? It's just that this wasn't always the case for systems programming languages, and Rust going this way makes me slightly sad. Or maybe I'm just too indoctrinated in Go at this point, where all symbol access happens as package.Symbol, packages are imported explicitly and there is no magic name injection anywhere (almost certainly by design).

I can't exactly put my finger on why this is bothering me; perhaps I'm just yelling at clouds here. While I'm at it, I should finally write that post about printf-based debugging...

[1]	Note that it could be simpler to use `serde`'s `to_json` function here, but I opted for the explicit serializer because I wanted to show how we invoke a new method on an integer literal.

[2]	Apparently, not all tooling has access to sophisticated language servers; for example, as far as I can tell GitHub source analysis won't be able to find where `write_u16` is coming from, and the same is true of Sourcegraph's default configuration (though I've been told this is in the works).