typer.ml

(* TODO:
 - rewrite type-checking of operators to propagate constraint
 - optimize computation of pattern free variables
 - check whether it is worth using recursive hash-consing internally
*)


let warning loc msg =
  Format.fprintf !Location.warning_ppf "Warning %a:@\n%a%s@\n" 
    Location.print_loc loc
    Location.html_hilight loc
    msg

(* I. Transform the abstract syntax of types and patterns into
      the internal form *)

open Location
open Ast
open Ident

module S = struct type t = string let compare = compare end
module StringSet = Set.Make(S)
module TypeEnv = Map.Make(S)
module Env = Map.Make(Id)


exception NonExhaustive of Types.descr
exception Constraint of Types.descr * Types.descr * string
exception ShouldHave of Types.descr * string
exception ShouldHave2 of Types.descr * string * Types.descr
exception WrongLabel of Types.descr * label
exception UnboundId of string

let raise_loc loc exn = raise (Location (loc,exn))


(* Eliminate Recursion, propagate Sequence Capture Variables *)

let rec seq_vars accu = function
  | Epsilon | Elem _ -> accu
  | Seq (r1,r2) | Alt (r1,r2) -> seq_vars (seq_vars accu r1) r2
  | Star r | WeakStar r -> seq_vars accu r
  | SeqCapture (v,r) -> seq_vars (IdSet.add v accu) r

type derecurs_slot = {
  ploc : Location.loc;
  pid  : int;
  mutable ploop : bool;
  mutable pdescr : derecurs option
} and derecurs =
  | PAlias of derecurs_slot
  | PType of Types.descr
  | POr of derecurs * derecurs
  | PAnd of derecurs * derecurs
  | PDiff of derecurs * derecurs
  | PTimes of derecurs * derecurs
  | PXml of derecurs * derecurs
  | PArrow of derecurs * derecurs
  | POptional of derecurs
  | PRecord of bool * derecurs label_map
  | PCapture of id
  | PConstant of id * Types.const
  | PRegexp of derecurs_regexp * derecurs
and derecurs_regexp =
  | PEpsilon
  | PElem of derecurs
  | PSeq of derecurs_regexp * derecurs_regexp
  | PAlt of derecurs_regexp * derecurs_regexp
  | PStar of derecurs_regexp
  | PWeakStar of derecurs_regexp

let rec hash_derecurs = function
  | PAlias s -> s.pid
  | PType t -> 1 + 17 * (Types.hash_descr t)
  | POr (p1,p2) -> 2 + 17 * (hash_derecurs p1) + 257 * (hash_derecurs p2)
  | PAnd (p1,p2) -> 3 + 17 * (hash_derecurs p1) + 257 * (hash_derecurs p2)
  | PDiff (p1,p2) -> 4 + 17 * (hash_derecurs p1) + 257 * (hash_derecurs p2)
  | PTimes (p1,p2) -> 5 + 17 * (hash_derecurs p1) + 257 * (hash_derecurs p2)
  | PXml (p1,p2) -> 6 + 17 * (hash_derecurs p1) + 257 * (hash_derecurs p2)
  | PArrow (p1,p2) -> 7 + 17 * (hash_derecurs p1) + 257 * (hash_derecurs p2)
  | POptional p -> 8 + 17 * (hash_derecurs p)
  | PRecord (o,r) -> (if o then 9 else 10) + 17 * (LabelMap.hash hash_derecurs r)
  | PCapture x -> 11 + 17 * (Id.hash x)
  | PConstant (x,c) -> 12 + 17 * (Id.hash x) + 257 * (Types.hash_const c)
  | PRegexp (p,q) -> 13 + 17 * (hash_derecurs_regexp p) + 257 * (hash_derecurs q)
and hash_derecurs_regexp = function
  | PEpsilon -> 1
  | PElem p -> 2 + 17 * (hash_derecurs p)
  | PSeq (p1,p2) -> 3 + 17 * (hash_derecurs_regexp p1) + 257 * (hash_derecurs_regexp p2)
  | PAlt (p1,p2) -> 4 + 17 * (hash_derecurs_regexp p1) + 257 * (hash_derecurs_regexp p2)
  | PStar p -> 5 + 17 * (hash_derecurs_regexp p)
  | PWeakStar p -> 6 + 17 * (hash_derecurs_regexp p)

let rec equal_derecurs p1 p2 = (p1 == p2) || match p1,p2 with
  | PAlias s1, PAlias s2 -> s1 == s2
  | PType t1, PType t2 -> Types.equal_descr t1 t2
  | POr (p1,q1), POr (p2,q2)
  | PAnd (p1,q1), PAnd (p2,q2)
  | PDiff (p1,q1), PDiff (p2,q2)
  | PTimes (p1,q1), PTimes (p2,q2)
  | PXml (p1,q1), PXml (p2,q2)
  | PArrow (p1,q1), PArrow (p2,q2) -> (equal_derecurs p1 p2) && (equal_derecurs q1 q2)
  | POptional p1, POptional p2 -> equal_derecurs p1 p2
  | PRecord (o1,r1), PRecord (o2,r2) -> (o1 == o2) && (LabelMap.equal equal_derecurs r1 r2)
  | PCapture x1, PCapture x2 -> Id.equal x1 x2
  | PConstant (x1,c1), PConstant (x2,c2) -> (Id.equal x1 x2) && (Types.equal_const c1 c2)
  | PRegexp (p1,q1), PRegexp (p2,q2) -> (equal_derecurs_regexp p1 p2) && (equal_derecurs q1 q2)
  | _ -> false
