Source file Generalization.ml

(******************************************************************************)
(*                                                                            *)
(*                                  Inferno                                   *)
(*                                                                            *)
(*                       François Pottier, Inria Paris                        *)
(*                                                                            *)
(*  Copyright Inria. All rights reserved. This file is distributed under the  *)
(*  terms of the MIT License, as described in the file LICENSE.               *)
(*                                                                            *)
(******************************************************************************)

open UnifierSig

module Make (S : STRUCTURE) = struct

(* -------------------------------------------------------------------------- *)

(* The [Generalization] module manages the [rank] fields of the unification
   variables, as well as a global notion of ``current rank'', stored in the
   field [state.young]. Ranks can be thought of as de Bruijn levels, in the
   following sense: whenever the left-hand side of a [CLet] constraint is
   entered, the current rank is incremented by one. Thus, the rank of a
   variable indicates where (i.e., at which [CLet] construct) this variable is
   (existentially) bound. *)
type rank = int
module Rank = struct
  type t = rank
  let merge r1 r2 : t = min r1 r2
  let equal (r1 : t) (r2 : t) = Int.equal r1 r2
  let base = 0
  let dummy = (-1)
  let is_dummy r = equal r dummy

  let () =
    (* we index into an array starting from the base rank,
       so we better have (base >= 0). *)
    assert (base >= 0)
end

module Data = struct
  type t = {
    mutable rank: int;
    mutable generic: bool;
  }

  let dummy () = { rank = Rank.dummy; generic = false }

  let merge d1 d2 =
    assert (not (Rank.is_dummy d1.rank));
    assert (not (Rank.is_dummy d2.rank));
    assert (not d1.generic);
    assert (not d2.generic);
    let rank = Rank.merge d1.rank d2.rank in
    { rank; generic = false; }
end

module U = Unifier.Make(S)(Data)

let fresh s =
  U.fresh s (Data.dummy ())

let rank (v : U.variable) : rank = (U.data v).Data.rank
let set_rank (v : U.variable) (r : rank) = (U.data v).Data.rank <- r
let adjust_rank (v : U.variable) (r : rank) =
  (* equivalent to [set_rank v (min r (rank v))] *)
  if r < rank v then set_rank v r

let is_generic (v : U.variable) : bool = (U.data v).Data.generic
let set_generic (v : U.variable) = (U.data v).Data.generic <- true

(* The rank of a variable is set to the current rank when the variable is
   first created. During the lifetime of a variable, its rank can only
   decrease. Decreasing a variable's rank amounts to hoisting out the
   existential quantifier that binds this variable. *)

(* Ranks are updated in a lazy manner. Only one rank maintenance operation
   takes place during unification: when two variables are unified, the rank of
   the merged variable is set to the minimum of the ranks of the two
   variables. (This operation is performed by the unifier.) Two other rank
   maintenance operations are performed here, namely downward propagation and
   upward propagation. Downward propagation updates a child's rank, based on
   its father rank; there is no need for a child's rank to exceed its father's
   rank. Upward propagation updates a father's rank, based the ranks of all of
   its children: there is no need for a father's rank to exceed the maximum of
   its children's ranks. These operations are performed at generalization time
   because it would be costly (and it is unnecessary) to perform them during
   unification. *)

(* The [rank] field maps every variable to the [CLet] construct where it is
   bound. Conversely, the [Generalization] module keeps track, for every
   active [CLet] construct, of a (complete) list of variables that are bound
   there. This takes the form of an array, stored in the field [state.pool].
   For every rank comprised between 1 and [state.young], both included, this
   array stores a list of the variables that are bound there. This array is
   again updated in a lazy manner, at generalization time. Because the unifier
   updates the ranks, but does not know about this array, the property that
   holds in general is: if a variable [v] has rank [i], then it appears in
   pool number [j], where [i <= j] holds. Immediately after generalization has
   been performed, the array has been updated, so [i = j] holds. *)

type state = {
  (* An array of pools (lists of variables), indexed by ranks. *)
  pool: U.variable list InfiniteArray.t;
  (* The current rank. *)
  mutable young: int;
}

(* -------------------------------------------------------------------------- *)

(* The [Generalization] module is in charge of constructing and instantiating
   type schemes, or graph fragments that contain universally quantified (i.e.,
   to-be-copied) variables as well as free (i.e., not-to-be-copied) variables.
   This happens when we exit the left-hand side of a [CLet] constraint, i.e.,
   when we move from a context of the form [let x v = <hole> in c] to a
   context of the form [let x = scheme in <hole>]. At this moment, the current
   rank [state.young] is decremented by one, and all variables whose rank was
   precisely [state.young] become universally quantified, or generic. These
   variables are no longer stored in any pool, as they are no longer
   existentially quantified. *)

