Writing Einsum in Depth (in OCaml)
Intro
I recently released Einsum in Depth, which I decided to write in OCaml. Perhaps not my wisest decision, but it does at least provide for an interesting experience report.
First, why did I decide to write it in OCaml? When I started working on the project back in 2022 (after which it languished for two years), OCaml made a lot of sense. I can think far more clearly in OCaml than I can in JavaScript, and it allowed me to understand the basics quickly. I initially expected parsing to be a bigger part of the project. If you look back at the initial commit, you can see that I was already using OCamllex and Menhir. For serious parsing projects these tools are fantastic. However, if you look at the code today you can see that I'm doing all of the lexing / parsing by hand. As it turns out, parsing the simple Einsum expressions that I support is really easy and those libraries aren't necessary.
On the Merits of JavaScript vs OCaml
As a language, I prefer OCaml, no question. I also prefer OCaml's tooling. The reason for preferring JavaScript is (a) its libraries, and (b) practicalities of actually running code in a web page. I'll skip discussion of the languages themselves, since, as it's a somewhat religious topic, readers have likely already made up their minds on the issue. However, I'll talk about tooling, libraries, and practicalities.
OCaml Tooling
OCaml's Dune build system is great. It works reliably, is very fast, and can both build and test your code.
The other big highlight on the OCaml side was Jane Street's expect-testing framework (it's not clear whether the official name is ppx_expect or expect-test). This thing is great. You can get in a quick feedback loop by running dune test -w
(where -w
means "watch"). Every time you save it'll rebuild / retest your code, which in OCaml is fast. I mean feedback in a fraction of a second fast.
Here's what a typical expect test looks like:
And here's what it looks like when I run dune test
:
Now the magical part. Since I agree that my test was wrong and needs updating, I can run dune promote
to accept the change, which will update my source file automatically.
One interesting note about expect testing -- it can change the way you write not just tests but code as well. It increases the importance of having nice pretty-printers (of which I'll say more in a moment) since you need to print your data structures nicely for testing. This is a blessing, since it's sometimes useful to have these pretty-printers, but also a curse since you might not have otherwise needed them. All of those pretty-printers need a place to go so now you're encouraged to create a module for every type:
001(* I might have written this *)002type group = string list003004(* What I end up writing *)005module Group = struct006type t = string list007008let pp_friendly = Fmt.(box (list ~sep:sp string))009let pp_original ppf group = Fmt.string ppf (String.concat "" group)010end
Finally, there's js_of_ocaml (AKA JSOO), which compiles OCaml to JavaScript. I honestly don't have much to say about it (and that's a good thing). Again Dune is a hero. You basically just need to tell it to produce JS and you don't have to worry about a thing. Three possible concerns:
- Performance. This was never an issue for me. Apparently OCaml runs faster on JavaScript engines than natively.
- Bundle size. The bundled JavaScript for Einsum in Depth is 319 kB, which I think is not too bad! I had to be careful to keep it that small, by not pulling in libraries unnecessarily. You might have asked above, why did I write my own
intersperse
when Base already includes an implementation? I've previously used Base in another project and was surprised at how much weight it added to JavaScript bundles, so I decided to rough it with the (vastly inferior) OCaml standard library this time. - Binding JS values / interoperability. I used Daniel Bünzli's Brr toolkit and was pleasantly surprised at how easy it made binding JS libraries. So, no complaints, but it still loses hard to just writing JavaScript and not having to think about this at all.
JavaScript Tooling
I only needed a couple of tools from the JS world this time.
I used Tailwind instead of writing my own CSS. Normally Tailwind works by watching your JavaScript files and extracting anything that looks like a Tailwind class name (e.g. rounded-md bg-white px-3
) from strings. It turns out you can also point it at .ml
files and it just works.
I knew to avoid Webpack and decided to use Parcel to bundle my JS. Wait, why do we need a JavaScript bundler at all? Well, js_of_ocaml still doesn't know about JavaScript modules. It operates under the assumption that everything is a global. Yuck.
Anyway, Parcel. It works great, except for the weird way it consistently fails the first time it's run but succeeds the second time.
Don't ask me why: I'm sure there's an interesting yak shave here but I already have plenty of quests on my plate.
One final thing which doesn't exactly fall under "JavaScript Tooling", but close enough: LLMs know JavaScript much better than they know OCaml. They're not bad at OCaml, but they're much more likely to produce an acceptable block of JS.
Pretty-printing
One of the highlights of using OCaml is pretty-printing. The Format
module is complex / comprehensive, though there are various helpful guides. I ended up mostly using Daniel Bünzli's fmt library (yes this is the second time I've mentioned him -- his libraries are great).
Format has this strange, obscure feature called semantic tags, which I used to color variables in code blocks (which are dynamically generated and then pretty-printed) the same as in the text and diagrams.
The way it works is that we can embed a tag (Colored
) with semantic information that the formatter won't do anything with by default.
I then define my own formatter (this isn't the whole thing but it's enough to show the gist), where we maintain a stack of the semantic tags we've entered.
The details don't matter too much. The point I'm trying to convey is that this allows us to add coloring to our pretty-printer. I also abused semantic tags a few years ago to do this:
Libraries
On the OCaml side, I already mentioned Daniel Bünzli's fmt and brr libraries. Well, it turns out he has more obscurely named libraries. The final library I used was his note, for functional reactive programming. Overall I feel pretty positively about it, despite the code getting a bit hairy in places (probably due to my own lack of skill).
On the JavaScript side I used isometric and d3 (despite my reservations) for diagrams.
Yak Shaving
Isometric was the source of my biggest yak shave of this project. I noticed two shortcomings of the library after I started using it.
1. The colors of different axes should change to be readable whether the user has a light or dark color scheme. This is easy to do with tailwind CSS classes (e.g. fill-indigo-800 dark:fill-rose-200
) but much more of a chore if you have to handle it in code (window.matchMedia('(prefers-color-scheme: dark)').addEventListener('change', event => { ... } )
). I noticed that the library didn't have support for classes (necessary for the Tailwind solution), so I added it myself.
2. The library only supported animating one property at a time. If you try an einsum like a b c ->
, you'll see that we might have to animate all three dimensions simultaneously. The maintainer of isometric indicated that they didn't want to support this, so I'm maintaining a fork just for my use which does.
It's hard to tell ahead of time whether a library is going to end up saving you time. Usually it does but sometimes it turns out that it would have been better to write something yourself, but it's never locally the right decision to do so.
Overall a surprisingly strong showing for OCaml. But in the end, Worse is Better, JavaScript eats the world, and I wrote a web page in OCaml so you don't have to.