and equal_derecurs_regexp r1 r2 = match r1,r2 with
  | PEpsilon, PEpsilon -> true
  | PElem p1, PElem p2 -> equal_derecurs p1 p2
  | PSeq (p1,q1), PSeq (p2,q2) 
  | PAlt (p1,q1), PAlt (p2,q2) -> (equal_derecurs_regexp p1 p2) && (equal_derecurs_regexp q1 q2)
  | PStar p1, PStar p2
  | PWeakStar p1, PWeakStar p2 -> equal_derecurs_regexp p1 p2
  | _ -> false

module DerecursTable = Hashtbl.Make(
  struct 
    type t = derecurs 
    let hash = hash_derecurs
    let equal = equal_derecurs
  end
)

module RE = Hashtbl.Make(
  struct 
    type t = derecurs_regexp * derecurs 
    let hash (p,q) = (hash_derecurs_regexp p) + 17 * (hash_derecurs q)
    let equal (p1,q1) (p2,q2) = (equal_derecurs_regexp p1 p2) && (equal_derecurs q1 q2)
  end
)

  
let counter = State.ref "Typer.counter - derecurs" 0
let mk_slot loc = 
  incr counter; 
  { ploop = false; ploc = loc; pid = !counter; pdescr = None }
  
let rec derecurs env p = match p.descr with
  | PatVar v ->
      (try PAlias (TypeEnv.find v env)
       with Not_found -> raise_loc_generic p.loc ("Undefined type/pattern " ^ v))
  | Recurs (p,b) -> derecurs (derecurs_def env b) p
  | Internal t -> PType t
  | Or (p1,p2) -> POr (derecurs env p1, derecurs env p2)
  | And (p1,p2) -> PAnd (derecurs env p1, derecurs env p2)
  | Diff (p1,p2) -> PDiff (derecurs env p1, derecurs env p2)
  | Prod (p1,p2) -> PTimes (derecurs env p1, derecurs env p2)
  | XmlT (p1,p2) -> PXml (derecurs env p1, derecurs env p2)
  | Arrow (p1,p2) -> PArrow (derecurs env p1, derecurs env p2)
  | Optional p -> POptional (derecurs env p)
  | Record (o,r) -> PRecord (o, LabelMap.map (derecurs env) r)
  | Capture x -> PCapture x
  | Constant (x,c) -> PConstant (x,c)
  | Regexp (r,q) -> 
      let constant_nil t v = PAnd (t, PConstant (v, Types.Atom Sequence.nil_atom)) in
      let vars = seq_vars IdSet.empty r in
      let q = IdSet.fold constant_nil (derecurs env q) vars in
      let r = derecurs_regexp (fun p -> p) env r in
      PRegexp (r, q)
and derecurs_regexp vars env = function
  | Epsilon -> PEpsilon
  | Elem p -> PElem (vars (derecurs env p))
  | Seq (p1,p2) -> PSeq (derecurs_regexp vars env p1, derecurs_regexp vars env p2)
  | Alt (p1,p2) -> PAlt (derecurs_regexp vars env p1, derecurs_regexp vars env p2)
  | Star p -> PStar (derecurs_regexp vars env p)
  | WeakStar p -> PWeakStar (derecurs_regexp vars env p)
  | SeqCapture (x,p) -> derecurs_regexp (fun p -> PAnd (vars p, PCapture x)) env p


and derecurs_def env b =
  let b = List.map (fun (v,p) -> (v,p,mk_slot p.loc)) b in
  let env = List.fold_left (fun env (v,p,s) -> TypeEnv.add v s env) env b in
  List.iter (fun (v,p,s) -> s.pdescr <- Some (derecurs env p)) b;
  env

(* Stratification and recursive hash-consing *)

type descr = 
  | IType of Types.descr
  | IOr of descr * descr
  | IAnd of descr * descr
  | IDiff of descr * descr
  | ITimes of slot * slot
  | IXml of slot * slot
  | IArrow of slot * slot
  | IOptional of descr
  | IRecord of bool * slot label_map
  | ICapture of id
  | IConstant of id * Types.const
and slot = {
  mutable fv : fv option;
  mutable hash : int option;
  mutable rank1: int; mutable rank2: int;
  mutable gen1 : int; mutable gen2: int;
  mutable d    : descr option
}
    
let descr s = 
  match s.d with
    | Some d -> d
    | None -> assert false
	
let gen = ref 0
let rank = ref 0
	     