(* The generic variables that have no structure are the ``quantifiers'' of the
   type scheme. A type scheme is internally represented as a pair of a root variable,
   and the rank at which generalization took place. In particular, the quantifiers
   can be reconstructed by traversing the root to find all structureless variable
   at the current level. Generic variables of a lower level are not part of the
   current scheme, they were generalized later during inference. *)
type scheme = { root: U.variable; rank : Rank.t; }

(* -------------------------------------------------------------------------- *)

(* The initial state. *)

(* The pool array is initially populated with empty pools. The rank is
   chosen so that the rank of its first child, the first rank that is
   actually exploited, is Rank.base. *)
let init () = {
  pool = InfiniteArray.make 8 [];
  young = Rank.base - 1;
}

(* -------------------------------------------------------------------------- *)

(* To get a "trivial" (fully monomorphic) type scheme,
   we pick a rank strictly above the root variable rank. *)
let trivial body =
  { root = body; rank = rank body + 1  }

let body { root; _ } = root

(* The quantifiers of a type scheme are exactly the generic
   structureless variables that are reachable from the root, at the
   scheme rank. *)
let quantifiers { root; rank = scheme_rank } =

  (* Prepare to mark which variables have been visited. *)
  let visited : unit U.VarMap.t = U.VarMap.create 128 in

  let rec traverse v quantifiers =

    (* If this variable is not generic or has been discovered already, then
       we must stop. *)

    if not (is_generic v)
    || not (Rank.equal (rank v) scheme_rank)
    || U.VarMap.mem visited v
    then
      quantifiers
    else begin

      (* Mark this variable as visited. If it carries no structure, then it is
         a leaf in the generic part of this type scheme, that is, a
         quantifier: add it to the list of quantifiers. Otherwise, traverse
         its descendants. Note that the variable must be marked before the
         recursive call, so as to guarantee termination in the presence of
         cyclic terms. *)

      U.VarMap.add visited v ();
      (* The order in which the quantifiers appear is determined in an arbitrary
         manner. *)
      match U.structure v with
      | None ->
          v :: quantifiers
      | Some t ->
          S.fold traverse t quantifiers

    end

  in
  traverse root []

(* -------------------------------------------------------------------------- *)

(* The internal function [register_at_rank] assumes that [v]'s rank is already
   a valid positive rank, and registers [v] by inserting it into the appropriate
   pool. *)

let register_at_rank ({ pool; _ } as state) v =
  let r = rank v in
  assert (Rank.base <= r && r <= state.young);
  InfiniteArray.set pool r (v :: InfiniteArray.get pool r)

(* The external function [register] assumes that [v]'s rank is uninitialized.
   It sets this rank to the current rank, [state.young], then registers [v]. *)

let registered v =
  not (Rank.is_dummy (rank v))

let register state v =
  assert (not (registered v));
  set_rank v state.young;
  register_at_rank state v

(* -------------------------------------------------------------------------- *)

(* Debugging utilities. *)

let show_variable v =
  Printf.printf "id = %d, rank = %d\n" (U.id v) (rank v)

let show_pool state k =
  Printf.printf "Pool %d:\n" k;
  List.iter show_variable (InfiniteArray.get state.pool k)

let show_young state =
  Printf.printf "state.young = %d\n" state.young

let show_pools state =
  for k = Rank.base to state.young do
    show_pool state k
  done

let show_state label state =
  Printf.printf "%s:\n" label;
  show_young state;
  show_pools state

(* -------------------------------------------------------------------------- *)

(* [enter] simply increments the current rank by one. The corresponding pool is
   in principle already empty. *)

let enter state =
  state.young <- state.young + 1;
  assert (InfiniteArray.get state.pool state.young = [])

(* -------------------------------------------------------------------------- *)

(* [exit] is where the moderately subtle generalization work takes place. *)

(* A data structure to represent all the variables created at the youngest level.
   (Since their creation some of them have been given an older level.)
*)
type variable_generation = {
  level: int;
  (* Current level. *)

  list: U.variable list;
  (* The list [vs] of all variables in the young generation. *)

  table: unit U.VarMap.t;
  (* This hash table stores all of these variables, so that we may check
     membership in the young generation in constant time. *)

  by_rank: U.variable list array;
  (* This array stores all of these variables, indexed by rank. The use
     of a bucket sort is theoretically costly if the [CLet]-nesting depth
     is not considered a constant, because of the need to walk through
     possibly-empty buckets; in that case, a standard sort algorithm, or
     (even better) no sort at all would suffice. (Sorting helps us compute
     better ranks; but distinguishing between [young] and non-[young] would
     be enough.) In practice, the [CLet]-nesting depth should remain low,
     and walking through empty buckets (in the loop that follows) should
     cost almost nothing. So we adopt this approach, even though it violates
     the complexity claim of the paper. *)

  is_young: U.variable -> bool;
  (* A membership test for the young generation. *)
}

