flatter closure conversion

[Stephen Weeks]

I have a new (I think) closure conversion algorithm that I would like
to propose.  It doesn't have anything to do with the control-flow
aspects of closure conversion -- it only has to do with the
environment record representation.  I claim that it produces no larger
closure representations than the usual flat closure conversion, and
often smaller representations.  The idea is in closures that are
closed over other closures, to flatten out entirely the inner closure
within the outer one (and transitively any closures it is closed over, 
etc.).

The effect can be achieved by a slight modification of the usual free
variable function that treats lambdas bound in lexically enclosing
scopes "transparently".

Consider the following language:

e ::= x | fn x => e | e e | let x = e in e end

In the following, assume that all bound variables in a program are
distinct.

Here is the usual free variable function.

FV: Exp -> P(Var)
FV(x) = {x}
FV(fn x => e) = FV(e) - {x}
FV(e1 e2) = FV(e1) U FV(e2)
FV(let x = e1 in e2 end) = FV(e1) U (FV(e) - {x})

Here is my new "closure variables" function.  The intuition is that
the closure variables of lambda is the set of variables that will be
stored in the closure record for the lambda.

CV: Exp x (Var -> P(Var)) -> P(Var)
CV(x, m) = m(x) if x in dom(m)
           {x}  otherwise
CV(fn x => e, m) = CV(e, m) - {x}
CV(e1 e2, m) = CV(e1, m) U CV(e2, m)
CV(let x = e1 in e2, m) = V U (CV(e2, m') - {x})
                          where V = CV(e1, m)
                                m' = if e1 is fn y => e1'
                                     then m[x -> V]
                                     else m

Let CV(fn x => e) denote the set of closure variables for the lambda
as computed above.  Also let T(fn x => e) denote (... yi, ...), where
{... yi, ...} is CV(fn x => e). 

Now, the closure conversion of fn x => e is just T(fn x => e)
(possibly with a code pointer, but I'm ignoring those).  The code
built for fn x => e is

fun f(r, x) =
	let ...
		yi = #i r
	    ...
            ...
	    fi = T(fn xi => ei)
	    ...
	in <closure conversion of e>
	end

where the fi's are the lambdas that are in an enclosing scope of e and 
fi occurs in e.  I.e. the input program looks like
	let fi = fn xi => ei
	in ... (fn x => e)
	end

Here is an example.

val f = fn _ => ... a ... b ...
val g = fn _ => ... b ...
val h = fn _ => ... f ... a ... g ... c ...
val i = fn _ => ... h ... d ...

The closure conversion is

val f = (a, b)
val g = (b)
val h = (a, b, c)
val i = (a, b, c, d)

fun F((a, b), _) = ... a ... b ...
fun G((b), _) = ... b ...
fun H((a, b, c), _) =	let val f = (a, b)
			    val g = (b)
			in ... f ... a ... g ... c ...
			end
fun I((a, b, c, d), _) = let val h = (a, b, c)
			 in ... h ... d ...
			 end

I remind you that traditional "flat" closure conversion would do
something like the following.

val f = (a, b)
val g = (b)
val h = (a, c, f, g)
val i = (h, d)

If you look at box and pointer diagrams of these "flat" closures, it
hopefully becomes clear why I call my version flatter.

In fact one can imagine a whole family of strategies ranging from
traditional flat closure conversion to the flatter variant I mention
here, depending on the level of transparency of lambdas.

I don't think this has any problems with space safety (in my usual
weakened sense where I get to pick the constant blowup per program).
Unlike the infamous Shao closure conversion paper, I'm not doing any
sharing of closures, so I don't have to do any lifetime analysis.

While it may appear that this approach will do significantly more
allocation than flat closure conversion, I believe that subsequent Cps
optimization (esp flattening) will eliminate most of the allocation.

Comments are appreciated.  If it sounds sensible, I may move it near
the front of my MLton queue, because I think it sounds pretty easy to
implement (1 day).