let rec hash_descr = function
  | IType x -> Types.hash_descr x
  | IOr (d1,d2) -> 1 + 17 * (hash_descr d1) + 257 * (hash_descr d2)
  | IAnd (d1,d2) -> 2 + 17 * (hash_descr d1) + 257 * (hash_descr d2)
  | IDiff (d1,d2) -> 3 + 17 * (hash_descr d1) + 257 * (hash_descr d2)
  | IOptional d -> 4 + 17 * (hash_descr d)
  | ITimes (s1,s2) -> 5 + 17 * (hash_slot s1) + 257 * (hash_slot s2)
  | IXml (s1,s2) -> 6 + 17 * (hash_slot s1) + 257 * (hash_slot s2)
  | IArrow (s1,s2) -> 7 + 17 * (hash_slot s1) + 257 * (hash_slot s2)
  | IRecord (o,r) -> (if o then 8 else 9) + 17 * (LabelMap.hash hash_slot r)
  | ICapture x -> 10 + 17 * (Id.hash x)
  | IConstant (x,y) -> 11 + 17 * (Id.hash x) + 257 * (Types.hash_const y)
and hash_slot s =
  if s.gen1 = !gen then 13 * s.rank1
  else (
    incr rank;
    s.rank1 <- !rank; s.gen1 <- !gen;
    hash_descr (descr s)
  )
    
let rec equal_descr d1 d2 = 
  match (d1,d2) with
  | IType x1, IType x2 -> Types.equal_descr x1 x2
  | IOr (x1,y1), IOr (x2,y2) 
  | IAnd (x1,y1), IAnd (x2,y2) 
  | IDiff (x1,y1), IDiff (x2,y2) -> (equal_descr x1 x2) && (equal_descr y1 y2)
  | IOptional x1, IOptional x2 -> equal_descr x1 x2
  | ITimes (x1,y1), ITimes (x2,y2) 
  | IXml (x1,y1), IXml (x2,y2) 
  | IArrow (x1,y1), IArrow (x2,y2) -> (equal_slot x1 x2) && (equal_slot y1 y2)
  | IRecord (o1,r1), IRecord (o2,r2) -> (o1 = o2) && (LabelMap.equal equal_slot r1 r2)
  | ICapture x1, ICapture x2 -> Id.equal x1 x2
  | IConstant (x1,y1), IConstant (x2,y2) -> (Id.equal x1 x2) && (Types.equal_const y1 y2)
  | _ -> false
and equal_slot s1 s2 =
  ((s1.gen1 = !gen) && (s2.gen2 = !gen) && (s1.rank1 = s2.rank2))
  ||
  ((s1.gen1 <> !gen) && (s2.gen2 <> !gen) && (
     incr rank;
     s1.rank1 <- !rank; s1.gen1 <- !gen;
     s2.rank2 <- !rank; s2.gen2 <- !gen;
     equal_descr (descr s1) (descr s2)
   ))
  
module Arg = struct
  type t = slot
      
  let hash s =
    match s.hash with
      | Some h -> h
      | None ->
	  incr gen; rank := 0; 
	  let h = hash_slot s in
	  s.hash <- Some h;
	  h
	    
  let equal s1 s2 = 
    (s1 == s2) || 
    (incr gen; rank := 0; 
     let e = equal_slot s1 s2 in
(*     if e then Printf.eprintf "Recursive hash-consig: Equal\n";  *)
     e)
end
module SlotTable = Hashtbl.Make(Arg)
  
let rec fv_slot s =
  match s.fv with
    | Some x -> x
    | None ->
	if s.gen1 = !gen then IdSet.empty 
	else (s.gen1 <- !gen; fv_descr (descr s))
and fv_descr = function
  | IType _ -> IdSet.empty
  | IOr (d1,d2)
  | IAnd (d1,d2)  
  | IDiff (d1,d2) -> IdSet.cup (fv_descr d1) (fv_descr d2)
  | IOptional d -> fv_descr d
  | ITimes (s1,s2)  
  | IXml (s1,s2)  
  | IArrow (s1,s2) -> IdSet.cup (fv_slot s1) (fv_slot s2)
  | IRecord (o,r) -> List.fold_left IdSet.cup IdSet.empty (LabelMap.map_to_list fv_slot r)
  | ICapture x | IConstant (x,_) -> IdSet.singleton x

      
let compute_fv s =
  match s.fv with
    | Some x -> ()
    | None ->
	incr gen;
	let x = fv_slot s in
	s.fv <- Some x
	  

let todo_fv = ref []
	  
let mk () =   
  let s = 
    { d = None;
      fv = None;
      hash = None;
      rank1 = 0; rank2 = 0;
      gen1 = 0; gen2 = 0 } in
  todo_fv := s :: !todo_fv;
  s

let flush_fv () =
  List.iter compute_fv !todo_fv;
  todo_fv := []
    
let compile_slot_hash = DerecursTable.create 67
let compile_hash = DerecursTable.create 67

let defs = ref []

let rec compile p =
  try DerecursTable.find compile_hash p
  with Not_found ->
    let c = real_compile p in
    DerecursTable.replace compile_hash p c;
    c