let young_variable_generation state =
  (* Get the list [vs] of all variables in the young generation. *)
  let vs = InfiniteArray.get state.pool state.young in

  let table = U.VarMap.create 128 in

  let by_rank = Array.make (state.young + 1) [] in

  (* Initialize the by_rank array *)
  List.iter (fun v ->
    U.VarMap.add table v ();
    let r = rank v in
    assert (Rank.base <= r && r <= state.young);
    by_rank.(r) <- v :: by_rank.(r)
  ) vs;

  let is_young v =
    U.VarMap.mem table v
  in

  { level = state.young; list = vs; table; by_rank; is_young }

let update_generation_ranks young_generation =
  (* Now, update the rank of every variable in the young generation. Downward
     propagation and upward propagation, as described above, are performed. A
     single depth-first traversal of the young generation achieves
     both. Roughly speaking, downward propagation is achieved on the way down,
     while upward propagation is achieved on the way up. (In reality, all rank
     updates takes place during the upward phase.)

     It may be worth noting that downward propagation is required, as (for
     instance) [instantiate] assumes that a non-generic variable cannot have
     generic children. Upward propagation is an optional optimization; without
     it, we would perform slightly more copying, but that would be harmless.

     During each traversal, every visited variable is marked as such, so as to
     avoid being visited again. To ensure that visiting every variable once is
     enough, the roots must be processed by increasing order of rank. In the
     absence of cycles, this enforces the following invariant: when performing
     a traversal whose starting point has rank [k], every variable marked as
     visited has rank [k] or less already. (In the presence of cycles, this
     algorithm is incomplete and may compute ranks that are slightly higher
     than necessary.) Conversely, every non-visited variable must have rank
     greater than or equal to [k]. This explains why [k] remains constant as
     we go down (i.e., discovering [v] does not improve the value of [k] that
     we are pushing down). *)

  let visited : unit U.VarMap.t = U.VarMap.create 128 in

  for k = Rank.base to young_generation.level do

    (* A postcondition of [traverse v] is [U.rank v <= k]. (This is downward
       propagation.) *)
    let rec traverse v =
      assert (rank v >= Rank.base);
      (* If [v] was visited before, then its rank must be below [k], as we
         adjust ranks on the way down already. *)
      if U.VarMap.mem visited v then
        assert (rank v <= k)
      else begin
        (* Otherwise, immediately mark it as visited, and immediately adjust
           its rank so as to be at most [k]. (This is important if cyclic
           graphs are allowed.) *)
        U.VarMap.add visited v ();
        adjust_rank v k;
        (* If [v] is part of the young generation, and if it has structure,
           then traverse its children (updating their ranks) and on the way
           back up, adjust [v]'s rank again (this is upward propagation). If
           [v] has structure but no children, then it is a constant, and it
           receives the base rank; it will be moved to the oldest pool. If
           [v] has no structure, do nothing; it would be wrong to move its
           rank down to the base rank. *)
        if young_generation.is_young v then begin
          (* The rank of this variable can't have been below [k], because
             it is young but was not visited yet. Thus, it must have been
             at or above [k], and since we have just adjusted it, it must
             now be [k]. *)
          assert (rank v = k);
          Option.iter (fun t ->
            adjust_rank v (
              S.fold (fun child accu ->
                traverse child;
                max (rank child) accu
              ) t Rank.base (* the base rank is neutral for [max] *)
            )
          ) (U.structure v)
        end
        (* If [v] is old, stop. *)
        else
          assert (rank v < young_generation.level)
      end

    in
    List.iter traverse young_generation.by_rank.(k)

  done

let generalize state young_generation : U.variable list =
  (* Every variable that has become an alias for some other (old or young)
     variable is dropped. We keep only one representative of each class.

     Every variable whose rank has become strictly less than [young] may be
     safely turned into an old variable. It is moved into the pool that
     corresponds to its rank.

     Every variable whose rank is still [young] must be generalized. That is,
     it becomes universally quantified in the type scheme that is being
     created. We set its rank to [generic]. By convention, a variable of rank
     [generic] is considered universally quantified. *)
  let generalizable =
    List.filter (fun v ->
      U.is_representative v && begin
        if rank v < state.young then begin
          register_at_rank state v;
          false
        end
        else begin
          assert (rank v = state.young);
          set_generic v;
          U.structure v = None
        end
      end
    ) young_generation.list
  in

  (* Update the state by emptying the current pool. *)
  InfiniteArray.set state.pool state.young [];

  (* The generic variables are now unreachable from the variables that still
     have positive rank and inhabit one of the pools. *)
  assert (
    (* For every [v] in the young generation, *)
    U.VarMap.fold (fun v () ok ->
      ok && (
        (* If [v] is not generic, *)
        is_generic v ||
        match U.structure v with
        | None ->
            true
        | Some t ->
            (* then its child [w] is not generic. *)
            S.fold (fun w ok -> ok && not (is_generic w)) t true
      )
    ) young_generation.table true
  );
  generalizable

