Source file Solver.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 Signatures

module Make
  (X : TEVAR)
  (S : HSTRUCTURE)
  (O : OUTPUT with type 'a structure = 'a S.structure)
= struct

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

(* The type [tevar] of term variables is provided by [X]. *)

type tevar =
  X.t

type tevars =
  tevar list

module TeVarMap =
  Hashtbl.Make(X)

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

(* The type variables that appear in constraints are immutable: they
   are integer identifiers. *)

type variable =
  int

type variables =
  variable list

let fresh : unit -> variable =
  Utils.gensym()

module VarTable = Hashtbl

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

(* The syntax of constraints. *)

type range =
  Lexing.position * Lexing.position

type shallow_ty =
  variable O.structure

type tyvars =
  O.tyvar list

type tys =
  O.ty list

type scheme =
  tyvars * O.ty

type schemes =
  scheme list

(**A constraint of type ['a co] is a constraint whose resolution, if
   successful, produces a value of type ['a]. The ability for the
   resolution process to produce a result can be used to express
   elaboration in an elegant way. *)
type _ co =

| CRange : range * 'a co -> 'a co
    (**The constraint [CRange (range, c)] is equivalent to the constraint [c],
       and is annotated with [range], a range of positions in some piece of
       source code. *)

| CTrue : unit co
    (**The trivial constraint [CTrue] is always true. *)

| CMap : 'a co * ('a -> 'b) -> 'b co
    (**The constraint [CMap (c, f)] is equivalent to the constraint [c]. The
       transformation function [f] is used to postprocess the result of solving
       this constraint. *)

| CPure : 'a -> 'a co
    (**The constraint [CPure x] is always true, and produces the value [x].
       It is equivalent to [CMap (CTrue, fun () -> x)]. *)

| CConj : 'a co * 'b co -> ('a * 'b) co
    (**[CConj (c1, c2)] is the conjunction of the constraints [c1] and [c2]. *)

| CEq : variable * variable -> unit co
    (**[CEq (v1, v2)] is an equality constraint between the type variables [v1]
       and [v2]. *)

| CExist : variable * variable S.structure option * 'a co -> 'a co
    (**An existential quantification [CExist (v, so, c)] binds the type
       variable [v] in the constraint [c]. This type variable carries an
       optional equality constraint [so] whose right-hand side is a shallow
       term. *)

| CDecode : variable -> O.ty co
    (**The constraint [CDecode v] is logically equivalent to
       a [true] constraint, so it is always satisfied. Its result is
       a decoded type which represents the value assigned to the
       type variable [v] at the end of the resolution process. *)

| CInstance : tevar * variable -> tys co
    (**An instantiation constraint [CInstance (x, v)] requires the type
       variable [v] to be an instance of the type scheme denoted by the
       term variable [x]. Its result is a list of types that indicates
       how the type scheme was instantiated. *)

| CDef : tevar * variable * 'a co -> 'a co
    (**The constraint [CDef (x, v, c)] binds the term variable [x] to
       the trivial (monomorphic) type scheme [v] in the constraint [c]. *)

| CLet :
    variables * tevars * variables * 'a co * 'b co ->
    (tyvars * schemes * 'a * 'b) co
    (**The constraint [CLet rs xs vs c1 c2] is solved as follows:
       - First, the constraint [∀rs.∃vs.c1] is solved.
       - Then, for each pair [(x, v)] drawn from the lists [xs] and [vs]
         (which must have the same length), the term variable [x] is
         bound to the constraint abstraction [λv.∃rs.∃(vs\v).c1].
       - Finally, under these bindings, the constraint [c2] is solved.

       Thus, the variables [rs] are initially regarded as rigid while [c1]
       is solved and simplified; but, once this is done, their rigid
       nature disappears and they become flexible before generalization
       is performed. The scheme associated with the term variable [x]
       has the corresponding type variable [v] as its root.

       The semantic value of this constraint is a tuple of:
       - a list [γs] of decoded type variables that may appear in the semantic
         value [v1].
       - a list of decoded schemes, in correspondence with the list [xs].
       - the semantic value [v1] of the constraint [c1].
       - the semantic value [v2] of the constraint [c2]. *)

(* -------------------------------------------------------------------------- *)
(* Exceptions. *)

exception Unbound of range * tevar
exception Unify of range * O.ty * O.ty
exception Cycle of range * O.ty
exception VariableScopeEscape of range

(* A pretty-printer for constraints, used while debugging. *)

module Printer = struct

  [@@@warning "-4"] (* allow fragile pattern matching *)

  open PPrint

  let string_of_var v =
    string_of_int v

  let var v =
    string (string_of_var v)

  let tevar x =
    string (X.to_string x)

  let annot (x, v) =
    tevar x ^^ string " = " ^^ var v

  let shallow_equality s =
    string " = " ^^ S.pprint var s

  (* The grammar of constraints is organized in several layers: [binders],
     [conj], [simple]. The [entry] layer is a synonym for [binders] and is
     the entry point. *)

  (* [CRange] and [CMap] constructors are ignored by this printer. *)

  let rec entry : type a . a co -> document = fun c ->
    binders c

  and binders : type a . a co -> document = fun c ->
    let next = conj in
    let self = binders in
    group @@ match c with
    | CRange (_, c) ->
        self c
    | CMap (c, _) ->
        self c
    | CExist (v, so, c) ->
        string "exi " ^^
        var v ^^
        optional shallow_equality so ^^
        string " in" ^/^
        self c
    | CDef (x, v, c) ->
        string "def " ^^
        annot (x, v) ^^
        string " in" ^/^
        self c
    | CLet (rs, xs, vs, c1, c2) ->
        assert (List.length xs = List.length vs);
        let xvs = List.combine xs vs in
        string "let" ^^
        (if rs = [] then empty else brackets (separate_map space var rs)) ^^
        separate_map (string " and ") annot xvs ^^
        string " where" ^^ group (nest 2 (break 1 ^^
        next c1 ^^
        string " in")) ^/^
        self c2
    | _ ->
        next c

  and conj : type a . a co -> document = fun c ->
    let self = conj in
    let next = simple in
    group @@ match c with
    | CRange (_, c) ->
        self c
    | CMap (c, _) ->
        self c
    | CConj (c1, c2) ->
        self c1 ^^ string " *" ^/^ self c2
    | _ ->
        next c

  and simple : type a . a co -> document = fun c ->
    let self = simple in
    group @@ match c with
    | CRange (_, c) ->
        self c
    | CMap (c, _) ->
        self c
    | CTrue | CPure _ ->
        string "True"
    | CDecode v ->
        separate space [string "wit"; var v]
    | CEq (v1, v2) ->
        separate space [var v1; string "="; var v2]
    | CInstance (x, v) ->
        tevar x ^^ utf8string " ≤ " ^^ var v
    | CExist _
    | CDef _
    | CLet _
    | CConj _
      ->
        (* Introduce parentheses. *)
        surround 2 0 lparen (entry c) rparen

end (* Printer *)

let pprint =
  Printer.entry

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

(* The functor [Solve] initializes the solver's internal mutable state and
   runs the solver. This functor is later re-packaged as a function [solve]. *)

module Solve (C : sig
  (**The flag [rectypes] indicates whether equirecursive types are permitted. *)
  val rectypes: bool
end) = struct
open C

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

(* Initialize the generalization engine [G], which itself initializes the
   unifier [U]. *)

(* TODO explain why we use [Structure.Option] *)

module OS = Structure.Option(S)
module G = Generalization.Make(OS)
module U = G.U

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

(* The solver carries an environment, an immutable mapping of term variables
   to type schemes. *)

module Env = struct

  let table =
    TeVarMap.create 8

  let lookup range x =
    try
      TeVarMap.find table x
    with Not_found ->
      raise (Unbound (range, x))

  let bind x s =
    (* Shadowing allowed. *)
    TeVarMap.add table x s

  let unbind x =
    TeVarMap.remove table x

end (* Env *)

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

(* The solver maintains a mutable mapping of immutable type variables to
   unifier variables. *)

(* [uvar] looks up the mapping. [flexible] and [rigid] create a new unifier
   variable and extend the mapping. [uunbind] removes a mapping. *)

module UVar = struct

  let table =
    VarTable.create 8

  let uvar (v : variable) : U.variable =
    try
      VarTable.find table v
    with Not_found ->
      (* This should never happen. *)
      Printf.ksprintf failwith "Unbound variable %d" v

  let flexible (v : variable) so =
    assert (not (VarTable.mem table v));
    let uv = G.flexible (OS.map uvar so) in
    VarTable.add table v uv;
    uv

  let rigid (v : variable) =
    assert (not (VarTable.mem table v));
    let uv = G.rigid() in
    VarTable.add table v uv;
    uv

  let uunbind v =
    VarTable.remove table v

end (* UVar *)

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

(* Set up type decoding. *)

(* A type decoder is a function that transforms a unifier variable into an
   output type. We choose to decode types in an eager manner; that is, we
   take care of invoking the decoder, so that the client never needs to
   perform this task. As a result, we do not even need to expose the decoder
   to the client (although we could, if desired). *)

(* Because [O.ty] is a nominal representation of types, a type is decoded in
   the same way, regardless of how many type binders we have entered. This
   makes it possible for a decoder to exploit memoization. Thanks to this
   property, the type decoding process requires only linear time and space
   (in the acyclic case), regardless of how many calls to the decoder are
   performed. *)

(* TODO explain *)

module D =
  Decoder.Make (G.Data) (U) (struct
    include O
    let structure (s : O.ty U.structure) : O.ty =
      O.structure (OS.project_nonleaf (G.Data.project s))
  end)

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

(* The exceptions [Unify] and [Cycle], raised by the unifier, must be caught
   and re-raised in a slightly different format.

   A cyclic decoder is used even if [rectypes] is [false]. Indeed, recursive
   types can appear before the occurs check has been performed. *)

let unify range v1 v2 =
  try
    U.unify v1 v2
  with
  | U.Unify (v1, v2) ->
     let decode = D.new_cyclic_decoder() in
     raise (Unify (range, decode v1, decode v2))
  | G.VariableScopeEscape ->
     raise (VariableScopeEscape range)

let exit range ~rectypes vs =
  try
    G.exit ~rectypes vs
  with
  | G.Cycle v ->
     let decode = D.new_cyclic_decoder() in
     raise (Cycle (range, decode v))
  | G.VariableScopeEscape ->
     raise (VariableScopeEscape range)

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

(* Initialize a decoder for use during elaboration. If [rectypes] is on, a
   cyclic decoder is used; otherwise, an acyclic decoder is used. *)

let decode : U.variable -> O.ty =
  if rectypes then D.new_cyclic_decoder() else D.new_acyclic_decoder()

let decode_scheme (s : G.scheme) : scheme =
  List.map D.decode_variable (G.quantifiers s),
  decode (G.body s)

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

(* The toplevel constraint that is passed to the solver must have been
   constructed by [let0], otherwise we risk encountering type variables that
   we cannot register. Indeed, [G.register] must not be called unless
   [G.enter] has been invoked first.

   (If we accepted an arbitrary constraint from the user and silently
   wrapped it in [let0], what would we do with the toplevel quantifiers?) *)

(* The function [ok] determines whether a constraint is a suitable argument
   to the function [solve]. It is used only inside an assertion. *)

let rec ok : type a . a co -> bool =
  function
  | CRange (_, c) ->
      ok c
  | CTrue | CPure _ ->
      true
  | CMap (c, _) ->
      ok c
  | CLet (_rs, _xs, _vs, _c1, c2) ->
      (* The left-hand constraint [c1] does not need to be [ok], since it is
         examined after a call to [G.enter]. *)
      ok c2
  | CConj (c1, c2) ->
      ok c1 && ok c2
  | CEq _
  | CExist _
  | CDecode _
  | CInstance _
  | CDef _ ->
      (* These forms are not [ok], as they involve (free or binding
         occurrences of) type variables. *)
      false

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

(* A value of type ['a on_sol] is a delayed computation that can be run
   only after the solver has succeeded in finding a global solution. It
   produces a result of type ['a]. Defunctionalizing the solver would
   reveal that a value of type ['a on_sol] is in fact a tree of closures
   and that its evaluation is performed bottom-up. *)
type 'a on_sol =
  On_sol of (unit -> 'a) [@@unboxed]

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

(* We wish to solve a constraint of type [a co], producing a
   result of type [a]. We proceed in two phases:

   - Resolution: the constraint is traversed and solved;

   - Elaboration: if the resolution was successful, then a global solution
     exists, and a result can be computed (bottom-up) out of it. *)

(* The recursive solution [solve] implements the first phase. It returns a
   closure of type ['a on_sol], which implements the second phase: it
   describes how to compute a result once resolution is successful.

   [solve] is parameterized with an environment (which maps term variables to
   type schemes) and with a range (the range annotation that was most recently
   encountered on the way down). *)

open UVar

let rec solve : type a . range -> a co -> a on_sol =
  fun range c -> match c with
  | CRange (range, c) ->
      solve range c
  | CTrue ->
      On_sol (fun () -> ())
  | CPure x ->
      On_sol (fun () -> x)
  | CMap (c, f) ->
      let (On_sol r) = solve range c in
      On_sol (fun () -> f (r ()))
  | CConj (c1, c2) ->
      let (On_sol r1) = solve range c1 in
      let (On_sol r2) = solve range c2 in
      On_sol (fun () ->
        (* Even though we recommend using semantic actions without side
           effects, it seems wise to impose left-to-right evaluation of
           the semantic actions. *)
        let a1 = r1() in
        let a2 = r2() in
        a1, a2
      )
  | CEq (v, w) ->
      unify range (uvar v) (uvar w);
      On_sol (fun () -> ())
  | CExist (v, s, c) ->
      ignore (flexible v s);
      let result = solve range c in
      uunbind v;
      result
  | CDecode v ->
      let uv = uvar v in
      On_sol (fun () -> decode uv)
  | CInstance (x, w) ->
      (* The environment provides a type scheme for [x]. *)
      let s = Env.lookup range x in
      (* Instantiating this type scheme yields a variable [v], which we
         unify with [w]. It also yields a list of witnesses, which we
         record, as they will be useful during the decoding phase. *)
      let witnesses, v = G.instantiate s in
      unify range v (uvar w);
      On_sol (fun () -> List.map decode witnesses)
  | CDef (x, v, c) ->
      Env.bind x (G.trivial (uvar v));
      let a = solve range c in
      Env.unbind x;
      a
  | CLet (rs, xs, vs, c1, c2) ->
      (* Warn the generalization engine that we are entering the left-hand
         side of a [let] construct. *)
      G.enter();
      (* Register the rigid prefix [rs] with the generalization engine. *)
      let urs = List.map (fun v -> rigid v) rs in
      (* Register the variables [vs] with the generalization engine, just as
         if they were existentially bound in [c1]. This is what they are,
         basically, but they also serve as named entry points. *)
      let uvs = List.map (fun v -> flexible v None) vs in
      (* Solve the constraint [c1]. *)
      let (On_sol r1) = solve range c1 in
      (* Ask the generalization engine to perform an occurs check, to adjust
         the ranks of the type variables in the young generation (i.e., all
         of the type variables that were registered since the call to
         [G.enter] above), and to construct a list [ss] of type schemes for
         our entry points. The generalization engine also produces a list
         [gammas] of the young variables that should be universally
         quantified here. *)
      let gammas, ss = exit range ~rectypes uvs in
      (* All rigid variables of [rs] must be generalizable. This assertion
         may be costly and should be removed or disabled in the future. *)
      assert (urs |> List.for_all (fun r -> List.mem r gammas));
      List.iter uunbind vs;
      (* Extend the environment [env]. *)
      List.iter2 Env.bind xs ss;
      (* Proceed to solve [c2] in the extended environment. *)
      let (On_sol r2) = solve range c2 in
      List.iter Env.unbind xs;
      On_sol (fun () ->
        List.map D.decode_variable gammas,
        List.map decode_scheme ss,
        r1(),
        r2()
      )

(* The solver's state is now ready. The following function (which must be
   called at most once, as it affects the solver's state) combines solving
   and elaboration. *)

let main (type a) (c : a co) : a =

  (* Check that [c] is a well-formed toplevel constraint. If it is not,
     then the user is at fault. *)
  if not (ok c) then
    invalid_arg "solve: ill-formed toplevel constraint";

  let range = Lexing.(dummy_pos, dummy_pos) in
  (* Phase 1: solve the constraint. *)
  let (On_sol r) = solve range c in
  (* Phase 2: elaborate. *)
  r()

end (* Solve *)

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

(* Re-package the functor [Solve] as a function. *)

let solve ~rectypes c =
  let module S = Solve(struct let rectypes = rectypes end) in
  S.main c

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

(* Combinators for building constraints. *)

(* The type ['a co] forms an applicative functor. *)

let pure x =
  CPure x

let (let+) c f = CMap(c, f)
let (and+) c1 c2 = CConj(c1, c2)

(* [map] and [conjunction] could be defined as follows, if desired: *)

let _map f c =
  let+ x = c in
  f x

let _conjunction c1 c2 =
  let+ x1 = c1
  and+ x2 = c2
  in (x1, x2)

(* The type ['a co] does not form a monad. Indeed, there is no way of defining
   a [bind] combinator. *)

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

(* Existential quantification. *)

type ('var, 'a) binder =
  ('var -> 'a co) -> 'a co

let (let@) m f =
  m f

let decode v =
  CDecode v

let exist_aux t f =
  (* Create a fresh constraint variable [v]. *)
  let v = fresh () in
  (* Pass [v] to the client, *)
  let c = f v in
  (* and wrap the resulting constraint [c] in an existential quantifier. *)
  CExist (v, t, c)

let exist f =
  exist_aux None f

let shallow t f =
  exist_aux (Some t) f

let lift f v1 t2 =
  shallow t2 (fun v2 ->
    f v1 v2
  )

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

(* Deep types. *)

type deep_ty =
  | DeepVar of variable
  | DeepStructure of deep_ty S.structure

(* Conversion of deep types to shallow types. *)

(* Our API is so constrained that this seems extremely difficult to implement
   from the outside. So, we provide it, for the user's convenience. In fact,
   even here, inside the abstraction, implementing this conversion is slightly
   tricky. *)

let deep dty f =
  (* Accumulate a list of the fresh variables that we create. *)
  let vs = ref [] in
  (* [convert] converts a deep type to a variable. *)
  let rec convert dty =
    match dty with
    | DeepVar v ->
        v
    | DeepStructure s ->
        (* First recursively convert our children, then allocate a fresh
           variable [v] to stand for the root. Record its existence in the
           list [vs]. *)
        let v = fresh () in
        let s =
          (* This definition modifies [vs], so it should not be inlined away. *)
          S.map convert s in
        vs := (v, s) :: !vs;
        v
  in
  (* Convert the deep type [dty] and pass the variable that stands for its
     root to the user function [f]. *)
  let c = f (convert dty) in
  (* Then, create a bunch of existential quantifiers, in an arbitrary order. *)
  List.fold_left (fun rc (v, s) -> CExist (v, Some s, rc)) c !vs

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

(* Equations. *)

let (--) v1 v2 =
  CEq (v1, v2)

let (---) v t =
  lift (--) v t

(* If [shallow] was not exposed, [lift] could also be defined (outside this
   module) in terms of [exist] and [---], as follows. This definition seems
   slower, though; its impact on the test suite is quite large. *)

let _other_lift f v1 t2 =
  exist (fun v2 ->
    let+ () = v2 --- t2
    and+ v = f v1 v2
    in v
  )

(* [shallow] could in principe be defined as follows: *)

let _shallow (t : variable O.structure) : (variable, 'a) binder =
  fun f ->
    let@ v = exist in
    let+ () = v --- t
    and+ y = f v in
    y

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

(* Instantiation constraints. *)

let instance x v =
  CInstance (x, v)

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

(* Constraint abstractions. *)

(* The [CDef] form is so trivial that it deserves its own syntax. Viewing it
   as a special case of [CLet] would be more costly (by a constant factor). *)

let def x v c =
  CDef (x, v, c)

(* The auxiliary function [single] asserts that its argument [xs] is a
   singleton list, and extracts its unique element. *)

let single xs =
  match xs with
  | [ x ] ->
      x
  | _ ->
      assert false

(* [letrn] is our most general combinator. It offers full access to the
   expressiveness of [CLet] constraints. *)

(* The integer parameter [k] indicates how many rigid variables the user
   wishes to create. These variables are rigid while the left-hand constraint
   is solved. Once generalization has taken place, they become generic
   variables in the newly-created schemes. *)

(* The general form of [CLet] involves two constraints, the left-hand side and
   the right-hand side, yet it defines a *family* of constraint abstractions,
   which become bound to the term variables [xs]. *)

let letrn k xs f1 c2 =
  (* Allocate a list [rs] of [k] fresh type variables. *)
  let rs = List.init k (fun _ -> fresh()) in
  (* Allocate a fresh type variable for each term variable in [xs]. *)
  let vs = List.map (fun _ -> fresh()) xs in
  (* Apply [f1] to the lists [rs] and [vs], to obtain a constraint [c1]. *)
  let c1 = f1 rs vs in
  (* Done. *)
  CLet (rs, xs, vs, c1, c2)

(* [letr1] is a special case of [letrn] where only one term variable
   is bound, so the lists [xs], [vs], [ss] are singletons. *)

let letr1 k x f1 c2 =
  let+ gammas, ss, v1, v2 = letrn k [x] (fun rs vs -> f1 rs (single vs)) c2
  in gammas, single ss, v1, v2

(* [letn] is a special case of [letrn] where [k] is zero, so no rigid
   variables are created. *)

let letn xs f1 c2 =
  letrn 0 xs (fun _rs vs -> f1 vs) c2

(* [let1] is a special case of [letn] where only one term variable
   is bound, so the lists [xs], [vs], [ss] are singletons. *)

let let1 x f1 c2 =
  let+ gammas, ss, v1, v2 = letn [x] (fun vs -> f1 (single vs)) c2
  in gammas, single ss, v1, v2

(* [let0] is a special case of [letn], where no term variable is bound, and
   the right-hand side is [CTrue]. The constraint produced by [let0] is [ok]
   by construction; in other words, it is a suitable argument to [solve]. *)

let let0 c1 =
  let+ gammas, _, v1, () = letn [] (fun _ -> c1) CTrue
  in gammas, v1

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

(* Correlation with the source code. *)

let correlate range c =
  CRange (range, c)

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

end (* Make *)