and real_compile = function
  | PAlias v ->
      if v.ploop then
	raise_loc_generic v.ploc ("Unguarded recursion on type/pattern");
      v.ploop <- true;
      let r = match v.pdescr with Some x -> compile x | _ -> assert false in
      v.ploop <- false;
      r
  | PType t -> IType t
  | POr (t1,t2) -> IOr (compile t1, compile t2)
  | PAnd (t1,t2) -> IAnd (compile t1, compile t2)
  | PDiff (t1,t2) -> IDiff (compile t1, compile t2)
  | PTimes (t1,t2) -> ITimes (compile_slot t1, compile_slot t2)
  | PXml (t1,t2) -> IXml (compile_slot t1, compile_slot t2)
  | PArrow (t1,t2) -> IArrow (compile_slot t1, compile_slot t2)
  | POptional t -> IOptional (compile t)
  | PRecord (o,r) ->  IRecord (o, LabelMap.map compile_slot r)
  | PConstant (x,v) -> IConstant (x,v)
  | PCapture x -> ICapture x
  | PRegexp (r,q) -> compile_regexp r q
and compile_regexp r q =
  let memo = RE.create 17 in
  let rec aux accu r q =
    if RE.mem memo (r,q) then accu
    else (
      RE.add memo (r,q) ();
      match r with
	| PEpsilon -> (match q with PRegexp (r,q) -> aux accu r q | _ -> (compile q) :: accu)
	| PElem p -> ITimes (compile_slot p, compile_slot q) :: accu
	| PSeq (r1,r2) -> aux accu r1 (PRegexp (r2,q))
	| PAlt (r1,r2) -> aux (aux accu r1 q) r2 q
	| PStar r1 -> aux (aux accu r1 (PRegexp (r,q))) PEpsilon q
	| PWeakStar r1 -> aux (aux accu PEpsilon q) r1 (PRegexp (r,q))
    )
  in
  let accu = aux [] r q in
  match accu with
    | [] -> assert false
    | p::l -> List.fold_left (fun acc p -> IOr (p,acc)) p l
and compile_slot p =
  try DerecursTable.find compile_slot_hash p
  with Not_found ->
    let s = mk () in
    defs := (s,p) :: !defs;
    DerecursTable.add compile_slot_hash p s;
    s

      
let rec flush_defs () = 
  match !defs with
    | [] -> ()
    | (s,p)::t -> defs := t; s.d <- Some (compile p); flush_defs ()
	
let typ_nodes = SlotTable.create 67
let pat_nodes = SlotTable.create 67
		  
let rec typ = function
  | IType t -> t
  | IOr (s1,s2) -> Types.cup (typ s1) (typ s2)
  | IAnd (s1,s2) ->  Types.cap (typ s1) (typ s2)
  | IDiff (s1,s2) -> Types.diff (typ s1) (typ s2)
  | ITimes (s1,s2) -> Types.times (typ_node s1) (typ_node s2)
  | IXml (s1,s2) -> Types.xml (typ_node s1) (typ_node s2)
  | IArrow (s1,s2) -> Types.arrow (typ_node s1) (typ_node s2)
  | IOptional s -> Types.Record.or_absent (typ s)
  | IRecord (o,r) -> Types.record' (o, LabelMap.map typ_node r)
  | ICapture x | IConstant (x,_) -> assert false
      
and typ_node s : Types.node =
  try SlotTable.find typ_nodes s
  with Not_found ->
    let x = Types.make () in
    SlotTable.add typ_nodes s x;
    Types.define x (typ (descr s));
    x
      
let rec pat d : Patterns.descr =
  if IdSet.is_empty (fv_descr d)
  then Patterns.constr (typ d)
  else pat_aux d
    
    
and pat_aux = function
  | IOr (s1,s2) -> Patterns.cup (pat s1) (pat s2)
  | IAnd (s1,s2) -> Patterns.cap (pat s1) (pat s2)
  | IDiff (s1,s2) when IdSet.is_empty (fv_descr s2) ->
      let s2 = Types.neg (typ s2) in
      Patterns.cap (pat s1) (Patterns.constr s2)
  | IDiff _ ->
      raise (Patterns.Error "Difference not allowed in patterns")
  | ITimes (s1,s2) -> Patterns.times (pat_node s1) (pat_node s2)
  | IXml (s1,s2) -> Patterns.xml (pat_node s1) (pat_node s2)
  | IOptional _ -> 
      raise (Patterns.Error "Optional field not allowed in record patterns")
  | IRecord (o,r) ->
      let pats = ref [] in
      let aux l s = 
	if IdSet.is_empty (fv_slot s) then typ_node s
	else
	  ( pats := Patterns.record l (pat_node s) :: !pats;
	    Types.any_node )
      in
      let constr = Types.record' (o,LabelMap.mapi aux r) in
      List.fold_left Patterns.cap (Patterns.constr constr) !pats
	(* TODO: can avoid constr when o=true, and all fields have fv *)
  | ICapture x -> Patterns.capture x
  | IConstant (x,c) -> Patterns.constant x c
  | IArrow _ ->
      raise (Patterns.Error "Arrow not allowed in patterns")
  | IType _ -> assert false
      
and pat_node s : Patterns.node =
  try SlotTable.find pat_nodes s
  with Not_found ->
    let x = Patterns.make (fv_slot s) in
    SlotTable.add pat_nodes s x;
    Patterns.define x (pat (descr s));
    x
      