let exit ~rectypes state roots =
  let young_generation = young_variable_generation state in

  (* If the client would like us to detect and rule out recursive types, then
     now is the time to perform an occurs check over the young generation. *)
  if not rectypes then
    List.iter (U.new_occurs_check young_generation.is_young)
      young_generation.list;

  (* Determine the rank of every variable in the young generation as
     precisely as possible. *)
  update_generation_ranks young_generation;

  (* Move the old variables in their corresponding pool and make the
     young one generic. *)
  let generalizable = generalize state young_generation in

  (* Exit the current inference level. *)
  state.young <- state.young - 1;

  let make_scheme root = { root; rank = young_generation.level; } in

  (* Return the list of unique generalizable variables that was constructed
     above, and a list of type schemes, obtained from the list [roots]. *)
  generalizable,
  List.map make_scheme roots

(* -------------------------------------------------------------------------- *)

(* Instantiation amounts to copying a fragment of a graph. The fragment that
   must be copied is exactly the prefix which satisfies [is_generic]. *)

let instantiate state { root; rank = scheme_rank } =

  (* Prepare to mark which variables have been visited and record their copy. *)
  let visited : U.variable U.VarMap.t = U.VarMap.create 128 in

  let quantifier_instances = ref [] in

  (* If the variable [v] has rank [generic], then [copy v] returns a copy of
     it, and copies its descendants recursively. If [v] has positive rank,
     then [copy v] returns [v]. Only one copy per variable is created, even if
     a variable is encountered several times during the traversal. *)
  let rec copy v =

    (* If this variable has positive rank, then it is not generic: we must
       stop. *)

    if not (is_generic v) then
      v

    (* If a copy of this variable has been created already, return it. *)

    else begin
      try
        U.VarMap.find visited v
      with Not_found ->

        (* The variable must be copied, and has not been copied yet. Create a
           new variable, register it, and update the mapping. Then, copy its
           descendants. Note that the mapping must be updated before making a
           recursive call to [copy], so as to guarantee termination in the
           presence of cyclic terms. *)

        let v' =
          let data = { Data.rank = state.young; Data.generic = false } in
          U.fresh None data in
        register_at_rank state v';
        U.VarMap.add visited v v';
        begin match U.structure v with
          | None ->
            (* When inferring Core ML terms, the condition on ranks below is
               not necessary: a term (let x = t in u) will only ever generate
               instance constraints for (x) in the constraint elaborated for (u).

               These instance constraints always take place at the level right below
               the level of (t), before this level has been generalized. At this point,
               no lowever-level variables could ever be generic.

               The condition is necessary to support other language features,
               such as various forms of first-class polymorphism:
                 let p = fun y ->
                   { id : 'a. 'a -> 'a * 'b = fun x -> (x, y) }
                 in
                 p.id ()
               In this example, the use of
                  p.id
               at level 0 ccreates an instance constraint on a polymorphic scheme
               ('a . 'a -> 'a * 'b) coming from level 2 { id = ... }, and at this point
               the level 1 (let p = ...) also contains generic variables: if we
               collected all generic structureless variables, we would instead
               produce the scheme ('a 'b. 'a -> 'a * 'b), which is incorrect.
            *)
            if Rank.equal (rank v) scheme_rank then
              quantifier_instances := v' :: !quantifier_instances;
          | Some s ->
            U.set_structure v' (Some (S.map copy s));
        end;
        v'

    end
  in

  (* [Gabriel]: we are assuming here that quantifiers will be computed in the same order in this function
     and in 'quantifiers' above. This is fairly optimistic, as it relies on two facts:
     - the iteration order of the provided S.map and S.fold are the same
     - the 'copy' traversal traverses more than 'quantifiers', as it also descendes in
       generic nodes at lower level, but this does not affect the traversal order of
       current-level nodes as older nodes can never reach current-level nodes

     Caching the list of quantifiers may be a more robust design, but here we really wanted to verify
     our idea that the generic-with-levels design allows for after-the-fact computation of quantifiers.
  *)
  let instance = copy root in
  !quantifier_instances, instance
end