Getting a recursive function to typecheck

Hello,

I’m fought for a while with a piece of code that would not typecheck, and I’m curious if there’s an obvious way to get it to typecheck.

Take this code:

type tree =
  | Leaf of int
  | Node of (int * tree * tree)

let rec f (g : int -> 'a) tree : 'a =
   match tree with
  | Leaf i -> g i
  | Node (i, l, r) -> if i >= 0 then f g l else f g r

let x = f (fun i -> "s") (Leaf 0)
let y = f (fun i -> i) (Leaf 0)

Running this in the playground gives us, as expected,

type tree = Leaf of int | Node of (int * tree * tree)
val f : (int -> 'a) -> tree -> 'a = <fun>
val x : string = "s"
val y : int = 0

However, let’s say that I want f to call itself with a different g than the one passed as argument temporarily. For example:

type tree =
  | Leaf of int
  | Node of (int * tree * tree)

let rec f (g : int -> 'a) tree : 'a =
   match tree with
  | Leaf i -> g i
  | Node (i, l, r) ->
      if (i == 0) then
        let g' i = Leaf i in
        let ret = f g' (Node (12, Leaf(42), Leaf (24))) in
        f g ret
      else
      if i >= 0 then f g l else f g r

let x = f (fun i -> "s") (Leaf 0)
let y = f (fun i -> i) (Leaf 0)

Suddenly, this does not type-check anymore, since it infers the following, which is quite inconvenient:

val f : (int -> tree) -> tree -> tree = <fun>

Suddenly, the g argument must return a tree. Is there a way to get this function to keep inferring that g : int -> 'a ?

In my current code (which is much more complex, but let’s simplify here), my g function is of type
int -> ('a, int) and I only need the int in my internal recursive call, so I create a function g that returns ((), the_thing_I_need), and then discard the unit, and it fails to typecheck.
I did manage to get it to work by doing let g i = (Obj.magic (), the_thing_I_need) but it’s not pretty and I’m not even 100% sure it’s safe.

Writing the first line of f like this typechecks:

let rec f: type a. (int -> a) -> tree -> a = fun g tree ->

EDIT: Took me a while to find the documentation for this syntax: OCaml - Language extensions.

This also works (more similar to the linked manual section below):

let rec f: 'a. (int -> 'a) -> tree -> 'a = fun g tree ->

See OCaml - Polymorphism and its limitations for a longer explanation of the issue.

@sim642 's solution works, thank you!
I had tried

let rec f (type a) g tree = ...

which is thought was vaguely equivalent but that didn’t work

@vlaviron Thank you, I had read this also, but somehow missed that it did give the solution