let glb = State.ref "Typer.glb_env" TypeEnv.empty
	    
	    
let register_global_types b =
  List.iter 
    (fun (v,p) ->
       if TypeEnv.mem v !glb
       then raise_loc_generic p.loc ("Multiple definition for type " ^ v)
    ) b;
  glb := derecurs_def !glb b;
  let b = List.map (fun (v,p) -> (v,p,compile (derecurs !glb p))) b in
  flush_defs ();
  flush_fv ();
  List.iter 
    (fun (v,p,s) -> 
       if not (IdSet.is_empty (fv_descr s)) then
	 raise_loc_generic p.loc "Capture variables are not allowed in types";
       let t = typ s in
       if (p.loc <> noloc) && (Types.is_empty t) then
	 warning p.loc ("This definition yields an empty type for " ^ v);
       Types.Print.register_global v t) b

let dump_global_types ppf =
  TypeEnv.iter (fun v _ -> Format.fprintf ppf " %s" v) !glb
    
    
let typ p =
  let s = compile_slot (derecurs !glb p) in
  flush_defs ();
  flush_fv ();
  if IdSet.is_empty (fv_slot s) then typ_node s
  else raise_loc_generic p.loc "Capture variables are not allowed in types"
    
let pat p = 
  let s = compile_slot (derecurs !glb p) in
  flush_defs ();
  flush_fv ();
  try pat_node s
  with Patterns.Error e -> raise_loc_generic p.loc e
    | Location (loc,exn) when loc = noloc -> raise (Location (p.loc, exn))


(* II. Build skeleton *)

module Fv = IdSet

let all_branches = ref []

(* IDEA: introduce a node Loc in the AST to override nolocs
   in sub-expressions *)
   
let exp loc fv e =
  fv,
  { Typed.exp_loc = loc;
    Typed.exp_typ = Types.empty;
    Typed.exp_descr = e;
  }


let rec expr loc = function
  | LocatedExpr (loc,e) -> expr loc e
  | Forget (e,t) ->
      let (fv,e) = expr loc e and t = typ t in
      exp loc fv (Typed.Forget (e,t))
  | Var s -> 
      exp loc (Fv.singleton s) (Typed.Var s)
  | Apply (e1,e2) -> 
      let (fv1,e1) = expr loc e1 and (fv2,e2) = expr loc e2 in
      exp loc (Fv.cup fv1 fv2) (Typed.Apply (e1,e2))
  | Abstraction a ->
      let iface = List.map (fun (t1,t2) -> (typ t1, typ t2)) 
		    a.fun_iface in
      let t = List.fold_left 
		(fun accu (t1,t2) -> Types.cap accu (Types.arrow t1 t2)) 
		Types.any iface in
      let iface = List.map 
		    (fun (t1,t2) -> (Types.descr t1, Types.descr t2)) 
		    iface in
      let (fv0,body) = branches a.fun_body in
      let fv = match a.fun_name with
	| None -> fv0
	| Some f -> Fv.remove f fv0 in
      let e = Typed.Abstraction 
		{ Typed.fun_name = a.fun_name;
		  Typed.fun_iface = iface;
		  Typed.fun_body = body;
		  Typed.fun_typ = t;
		  Typed.fun_fv = fv
		} in
      exp loc fv e
  | Cst c -> 
      exp loc Fv.empty (Typed.Cst c)
  | Pair (e1,e2) ->
      let (fv1,e1) = expr loc e1 and (fv2,e2) = expr loc e2 in
      exp loc (Fv.cup fv1 fv2) (Typed.Pair (e1,e2))
  | Xml (e1,e2) ->
      let (fv1,e1) = expr loc e1 and (fv2,e2) = expr loc e2 in
      exp loc (Fv.cup fv1 fv2) (Typed.Xml (e1,e2))
  | Dot (e,l) ->
      let (fv,e) = expr loc e in
      exp loc fv (Typed.Dot (e,l))
  | RemoveField (e,l) ->
      let (fv,e) = expr loc e in
      exp loc fv (Typed.RemoveField (e,l))
  | RecordLitt r -> 
      let fv = ref Fv.empty in
      let r = LabelMap.map 
		(fun e -> 
		   let (fv2,e) = expr loc e 
		   in fv := Fv.cup !fv fv2; e)
		r in
      exp loc !fv (Typed.RecordLitt r)
  | Op (op,le) ->
      let (fvs,ltes) = List.split (List.map (expr loc) le) in
      let fv = List.fold_left Fv.cup Fv.empty fvs in
      exp loc fv (Typed.Op (op,ltes))
  | Match (e,b) -> 
      let (fv1,e) = expr loc e
      and (fv2,b) = branches b in
      exp loc (Fv.cup fv1 fv2) (Typed.Match (e, b))
  | Map (e,b) ->
      let (fv1,e) = expr loc e
      and (fv2,b) = branches b in
      exp loc (Fv.cup fv1 fv2) (Typed.Map (e, b))
  | Xtrans (e,b) ->
      let b = b @ [ mknoloc (Internal Types.any), MatchFail ] in
      let (fv1,e) = expr loc e
      and (fv2,b) = branches b in
      exp loc (Fv.cup fv1 fv2) (Typed.Xtrans (e, b))
  | MatchFail -> 
      exp loc (Fv.empty) Typed.MatchFail
  | Try (e,b) ->
      let (fv1,e) = expr loc e
      and (fv2,b) = branches b in
      exp loc (Fv.cup fv1 fv2) (Typed.Try (e, b))

	      
  and branches b = 
    let fv = ref Fv.empty in
    let accept = ref Types.empty in
    let branch (p,e) = 
      let (fv2,e) = expr noloc e in
      let br_loc = merge_loc p.loc e.Typed.exp_loc in
      let p = pat p in
      let fv2 = Fv.diff fv2 (Patterns.fv p) in
      fv := Fv.cup !fv fv2;
      accept := Types.cup !accept (Types.descr (Patterns.accept p));
      let br = 
	{ 
	  Typed.br_loc = br_loc;
	  Typed.br_used = br_loc = noloc;
	  Typed.br_pat = p;
	  Typed.br_body = e } in
      all_branches := br :: !all_branches;
      br in
    let b = List.map branch b in
    (!fv, 
     { 
       Typed.br_typ = Types.empty; 
       Typed.br_branches = b; 
       Typed.br_accept = !accept;
       Typed.br_compiled = None;
     } 
    )

