ocaml: reform msr stuff

ocaml/libsyntax/expr.ml

@@ -8,10 +8,9 @@ type sel_ident = loc_span * stack_ref
 type max_stack_ref = stack_ref
 type 'a record = 'a Record.t
 
-type pattern_mt = Pattern.t * ((loc_span * expr) option)
-
-and lambda =
-  { pat       : pattern_mt
+type lambda =
+  { pat       : Pattern.t
+  ; pty       : (loc_span * expr) option
   ; body_span : loc_span
   ; body      : expr
   }
@@ -23,27 +22,36 @@ and expr =
   | RefOf of sel_ident
   | Lambda of max_stack_ref * lambda
   | TyLambda of max_stack_ref * expr * lambda (* note that the first element is the layout of the binders *)
-  | Apply of max_stack_ref * expr * loc_span * expr
-  | RefTy of max_stack_ref * stack_ref * expr
-  | Record of max_stack_ref * expr record
-  | TyRecord of max_stack_ref * expr record
+  | Apply of expr * loc_span * expr
+  | RefTy of stack_ref * expr
+  | Record of expr record
+  | TyRecord of expr record
 
 (* parser *)
 
 module FieldsReg = Set.Make(String)
 module FieldsReg2 = Map.Make(String)
 
-let msr_of_expr = function
+let rec msr_of_expr = function
   | Infer -> 0
   | ELiteral _ -> 0
   | Use (_, x) -> x + 1
   | RefOf (_, x) -> x + 1
   | Lambda (x, _) -> x
   | TyLambda (x, _, _) -> x
-  | Apply (x, _, _, _) -> x
-  | RefTy (x, _, _) -> x
-  | Record (x, _) -> x
-  | TyRecord (x, _) -> x
+  | Apply (x, _, y) -> Int.max (msr_of_expr x) (msr_of_expr y)
+  | RefTy (x, y) -> Int.max (x + 1) (msr_of_expr y)
+  | Record x -> msr_of_record x 0
+  | TyRecord x -> msr_of_record x 0
+
+and msr_of_record rcd acc = match rcd with
+  | [] -> acc
+  | (_, _, x)::xs -> (Int.max acc (msr_of_expr x)) |> msr_of_record xs
+
+let msr_of_lambda lam =
+  let msr_body = msr_of_expr lam.body in
+  let msr_pty = Option.fold ~none:0 ~some:(fun (_, x) -> (msr_of_expr x) + 1) lam.pty in
+  Int.max msr_body msr_pty
 
 (* a minimal expression are all parts which are "self contained" enough (e.g. not an apply) *)
 let rec parse_minexpr ps =