Tree calculus demo: fusion analysis
Illustrating the reflective powers and modularity of tree calculus
As the title of Barry Jay’s book Reflective Programs in Tree Calculus suggests, the ability to write reflective programs is a major feature of tree calculus. Concretely: While λ-calculus and most programming languages are extensional, treating programs as black boxes, tree calculus is intensional — it allows “looking inside” programs.
Program analyses (think optimizations, type checking) are one class of intensional algorithms — operating directly on programs, not source code or abstract syntax trees. The code snippet in the cover image demonstrates this: f is a simple example function that can be called, but also passed as an argument to optimize. optimize f behaves just like f, but is more efficient.
This post briefly motivates and then develops the optimization (fusion) performed by optimize. I’d recommend following and playing along in this live demo. That it implements the analysis (and evaluates on an example) in less than 200 lines of code, using no libraries, demonstrates the practical power of tree calculus.
Motivation: fusion optimization
Imagine a program that processes lists of data in a pipeline of multiple steps, mapping (applying a function to each element to produce a new list) and filtering the data progressively. A naive implementation of this pipeline may perform these steps sequentially, each step materializing an intermediate list, only for the next step to consume and iterate over. We can do better, fusing multiple pipeline steps into one, resulting in fewer intermediate lists to materialize or iterate. Example:
# A function that processes a list of people, first
# selecting out their name and then turning it to upper case
# Naive: Creating an intermediate list
f = map to_upper_case ○ map get_name
# Better: Doing both steps in one go
f = map (to_upper_case ○ get_name)
# Note: "○" is function composition, i.e. (g ○ h)(x) is g(h(x))Query languages like SQL schedule their work smartly to begin with, libraries using iterators rather than lists/arrays get there by deferring materialization, and it makes for a cool compiler optimization to automatically rewrite your code to do less. Here we do the latter. But not as a compiler optimization, but (thanks to tree calculus being intensional) as an higher-order function. It takes a function as argument, inspects it and returns a potentially faster version of it.
Concretely, we write an analysis that detects programs of shape f ○ (map g) ○ (map h) ○ i and rewrites them to f ○ map (g ○ h) ○ i, for arbitrary functions f, g, h and i. The demo supports fusing a few functions beyond map.
Matching on tree structure, explicitly
Key ingredients: All values (that includes programs) in tree calculus are unlabeled binary trees. See here or here for quick refreshers or Barry Jay’s book for an excellent deep dive and context .The “triage” reduction rules (3a-c) allow a tree of shape △ (△ on_leaf on_stem) on_fork to observe the structure of another tree. Concretely:
3a
△ (△ on_leaf on_stem) on_fork △ ---> on_leaf
3b
△ (△ on_leaf on_stem) on_fork (△ u) ---> on_stem u
3c
△ (△ on_leaf on_stem) on_fork (△ u v) ---> on_fork u vUsing triage, we build up primitives is_leaf, is_stem and is_fork that can be composed to match on larger trees. Specifically, is_leaf returns some △ if the argument tree is a leaf, otherwise none. is_stem takes a matcher as argument and returns some x if the tree is a stem and the matcher returns some x when applied to the child of the stem. is_fork is analogous, but taking two matchers as arguments, for left and right subtree. (We represent none as a leaf and some as a stem.) Here is a closer look at the definition and behavior of is_leaf:
# Definition as per demo
is_leaf = triage (some △) (\_ none) (\_ \_ none)
# \_ is syntax for "λ_.", i.e. abstraction of some dummy
# variable. Due to reduction rule 1, "△ △ x" behaves just
# like "λ_.x" so we can rewrite to
= triage (some △) (△ △ none) (△ △ (△ △ none))
# substituting the definition of "triage" (see above)
= △ (△ (some △) (△ △ none)) (△ △ (△ △ none))
# substituting the definitions of "some" and "none"
= △ (△ (△ △) (△ △ △)) (△ △ (△ △ △))
# Example reduction: Applying "is_leaf" to a leaf
3a
△ (△ (△ △) (△ △ △)) (△ △ (△ △ △)) △ ---> △ △
# ^^^ ^^^
# this is the subtree that makes it here
# We interpret the result "△ △" as "some △".
# Example reduction: Applying "is_leaf" to some stem "△ u"
3b 1
△ (△ (△ △) (△ △ △)) (△ △ (△ △ △)) (△ u) ---> △ △ △ u --> △
# ^^^^^ ^^^^^
# this is the subtree that makes it here
# We interpret the result "△" as "none".Function composition g ○ h can be represented by the tree △ (△ (△ △ g)) h (see cover image) and thus recognized by the matcher is_fork (is_stem (is_fork is_leaf some)) some. Equipped with this (called is_composition in the demo), we’re already able to decompose “multi-step pipeline functions” as motivated earlier into their individual pipeline steps! Next, we need to attempt making sense of individual pipeline steps — for instance recognizing them as the list processing function map.
map is a recursive function and has a slightly larger tree representation. See cover image for the tree shape of map f. It would be fine to just repeat our previous strategy and use is_leaf, is_stem and is_fork to build up a matcher, ideally automatically. We’ll instead pick a different strategy that does not require building up an explicit matcher, to convey more intuition about what tree calculus can do.
Matching on tree structure, generically
Note that map f has the displayed tree structure no matter what f is. Our goal is to detect if a given tree has shape map f for any f and if so, extract that f. How would you do this, as a human? Given a mystery tree, you could render it, put it next to the rendering of map f from the cover and visually compare them. Maybe by tracing through both trees simultaneously with two fingers.
We can achieve the same with a function find_mismatch that works as follows: It takes three trees as input, two reference trees that are largely the same, but mismatch in exactly one spot, and a third mystery tree. find_mismatch walks these three trees simultaneously, i.e. asserts that the tree structure of the mystery tree is the same as that of the two reference trees. Abort if the mystery tree mismatches. When the walk reaches the spot where the two reference trees mismatch, return the corresponding subtree of the mystery tree. How does this help?
Imagine calling find_mismatch (map △) (map (△ △)) x where x is our mystery tree. Now, if x is of shape map f, the simultaneous walk of the three trees will succeed, in that they have the exact same shape — until the subtrees △, △ △ and f are reached, respectively. Subtrees △ and △ △ in the two reference trees are like a marker that says “extract this subtree from the mystery tree”. This is a silly but sound way to represent a tree pattern. Rather than representing the pattern using a bespoke data structure (encoded as a tree), we use trees directly. Or rather two trees, encoding the pattern variable/hole using a mismatch between the two trees.
Effectively, this method allows inverting any function func where the tree func x contains x unmodified and is otherwise constant. This is the case for functions like map or filter. Their first argument (a function) is not being consumed/called until another argument (a list) is being passed.
Combining it all
We are able to decompose functions into steps and we are able to detect steps that have shape map f for some f. The rest is plumbing.
In the interactive demo, extract is the function that attempts to detect map and some similar functions and returns none on failure or some f on success. Ultimately, optimize is defined as follows, decorated with some comments:
optimize = fix $ \self \f
# 1. Check if f has shape a ○ b
let (is_composition f) $ triage f _ $ \a \b
# If so, 2. optimize a and b themselves (recursive call)
let (self a) $ \a
let (self b) $ \b
# 3. Check if a and b have familiar shape,
# try extract function args "fa" and "fb"
let (extract a) $ \fa
let (extract b) $ \fb
# 4. Check if extraction of both "fa" and "fb" succeeded
let (both fa fb) $ \fs
triage
# 5a. If not, don't fuse anything but emit a ○ b
(compose a b)
# 5b. If so, emit a function that uses "fa ○ fb"
(triage _ _ $ \fa \fb filter_map (compose_option fa fb))
_
fsThis is how easy it is to analyze and transform programs with tree calculus! Programs, like any other values, are trees, so optimize is an ordinary higher-order function (and tree with ~30k nodes) that can introspect and transform programs.