let expr = expr noloc

let let_decl p e =
  let (_,e) = expr e in
  { Typed.let_pat = pat p;
    Typed.let_body = e;
    Typed.let_compiled = None }

(* III. Type-checks *)

let int_cup_record = Types.cup Types.Int.any Types.Record.any


type env = Types.descr Env.t

let match_fail = ref Types.empty

open Typed


let check loc t s msg =
  if not (Types.subtype t s) then raise_loc loc (Constraint (t, s, msg))

let rec type_check env e constr precise = 
(*  Format.fprintf Format.std_formatter "constr=%a precise=%b@\n"
    Types.Print.print_descr constr precise; 
*)
  let d = type_check' e.exp_loc env e.exp_descr constr precise in
  e.exp_typ <- Types.cup e.exp_typ d;
  d

and type_check' loc env e constr precise = match e with
  | Forget (e,t) ->
      let t = Types.descr t in
      ignore (type_check env e t false);
      t
  | Abstraction a ->
      let t =
	try Types.Arrow.check_strenghten a.fun_typ constr 
	with Not_found -> 
	  raise_loc loc 
	  (ShouldHave
	     (constr, "but the interface of the abstraction is not compatible"))
      in
      let env = match a.fun_name with
	| None -> env
	| Some f -> Env.add f a.fun_typ env in
      List.iter 
	(fun (t1,t2) ->
	   ignore (type_check_branches loc env t1 a.fun_body t2 false)
	) a.fun_iface;
      t

  | Match (e,b) ->
      let t = type_check env e b.br_accept true in
      type_check_branches loc env t b constr precise

  | Try (e,b) ->
      let te = type_check env e constr precise in
      let tb = type_check_branches loc env Types.any b constr precise in
      Types.cup te tb

  | Pair (e1,e2) ->
      type_check_pair loc env e1 e2 constr precise
  | Xml (e1,e2) ->
      type_check_pair ~kind:`XML loc env e1 e2 constr precise

  | RecordLitt r ->
(* try to get rid of precise = true for values of fields *)
(* also: the use equivalent of need_second to optimize... *)
      let precise = true in
      if not (Types.Record.has_record constr) then
	raise_loc loc (ShouldHave (constr,"but it is a record."));
      let (rconstr,res) = 
	List.fold_left
	  (fun (rconstr,res) (l,e) ->
	     (* could compute (split l e) once... *)
	     let pi = Types.Record.project_opt rconstr l in
	     if Types.is_empty pi then 
	       raise_loc loc 
		 (ShouldHave (constr,(Printf.sprintf 
					"Field %s is not allowed here."
					(LabelPool.value l)
				     )
			     ));
	     let t = type_check env e pi true in
	     let rconstr = Types.Record.condition rconstr l t in
	     let res = if precise then (l,Types.cons t) :: res else res in
	     (rconstr,res)
	  ) (constr, []) (LabelMap.get r)
      in
      if not (Types.Record.has_empty_record rconstr) then
	raise_loc loc 
	  (ShouldHave (constr,"More field should be present"));
      let t = 
	Types.record' (false, LabelMap.from_list (fun _ _ -> assert false) res)
      in
      check loc t constr "";
      t
(*
      if precise then
	Types.record' (false, LabelMap.from_list (fun _ _ -> assert false) res)
      else constr
*)
  | Map (e,b) ->
      let t = type_check env e (Sequence.star b.br_accept) true in

      let constr' = Sequence.approx (Types.cap Sequence.any constr) in
      let exact = Types.subtype (Sequence.star constr') constr in
      (* Note: 
	 - could be more precise by integrating the decomposition
	 of constr inside Sequence.map.
      *)
      let res = 
	Sequence.map 
	  (fun t -> 
	     type_check_branches loc env t b constr' (precise || (not exact)))
	  t in
      if not exact then check loc res constr "";
      if precise then res else constr
  | Op ("@", [e1;e2]) ->
      let constr' = Sequence.star 
		      (Sequence.approx (Types.cap Sequence.any constr)) in
      let exact = Types.subtype constr' constr in
      if exact then
	let t1 = type_check env e1 constr' precise
	and t2 = type_check env e2 constr' precise in
	if precise then Sequence.concat t1 t2 else constr
      else
	(* Note:
	   the knownledge of t1 may makes it useless to
	   check t2 with 'precise' ... *)
	let t1 = type_check env e1 constr' true
	and t2 = type_check env e2 constr' true in
	let res = Sequence.concat t1 t2 in
	check loc res constr "";
	if precise then res else constr
  | Apply (e1,e2) ->
(*
      let constr' = Sequence.star 
		      (Sequence.approx (Types.cap Sequence.any constr)) in
      let t1 = type_check env e1 (Types.cup Types.Arrow.any constr') true in
      let t1_fun = Types.Arrow.get t1 in

      let has_fun = not (Types.Arrow.is_empty t1_fun)
      and has_seq = not (Types.subtype t1 Types.Arrow.any) in

      let constr' =
	Types.cap 
	  (if has_fun then Types.Arrow.domain t1_fun else Types.any)
	  (if has_seq then constr' else Types.any)
      in
      let need_arg = has_fun && Types.Arrow.need_arg t1_fun in
      let precise  = need_arg || has_seq in
      let t2 = type_check env e2 constr' precise in
      let res = Types.cup 
		  (if has_fun then 
		     if need_arg then Types.Arrow.apply t1_fun t2
		     else Types.Arrow.apply_noarg t1_fun
		   else Types.empty)
		  (if has_seq then Sequence.concat t1 t2
		   else Types.empty)
      in
      check loc res constr "";
      res
*)
      let t1 = type_check env e1 Types.Arrow.any true in
      let t1 = Types.Arrow.get t1 in
      let dom = Types.Arrow.domain t1 in
      let res =
	if Types.Arrow.need_arg t1 then
	  let t2 = type_check env e2 dom true in
	  Types.Arrow.apply t1 t2
	else
	  (ignore (type_check env e2 dom false); Types.Arrow.apply_noarg t1)
      in
      check loc res constr "";
      res
  | Op ("flatten", [e]) ->
      let constr' = Sequence.star 
		      (Sequence.approx (Types.cap Sequence.any constr)) in
      let sconstr' = Sequence.star constr' in
      let exact = Types.subtype constr' constr in
      if exact then
	let t = type_check env e sconstr' precise in
	if precise then Sequence.flatten t else constr
      else
	let t = type_check env e sconstr' true in
	let res = Sequence.flatten t in
	check loc res constr "";
	if precise then res else constr
  | Op ("atom_of", [e]) ->
      let t = type_check env e Sequence.string false in
      Types.atom Atoms.any
  | _ -> 
      let t : Types.descr = compute_type' loc env e in
      check loc t constr "";
      t

and type_check_pair ?(kind=`Normal) loc env e1 e2 constr precise =
  let rects = Types.Product.get ~kind constr in
  if Types.Product.is_empty rects then 
    (match kind with
      | `Normal -> raise_loc loc (ShouldHave (constr,"but it is a pair."))
      | `XML -> raise_loc loc (ShouldHave (constr,"but it is an XML element.")));
  let need_s = Types.Product.need_second rects in
  let t1 = type_check env e1 (Types.Product.pi1 rects) (precise || need_s) in
  let c2 = Types.Product.constraint_on_2 rects t1 in
  if Types.is_empty c2 then 
    raise_loc loc (ShouldHave2 (constr,"but the first component has type",t1));
  let t2 = type_check env e2 c2 precise in

  if precise then 
    match kind with
      | `Normal -> Types.times (Types.cons t1) (Types.cons t2)
      | `XML -> Types.xml (Types.cons t1) (Types.cons t2)
  else
    constr


