catala/compiler/shared_ast/var.ml
Louis Gesbert 035dff35a3 Preserve and propagate types through closure conversion
some of the types (in particular, in hoisted closures) could not be
reconstructed afterwards. This properly propagates the types, including to
closure deconstruction time, giving additional insurance; and allowing
monomorphisation not to choke on the result.
2024-05-30 17:00:07 +02:00

117 lines
4.0 KiB
OCaml

(* This file is part of the Catala compiler, a specification language for tax
and social benefits computation rules. Copyright (C) 2020-2022 Inria,
contributor: Louis Gesbert <louis.gesbert@inria.fr>
Licensed under the Apache License, Version 2.0 (the "License"); you may not
use this file except in compliance with the License. You may obtain a copy of
the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
License for the specific language governing permissions and limitations under
the License. *)
open Catala_utils
open Definitions
(** {1 Variables and their collections} *)
(** This module provides types and helpers for Bindlib variables on the [gexpr]
type *)
type 'e t = ('a, 't) naked_gexpr Bindlib.var constraint 'e = ('a, 't) gexpr
type 'e vars = ('a, 't) naked_gexpr Bindlib.mvar constraint 'e = ('a, 't) gexpr
let make (name : string) : 'e t = Bindlib.new_var (fun x -> EVar x) name
let compare = Bindlib.compare_vars
let equal = Bindlib.eq_vars
let hash = Bindlib.hash_var
let translate (v : 'e1 t) : 'e2 t =
Bindlib.copy_var v (fun x -> EVar x) (Bindlib.name_of v)
type 'e var = 'e t
(* The purpose of this module is just to lift a type parameter outside of
[Set.S] and [Map.S], so that we can have ['e Var.Set.t] for sets of variables
bound to the ['e = ('a, 't) gexpr] expression type. This is made possible by
the fact that [Bindlib.compare_vars] is polymorphic in that parameter; we
first hide that parameter inside an existential, then re-add a phantom type
outside of the set to ensure consistency. Extracting the elements is then
done with [Bindlib.copy_var] but technically it's not much different from an
[Obj] conversion.
If anyone has a better solution, besides a copy-paste of Set.Make / Map.Make
code... *)
module Generic = struct
(* Existentially quantify the type parameters to allow application of
Set.Make *)
type t = Var : 'e var -> t
(* Note: adding [[@@ocaml.unboxed]] would be OK and make our wrappers live at
the type-level without affecting the actual data representation. But
[Bindlib.var] being abstract, we can't convince OCaml it's ok at the moment
and have to hold it *)
let t v = Var v
let get (Var v) = Bindlib.copy_var v (fun x -> EVar x) (Bindlib.name_of v)
let compare (Var x) (Var y) = Bindlib.compare_vars x y
let eq (Var x) (Var y) = Bindlib.eq_vars x y [@@ocaml.warning "-32"]
let format ppf v = String.format ppf (Bindlib.name_of (get v))
end
(* Wrapper around Set.Make to re-add type parameters (avoid inconsistent
sets) *)
module Set = struct
open Generic
open Set.Make (Generic)
type nonrec 'e t = t
let empty = empty
let singleton x = singleton (t x)
let add x s = add (t x) s
let remove x s = remove (t x) s
let union s1 s2 = union s1 s2
let mem x s = mem (t x) s
let of_list l = of_list (List.map t l)
let elements s = elements s |> List.map get
let diff s1 s2 = diff s1 s2
let iter f s = iter (fun x -> f (get x)) s
(* Add more as needed *)
end
(* Wrapper around Map.Make to re-add type parameters (avoid inconsistent
maps) *)
module Map = struct
open Generic
module M = Map.Make (Generic)
open M
type k0 = M.key
exception Not_found = M.Not_found
type nonrec ('e, 'x) t = 'x t
let empty = empty
let singleton v x = singleton (t v) x
let add v x m = add (t v) x m
let remove v m = remove (t v) m
let update v f m = update (t v) f m
let find v m = find (t v) m
let find_opt v m = find_opt (t v) m
let bindings m = bindings m |> List.map (fun (v, x) -> get v, x)
let mem x m = mem (t x) m
let union f m1 m2 = union (fun v x1 x2 -> f (get v) x1 x2) m1 m2
let fold f m acc = fold (fun v x acc -> f (get v) x acc) m acc
let keys m = keys m |> List.map get
let values m = values m
let format_keys ?pp_sep m = format_keys ?pp_sep m
(* Add more as needed *)
end