Linearize an expression syntax tree - data-structures

Question

I defined a recursive type like this
type t = | Param of int | Add of t * t
and it allows me to write symbolic expressions like this
let x = Param(1)
let y = Param(2)
let z = Add(x, y)
# z;;
- : t = Add (Param 1, Param 2)
or something slightly more complex
let rec fib n x y = match n with
| 0 -> (x, y)
| _ -> (fib (n - 1) y (Add(x, y)))
let z5 = fib 5 x y
# z5;;
- : t * t =
(Add (Add (Param 2, Add (Param 1, Param 2)),
Add (Add (Param 1, Param 2), Add (Param 2, Add (Param 1, Param 2)))),
Add (Add (Add (Param 1, Param 2), Add (Param 2, Add (Param 1, Param 2))),
Add (Add (Param 2, Add (Param 1, Param 2)),
Add (Add (Param 1, Param 2), Add (Param 2, Add (Param 1, Param 2))))))
Now I want to convert this recursive thing into a linear form in which each unique Param and Add wold appears only once -- something like this:
(Step 0, (Param 1))
(Step 1, (Param 2))
(Step 2, (Add (Step 0, Step 1))
(Step 3, (Add (Step 1, Step 2))
(Step 4, (Add (Step 2, Step 3))
...
What is this transformation called? And how could it be implemented?

Answer 1

I usually call it flattening or unfolding, some other common names (that are sometimes too specific) are three address code or A-normal form.
It is probably more natural to represent steps with assignments to variables, rather that with the Step notion, e.g.,
type var = int
type const = int
type stmt = Set var * const
type expr = Cst of int | Var of int | Add of expr * expr | ...
So that your flattened fib will look like
Set (0, (Add (Cst 1, Cst 2)))
Set (1, (Add (Cst 2, Var 0)))
Set (2, (Add (Var 0, Var 1)))
...
The flattening itself is quite easy to implement using a recursive function that has type int -> expr -> stmt list -> int * expr * stmt list, e.g., (untested),
let rec flatten v exprs stmts = match expr with
| Cst _ | Var _ as expr -> v,expr,stmts
| Add (x,y) ->
let v,x,stmts = flatten v x stmts in
let v,y,stmts = flatten v x stmts in
v+1,Var (v+1),(Set (v+1,Add(x,y)) :: stmts)
The v parameter is the fresh generator for variable names (which we represent just with integers).
Besides, if you would like to see how flattening works in real-world AST, here is an example.
Implementing hash-consing and flattening
First we need to define a hash-consed representation of our AST, e.g.,
type exp = Cst of int | Ref of int | Add of hexp * hexp
and hexp = {ref : int; exp : exp}
We index each hash-consed expression with a unique ordinal so that physically equal expressions will have the same ref number. We need this in order to be able to store expressions in ordered data structures like maps and sets (OCaml doesn't allow us to order values by their physical address, and it makes sense, since OCaml uses generational GC so the addresses of values change over time).
Now, let's write the hashcons function that will hash-cons an expression, e.g.,
let hashes = Hashtbl.create 100
let hashcons exp =
match Hashtbl.find_opt hashes exp with
| None ->
let ref = Hashtbl.length hashes + 1 in
Hashtbl.add hashes exp ref;
{ref; exp}
| Some ref -> {ref; exp}
Now we can write our fib function that uses the hashconsed representation,
let rec fib n x y = match n with
| 0 -> (x, y)
| _ -> fib (n - 1) y (hashcons (Add (x, y)))
let x = hashcons##Cst 1
let y = hashcons##Cst 2
let _,z = fib 999 x y
No big changes. Now, let's write the flatten function. But before that we need some representation for the program,
module Program = Map.Make(Int)
let set id x = Program.update id (function
| None -> Some x
| x -> x)
let rec get ref prog = match Program.find_opt ref prog with
| None -> None
| Some _ -> Some {ref; exp = Ref ref}
let (++) exp prog = {ref=exp.ref; exp=Ref exp.ref},
set exp.ref exp prog
Our program is a mapping from reference number (which acts both like a variable name and the place where variable is defined - we have a built-in SSA property in our representation), and finally, our flatten function is,
let unify {ref} exp prog =
{ref; exp=Ref ref},
set ref {ref; exp} prog
let rec flatten input prog =
match get input.ref prog with
| Some exp -> exp,prog
| None -> match input.exp with
| Cst _ | Ref _ -> input ++ prog
| Add (x,y) ->
let x,prog = flatten x prog in
let y,prog = flatten y prog in
unify input (Add (x,y)) prog
It is linear, because when we flatten a complex expression we use unify and store the flattened form of that expression under its id. So that the next time we see the same complex expression, we can easily extract its already flattened version from the program, using the integer key for lookup.
And the output is,
# let p = snd##flatten z Program.empty
val p : hexp Program.t = <abstr>
# print_program p;;
#1 := 1
#2 := 2
#3 := #1 + #2
#4 := #2 + #3
#5 := #3 + #4
#6 := #4 + #5
<snip>
#1001 := #999 + #1000
- : unit = ()
where the print_program function is defined as,
let rec pp_exp ppf {exp} = match exp with
| Cst x -> Format.fprintf ppf "%d" x
| Ref x -> Format.fprintf ppf "#%d" x
| Add (x,y) -> Format.fprintf ppf "%a + %a" pp_exp x pp_exp y
let print_program prog =
Program.to_seq prog |>
Seq.iter ## fun (id,exp) ->
Format.printf "#%d := %a#\n" id pp_exp exp
The flattening function works extremely fast, but as always it is not for free. The hash-consed version of fib is now much slower, as we have to hash-cons each expression and during it we have to lookup into a hash table and compare huge values structurally. The fib function is not exploiting the hash-consed representation of the expressions and it is possible to write a more efficient fib function.