and compute_type env e =
  type_check env e Types.any true

and compute_type' loc env = function
  | Var s -> 
      (try Env.find s env 
       with Not_found -> raise_loc loc (UnboundId (Id.value s))
      )
  | Cst c -> Types.constant c
  | Dot (e,l) ->
      let t = type_check env e Types.Record.any true in
         (try (Types.Record.project t l) 
          with Not_found -> raise_loc loc (WrongLabel(t,l)))
  | RemoveField (e,l) ->
      let t = type_check env e Types.Record.any true in
      Types.Record.remove_field t l
  | Op (op, el) ->
      let args = List.map (fun e -> (e.exp_loc, compute_type env e)) el in
      type_op loc op args
  | Xtrans (e,b) ->
      let t = type_check env e Sequence.any true in
      let r = 
	Sequence.map_tree 
	  (fun t -> 
	     let res = type_check_branches loc env t b Sequence.any true in
	     let resid = !match_fail in
	     match_fail := Types.empty;
	     (res,resid)
	  ) t
      in
      r

(* We keep these cases here to allow comparison and benchmarking ...
   Just comment the corresponding cases in type_check' to
   activate these ones.
*)
  | Map (e,b) ->
      let t = compute_type env e in
      Sequence.map (fun t -> type_check_branches loc env t b Types.any true) t
  | Pair (e1,e2) -> 
      let t1 = compute_type env e1 
      and t2 = compute_type env e2 in
      Types.times (Types.cons t1) (Types.cons t2)
  | RecordLitt r ->
      let r = LabelMap.map (fun e -> Types.cons (compute_type env e)) r in
      Types.record' (false,r)
  | _ -> assert false

