I am trying to get flambda to optimize code that uses what is essentially a simple state monad, with `let*`

and `let+`

operators, to the code that can be written by breaking the abstraction and directly writing the code in state-passing style. I have not been able to find attributes or compilation flags that enable the abstraction-preserving version to compile to close to as good code as the abstraction-breaking version, and wonder if anyone has ideas.

Here is a concrete example:

```
module State = struct
type s
type 'a m = s -> 'a * s
let return a s = (a, s)
let bind m k s =
let a, s = m s in
k a s
let map f m = bind m (fun a -> return (f a))
module Import = struct
let ( let* ) = bind
let ( let+ ) m f = map f m
end
let read s = (s, s)
let run s m = m s
let mfold ps a ~f s =
let f (a, s) p = f p a s in
List.fold_left f (a, s) ps
end
open State.Import
type t =
| True
| Not of t
| And of t list
| Or of t list
| If of {cnd: t; pos: t; neg: t}
let iter_dnf1 ~meet ~top fml ~f =
let rec add_conjunct fml (cjn, splits) =
match fml with
| True | Not _ ->
let+ cjn = meet fml cjn in
(cjn, splits)
| And ps -> State.mfold ~f:add_conjunct ps (cjn, splits)
| Or ps -> State.return (cjn, ps :: splits)
| If {cnd; pos; neg} ->
add_conjunct (Or [And [cnd; pos]; And [Not cnd; neg]]) (cjn, splits)
in
let rec add_disjunct (cjn, splits) fml =
let* cjn, splits = add_conjunct fml (cjn, splits) in
let+ s = State.read in
match splits with
| ps :: splits ->
List.iter
(fun fml -> fst (State.run s (add_disjunct (cjn, splits) fml)))
ps
| [] -> f (cjn, s)
in
add_disjunct (top, []) fml
let dnf1 ~meet ~top fml =
let+ s = State.read in
fun f -> fst (State.run s (iter_dnf1 ~meet ~top fml ~f))
let iter_dnf2 ~meet ~top fml ~f s =
let rec add_conjunct fml (cjn, splits) s =
match fml with
| True | Not _ ->
let cjn, s = meet fml cjn s in
((cjn, splits), s)
| And ps -> State.mfold ~f:add_conjunct ps (cjn, splits) s
| Or ps -> ((cjn, ps :: splits), s)
| If {cnd; pos; neg} ->
add_conjunct
(Or [And [cnd; pos]; And [Not cnd; neg]])
(cjn, splits) s
in
let rec add_disjunct (cjn, splits) fml s =
let (cjn, splits), s = add_conjunct fml (cjn, splits) s in
match splits with
| ps :: splits ->
(List.iter (fun fml -> fst (add_disjunct (cjn, splits) fml s)) ps, s)
| [] -> (f (cjn, s), s)
in
add_disjunct (top, []) fml s
let dnf2 ~meet ~top fml s =
((fun f -> fst (iter_dnf2 ~meet ~top fml ~f s)), s)
```

The `iter_dnf1`

and `dnf1`

functions preserve the abstraction while the `iter_dnf2`

and `dnf2`

functions break it, but are otherwise equivalent. The code generated for the latter version is rather better, with the former in particular allocating many closures. See this example at Compiler Explorer. That is for 4.12, and the relative situation is similar with 4.14.

It seems difficult for flambda to eliminate the closures for the anonymous functions that result from the elaboration of the `let*`

and `let+`

operators, but maybe something else is going on.

Another transformation that seems difficult is needed in the `add_conjunct`

function. The bodies of some cases need to be eta-expanded with the state argument, and then `fun`

needs to be hoisted over `match`

, such as transforming:

```
match e with A -> fun s -> a | B -> fun s -> b | C -> fun s -> c
```

to

```
fun s -> match e with A -> a | B -> b | C -> c
```

Any suggestions would be much appreciated!