Recursive inlining

Stream: general

[Jannis Limperg]

If I mark a recursive function @[inline], will it get inlined? I'm asking because afaik some compilers don't support recursive inlining even in relatively straightforward cases. (My use case is meta, in case that affects the answer.)

[Jason Rute]

Dumb question. It seems from this stack overflow that recursive inclining is something like hard coding the first few cases up to a fixed depth before calling the recursive function. Is that what you had in mind? https://stackoverflow.com/questions/190232/can-a-recursive-function-be-inline

[Jannis Limperg]

No, my use case is much simpler than general recursive inlining (as I realised when considering your question, so thanks for that). I have definitions of the form

@[inline] meta def f (x : X) : Y -> Z := ... (f ...) ...

meta def g : Y -> Z := f x'

and I'd like to be sure that the code generated for g is the same as if I'd replaced f by its definition manually. In GHC Haskell (which doesn't do recursive inlining ever), I think I'd have to write f as a non-recursive function with a recursive where clause, but Lean doesn't have those.

[Mario Carneiro]

Why not write

meta def f (x : X) : Y -> Z := ... (f ...) ...
@[inline] meta def f' (x : X) : Y -> Z := ... (f ...) ...

[Mario Carneiro]

you could also use a tactic to perform the copy paste job

[Jannis Limperg]

Let's get concrete; I'm getting lost in my own abstractions. Here's a simpler version of the code I'm writing:

open expr

@[inline] meta def match_pi : expr → option expr
| (pi _ _ _ body) := some body
| _ := none

-- This is the recursive 'f'.
@[inline] meta def count_binders (match_binder : expr → option expr) : expr → ℕ := λ e,
match match_binder e with
| none := 0
| some body := count_binders body + 1
end

-- This is 'g', where 'f' should get inlined.
meta def count_pis : expr → ℕ := count_binders match_pi

This should ideally generate the same code as this hand-inlined version:

meta def count_pis' : expr → ℕ
| (pi _ _ _ body) := count_pis' body + 1
| _ := 0

/-
Optimisation steps to get there:

meta def count_pis := count_binders match_pi

==> [inline count_binders]

meta def tmp : expr → ℕ := λ e,
match match_pi e with
| none := 0
| some body := tmp body + 1
end

==> [inline match_pi and beta-reduce]

meta def tmp : expr → ℕ := λ e,
match (match e with | (pi _ _ _ body) := some body | _ := none end) with
| none := 0
| some body := tmp body + 1
end

==> [nested matches ("case-of-case")]

meta def tmp : expr → ℕ := λ e,
match e with
| (pi _ _ _ body) := tmp body + 1
| _ := 0
end
-/

Mario, your suggestion, if I understand it correctly, achieves this goal, but only for the first iteration of the loop:

meta def count_binders' (match_binder : expr → option expr) : expr → ℕ := λ e,
match match_binder e with
| none := 0
| some body := count_binders' body + 1
end

@[inline] meta def count_binders'' (match_binder : expr → option expr) : expr → ℕ := λ e,
match match_binder e with
| none := 0
| some body := count_binders' match_binder body + 1
end

meta def count_pis'' : expr → ℕ := count_binders'' match_pi

/-
Optimisation steps for count_pis'':

meta def count_pis'' := count_binders'' match_pi

==> [inline count_binders'']

meta def count_pis'' := λ e,
match match_pi e with
| none := 0
| some body := count_binders' match_binder body + 1
end

==> [inline match_binder and beta-reduce]

meta def count_pis'' := λ e,
match (match e with | (pi _ _ _ body) := some body | _ := none end) with
| none := 0
| some body := count_binders' match_pi body + 1
end

==> [nested matches]

meta def count_pis'' := λ e,
match e with
| (pi _ _ _ body) := count_binders' match_pi body + 1
| _ := 0
end
-/

All iterations except the first will still call count_binders', which in turn calls match_pi, so we've gained little.

[Reid Barton]

I'd call this specialization more than inlining

[Jannis Limperg]

Fair. But then I'd call specialisation a form of inlining. :P

[Reid Barton]

Well, you're specializing count_binders, presumably with the goal of next inlining match_pi (otherwise there's no gain)

[Jannis Limperg]

If this is the usual terminology, sure, I'll call it specialisation.

[Gabriel Ebner]

The inliner doesn't work across definitions generated by the equation compiler (match, |). Aside from that, the @[inline] does what you want. That is, after preprocessing, all references to definitions with @[inline] should be gone.

[Gabriel Ebner]

It should work if you write count_binders using option.elim or option.rec (because then it's a normal definition and can be substituted by the inliner).

[Jannis Limperg]

That's very interesting, thanks! Then I should be able to get it to do what I want.

[Jannis Limperg]

Followup @Gabriel Ebner: does Lean actually optimise nested matching? Because otherwise the inlining doesn't do too much anyway.

[Jannis Limperg]

Or rather, nested eliminator applications. Silly example:

meta def test (x : option ℕ) : ℕ :=
option.cases_on (@option.cases_on _ (λ _, option ℕ) x none some) 0 id

/- Should optimise to: -/

meta def test' (x : option ℕ) : ℕ :=
option.cases_on x 0 id

[Gabriel Ebner]

I'm afraid that there enough bugs there that you can't rely on it. set_option trace.compiler.preprocess true and set_option trace.compiler.optimize_bytecode true are your friends.

[Jannis Limperg]

Thanks again! When I rewrite everything without the equation compiler, inlining seems to work as intended. Unfortunately, my use case is in the tactic monad and I get a lot of separately-compiled lambdas. This blocks any followup simplification, so the inlining doesn't do much. I think I'll just live with the inefficiency.