and type_check_branches loc env targ brs constr precise =
  if Types.is_empty targ then Types.empty 
  else (
    brs.br_typ <- Types.cup brs.br_typ targ;
    branches_aux loc env targ 
      (if precise then Types.empty else constr) 
      constr precise brs.br_branches
  )
    
and branches_aux loc env targ tres constr precise = function
  | [] -> raise_loc loc (NonExhaustive targ)
  | { br_body = { exp_descr = MatchFail } } :: _ ->
      match_fail := Types.cup !match_fail targ;
      tres
  | b :: rem ->
      let p = b.br_pat in
      let acc = Types.descr (Patterns.accept p) in

      let targ' = Types.cap targ acc in
      if Types.is_empty targ' 
      then branches_aux loc env targ tres constr precise rem
      else 
	( b.br_used <- true;
	  let res = Patterns.filter targ' p in
	  let env' = List.fold_left 
		       (fun env (x,t) -> Env.add x (Types.descr t) env) 
		       env res in
	  let t = type_check env' b.br_body constr precise in
	  let tres = if precise then Types.cup t tres else tres in
	  let targ'' = Types.diff targ acc in
	  if (Types.non_empty targ'') then 
	    branches_aux loc env targ'' tres constr precise rem 
	  else
	    tres
	)

and type_let_decl env l =
  let acc = Types.descr (Patterns.accept l.let_pat) in
  let t = type_check env l.let_body acc true in
  let res = Patterns.filter t l.let_pat in
  List.map (fun (x,t) -> (x, Types.descr t)) res

and type_rec_funs env l =
  let types = 
    List.fold_left
      (fun accu -> function  {let_body={exp_descr=Abstraction a}} as l ->
	 let t = a.fun_typ in
	 let acc = Types.descr (Patterns.accept l.let_pat) in
	 if not (Types.subtype t acc) then
	   raise_loc l.let_body.exp_loc (NonExhaustive (Types.diff t acc));
	 let res = Patterns.filter t l.let_pat in
	 List.fold_left (fun accu (x,t) -> (x, Types.descr t)::accu) accu res
	 | _ -> assert false) [] l
  in
  let env' = List.fold_left (fun env (x,t) -> Env.add x t env) env types in
  List.iter 
    (function  { let_body = { exp_descr = Abstraction a } } as l ->
       ignore (type_check env' l.let_body Types.any false)
       | _ -> assert false) l;
  types


and type_op loc op args =
  match (op,args) with
    | "+", [loc1,t1; loc2,t2] ->
	check loc1 t1 int_cup_record
	"The first argument of + must be an integer or a record";
	let int = Types.Int.get t1 in
	let int = if Intervals.is_empty int then None else Some int in
	let r = if Types.Record.has_record t1 then Some t1 else None in
	(match (int,r) with
	   | Some t1, None ->
	       if not (Types.Int.is_int t2) then
		 raise_loc loc2
		   (Constraint 
		      (t2,Types.Int.any,
		       "The second argument of + must be an integer"));
	       Types.Int.put
		 (Intervals.add t1 (Types.Int.get t2));
	   | None, Some r1 ->
	       check loc2 t2 Types.Record.any 
	       "The second argument of + must be a record";
	       Types.Record.merge r1 t2
	   | None, None ->
	       Types.empty
	   | Some t1, Some r1 ->
	       check loc2 t2 int_cup_record
	       "The second argument of + must be an integer or a record";
	       Types.cup 
		 (Types.Int.put (Intervals.add t1 (Types.Int.get t2)))
		 (Types.Record.merge r1 t2)
	)
    | "-", [loc1,t1; loc2,t2] ->
	type_int_binop Intervals.sub loc1 t1 loc2 t2
    | ("*" | "/" | "mod"), [loc1,t1; loc2,t2] ->
	type_int_binop (fun i1 i2 -> Intervals.any) loc1 t1 loc2 t2
    | "@", [loc1,t1; loc2,t2] ->
	check loc1 t1 Sequence.any
	  "The first argument of @ must be a sequence";
	Sequence.concat t1 t2
    | "flatten", [loc1,t1] ->
	check loc1 t1 Sequence.seqseq 
	  "The argument of flatten must be a sequence of sequences";
	Sequence.flatten t1
    | "load_xml", [loc1,t1] ->
	check loc1 t1 Sequence.string
	  "The argument of load_xml must be a string (filename)";
	Types.any
    | "load_file", [loc1,t1] ->
	check loc1 t1 Sequence.string
	  "The argument of load_file must be a string (filename)";
	Sequence.string
    | "load_html", [loc1,t1] ->
	check loc1 t1 Sequence.string
	  "The argument of load_html must be a string (filename)";
	Types.any
    | "raise", [loc1,t1] ->
	Types.empty
    | "print_xml", [loc1,t1] ->
	Sequence.string
    | "print", [loc1,t1] ->