29 March 2017

For a while now, people have been saying that Rust should appeal to Haskell programmers as an intermediate between Haskell and C. Rust gives you a much more low-level view of your program's execution, but with valuable language constructs such as algebraic data types that allow you to build up to a high-level view of your algorithm. Rust's community places a priority on the ability to create "zero-cost abstractions" so that writing high level code doesn't come at the cost of performance.

The idea of zero-cost abstractions is very appealing. In the Haskell community, there are many great libraries that build up useful abstractions. For some of them, the authors have put a great deal of work in making sure that the performance costs are minimal, such as by making sure the compiler will inline code in common use cases. However, there are also quite a few abstractions that sound very appealing but kill the asymptotics of your program. Free monads are one example, where some algorithms will have bad asymptotics unless you apply a non-obvious trick.

In any case, it's probably been two or three years since I first heard of Rust, and it was about time for me to actually write some code in it. I was especially interested in trying out Rust's style of generics. I've previously ranted about why I've been frustrated by the lack of generics in Go, so I was keen to see what generics in an imperative language could look like.

In order to get the picture I wanted, I chose to implement a data structure based on splay trees, because data structures are the realm in which generics truly shine. By writing the splay algorithm in an appropriately general form, it becomes possible to incorporate that existing code into a new, slightly different data structure.

The Problem

The problem I chose is to write a data structure that supports the following operations:

  • Set the value at a given index to a given value
  • Get the value at a given index
  • Given two indices l and r, reverse the interval [l, r)

The first two operations can be trivially implemented in constant time with an array. However, if you make the decision to use an array, then the third operation will (as far as I know) require linear time, as you need to iterate over the entire interval.

On the other hand, any balanced binary search tree will allow you to implement the first two operations in logarithmic time if you keep the size of the subtree at each node. Splay trees additionally have the property that with a proper sequence of splay operations, you can isolate an interval into its own subtree. By then using a technique called "lazy propagation," it is possible to implement all three operations in (amortized) logarithmic time.

The Code

For these purposes, it's not enough to simply create a splay tree parameterized by the type of data it contains. In a data structure with lazy propagation, each node has a marker for any transformation that has been applied to that subtree. Whenever we visit that node, we want to push that transformation down to its children. In the case of reversals, the pushing operation involves swapping the left and right children and marking each of them as reversed (or unmarking if they are already marked).

Since binary search tree operations and splay operations visit nodes, we want to write them in such a way that they can automatically call this push operation as needed without being unnecessarily aware of them. This can be done very cleanly with the ideas of F-Algebras and F-Coalgebras.

For this particular case, we start with a functor that describes the branching structure of a tree holding values of type a.

data TreeF a b = Empty | Branch a b b

A TreeF a-algebra is a type t with a function TreeF a t -> t. Similarly, a TreeF a-coalgebra is a type t with a function t -> TreeF a t. If we were to define a binary tree holding as in the usual fashion, we'd find it is both a TreeF a-algebra and a TreeF a-coalgebra.

data Tree a = Leaf | Node a (Tree a) (Tree a)

alg :: TreeF a (Tree a) -> Tree a
alg Empty = Leaf
alg (Branch v l r) = Node v l r

coalg :: Tree a -> TreeF a (Tree a)
coalg Leaf = Empty
coalg (Node v l r) = Branch v l r

These alg and coalg functions become the appropriate places to put extra computation that does things like maintain a size annotation or perform the push operations for lazy propagation. In both the Haskell and Rust code, I define all the tree operations in terms of these functions rather than direct pattern matching, which makes the actual lazy propagation relatively painless.

Differences In Rust

No Higher Kinded Types

If you read through the code, you might notice that the way I define an algebra in Haskell is more general than the way I define it in Rust. Haskell gives me the ability to talk about higher kinded types, so I can have an Algebra typeclass that's parameterized by a functor, which in this case will be TreeF a. In Rust, all type parameters need to be concrete types, so I instead define a TreeAlgebra trait that is parameterized by a instead. This can somewhat reduce the potential for code reuse, but I think Rust's system still gets you most of the way there.

Think About Memory

Rust is famously not garbage collected, and as a result memory concerns spread themselves through every function. For this use case I wanted to keep only one tree around at any given time, so it was convenient to have functions that take ownership of the whole tree, and return a structure (either the new tree or a zipper) that encapsulates all of the information needed to construct the next tree. This is a particularly big benefit for splay trees, because improper use of an old value could destroy the amortized time bounds.

Because I was writing my functions in a take-ownership style, I was able to avoid having to think too hard about lifetimes. In fact, none of my functions use explicit lifetime parameters. However, the move semantics did force me to use a few tricks, like when I pattern match on separate(node), if I found a value where I didn't want to change anything, I couldn't just go back and use node like I could in Haskell, but instead needed to recombine the result into a new node.

Overall, both implementations felt rather mathematical, with the same sort of pattern matching and formulas. However, Haskell's combination of lazy propagation and garbage collection make the experience a bit closer to the mathematics. For example, in Rust I have a right_zipper function and a left_zipper function to move down the tree. In Haskell, I can substitute these functions with an Coalgebra (TreeF a) instance for zippers, and because of lazy evaluation as long as I use only one of the branches the other won't be evaluated at all.

The Payoff

The most clear payoff is that the Rust code is dramatically more efficient. I ran both versions of the code (compiled with optimizations) on a randomly generated test case containing 1 million operations on a structure of size 100,000. The two programs yielded identical output (as they should!), but the Rust code ran almost 6 times faster and with 20 times less memory:

Elapsed (wall clock) time: 47.59 seconds
Maximum resident set size: 334368 kB

Elapsed (wall clock) time: 8.78 seconds
Maximum resident set size: 16856 kB

Although not something I observed in this experience, another benefit of Rust's memory management system is that there is no garbage collection. Haskell's garbage collector is optimized for throughput and as a result can cause significant execution pauses. This makes Haskell a questionable choice for real-time applications, as one of these pauses will likely blow right past any latency requirements you might have. If there's one thing that would make me hesitant to use Haskell in a production system right now, it's the garbage collector. Rust sidesteps the problem entirely.

Overall, Rust seems superior to Haskell for most end products. The performance gains are significant without any real increase in work. However, the Haskell code was very conducive to exploring the abstraction space, whereas I had the advantage of knowing the approach from the get-go when writing the Rust code. If I were trying to do that exploration in Rust, I think the memory management concerns would be quite annoying, because I'd be trying to focus on what the abstraction should be, rather than how it's implemented. In the future I expect a lot of my projects to start in Haskell for the exploratory phase, then get ported to Rust if I want a more performant product.