[@@inline always] attribute not working for monadic function

Learning
,
1

I have some simple monadic bindings, given below:

module Syntax = struct
  (* For ppx_let *)
  module Let_syntax = struct 
    let bind m ~f = fun sin ->
      let sout, res = m sin in
      f res sout
    [@@inline always]

    let return = return

    let map m ~f = fun sin ->
      let sout, res = m sin in
      (sout, f res)
    [@@inline always]
      
    let both m1 m2 = fun sin ->
      let s1, res1 = m1 sin in
      let s2, res2 = m2 s1 in
      (s2, (res1, res2))
    [@@inline always]
  end
    
  open Let_syntax
  
  let (let+) (m: 'c state -> 'c state * 'a) (f: 'a -> 'b) : ('c state -> 'c state * 'b) = map m ~f [@@inline always]
  let (and+) = both
  let (let* ) (m: 'c state -> 'c state * 'a) (f: 'a -> 'c state -> 'c state * 'b) : ('c state -> 'c state * 'b) = bind m ~f [@@inline always]
  let (and* ) = both
  
end

The following are the contents of the dune file at project root:

(env
  (dev
    (flags (:standard -warn-error -A -w -39 -w -27 -w -32))
    (ocamlopt_flags (:standard -O3 -inline 100 -inlining-report))
    ))

I have flambda installed.

However, it seems that the bind and map functions are not getting inlined. For instance, in cases of unexpected errors, the backtrace is full of references to these bind and map functions like so:

Raised at Util__Error.fail in file "lib/util/error.ml", line 10, characters 24-51
Called from Ast__Rewriter.resolve_and_find in file "lib/ast/rewriter.ml", line 811, characters 52-95
Called from Ast__Rewriter.Syntax.Let_syntax.map in file "lib/ast/rewriter.ml", line 26, characters 13-18
Called from Ast__Rewriter.Syntax.Let_syntax.map in file "lib/ast/rewriter.ml", line 25, characters 22-27
Called from Ast__Rewriter.Syntax.Let_syntax.map in file "lib/ast/rewriter.ml", line 25, characters 22-27
Called from Ast__Rewriter.Syntax.Let_syntax.bind in file "lib/ast/rewriter.ml", line 19, characters 22-27

Is this the expected behaviour? Is there a reason why these functions are not getting inlined? How can I go about troubleshooting this?

1 Like
2

It is kind of expected for most versions of OCaml, but the new 5.2 version should behave much better.
The issue is that bind, for example, is seen as taking three arguments. The anonymous function taking sin as parameter is merged with bind itself, so every call to f will happen under bind.
Merging functions is how OCaml handles currified functions efficiently, so it is rather important, but recently we have found a way to take into account the syntactic hints from the programmer to produce more reasonable code in cases similar to the one you showed.
If you can, I would suggest testing the 5.2 alpha release and see if you still have the same issue. On older versions you can also insert dummy let () = () in bindings between the parameters and the anonymous function, sometimes it is enough to trick the compiler into keeping the split between the functions.

1 Like
3

Do you think if I rewrite it to make sin an explicit parameter of map and bind, or other such tricks it might make a difference?

4

In versions of OCaml up to 5.1, let bind m ~f = fun sin -> ... is equivalent to let bind m ~f sin = .... And that definition looks consistent to me with what you’re observing.

But it looks like you might be hitting an issue that has nothing to do with inlining or merging functions. Take the following example:

let foo arg =
  let* x = <def_x> in
  let* y = <def_y> in
  <body>

This compiles (flambda or not) to:

let foo arg =
  bind <def_x> (fun x ->
    bind <def_y> (fun y -> <body>))

Inlining bind (assuming flambda), you can get:

let foo arg =
  fun sin_1 ->
    let m_1 = <def_x> in
    let f_1 = fun x ->
      fun sin_2 ->
        let m_2 = <def_y> in
        let f_2 = fun y -> <body> in
        let sout_2, res_2 = m_2 sin_2 in
        f_2 res_2 sout_2
    in
    let sout_1, res_1 = m_1 sin_1 in
    f_1 res_1 sout_1

This will be simplified further, as the functions f_1 and f_2 at least can be inlined, but there is already a major issue: the main piece of code (after fun sin_1) is considered to be part of the body of bind, not foo.
With OCaml 5.2, there is a small improvement because instead of bind it will be reported as part of the anonymous function returned by bind, but it not part of foo anymore.
In fact, as far as the compiler is concerned foo is a small function that either immediately tail-calls bind (if there is no inlining) or allocates a closure and immediately returns it (if bind is inlined). You will never see foo itself in the backtrace because it doesn’t really perform any work itself.

1 Like