<commit_before>(*i camlp4use: "q_MLast.cmo" i*)

open Names
open Term
open Environ
open Pre_env
open Pcaml
open Declarations


let id = ref 0
let uniq = ref 256

let loc = Ploc.dummy

let rec drop n xs = match n, xs with
  | 0, _ -> xs
  | _, [] -> []
  | n, x :: xs when n > 0 -> drop (n - 1) xs
  | _, _ -> raise (Invalid_argument "drop")

let rec take n xs = match n, xs with
  | 0, _ -> []
  | _, [] -> []
  | n, x :: xs when n > 0 -> x :: take (n - 1) xs
  | _, _ -> raise (Invalid_argument "take")

(* Flag to signal whether recompilation is needed. *)
let env_updated = ref false

let descend_ast f t = match t with
  | <:expr< Con $_$ >> -> t
  | <:expr< Lam (fun $x$ -> $body$) >> -> <:expr< Lam (fun $x$ -> $f body$) >>
  | <:expr< Prod $dom$ (fun $lid:x$ -> $cod$) >> -> <:expr< Prod $f dom$ (fun $lid:x$ -> $f cod$) >>
  | <:expr< app $t1$ $t2$ >> -> <:expr< app $f t1$ $f t2$ >>
  | <:expr< Const $i$ [| $list:args$ |] >> ->
      <:expr< Const $i$ [| $list:List.fold_right (fun t ts -> f t :: ts) args []$ |] >>
  | <:expr< Set >> -> t
  | <:expr< Prop >> -> t
  | <:expr< Type $_$ >> -> t
  | <:expr< $lid:x$ >> -> t
  | <:expr< match $scrutinee$ with [ $list:branches$ ] >> ->
      let bs = List.fold_right (fun (p, w, body) bs -> (p, w, f body) :: bs) branches [] in
	<:expr< match $f scrutinee$ with [ $list:bs$ ] >>
  | <:expr< let $flag:r$ $list:defs$ in $t$ >> ->
      let defs' = List.map (fun (p, t) -> (p, f t)) defs in
	<:expr< let $flag:r$ $list:defs'$ in $f t$ >>
  | <:expr< do { $list:_$ } >> -> t
  | _ -> let s = Eprinter.apply Pcaml.pr_expr Pprintf.empty_pc t
    in print_endline s; raise (Invalid_argument "descend_ast")

let subst rho x = try List.assoc x rho with Not_found -> x

let rec shrink rho = function
  | <:expr< app $t1$ $lid:y$ >> ->
    (match shrink rho t1 with
       | <:expr< Lam (fun $lid:x$ -> $body$) >> ->
	 shrink ((x, y) :: rho) body
       | t -> <:expr< app $t$ $lid:subst rho y$ >>)
  | <:expr< $lid:x$ >> -> <:expr< $lid:subst rho x$ >>
  | t -> descend_ast (shrink rho) t

let rec pop_abstractions n =
  if n > 0 then function
    | <:expr< Lam (fun $x$ -> $t$) >> ->
      let vars, body = pop_abstractions (n - 1) t in
	(x :: vars, body)
    | t -> ([], t)
  else function t -> ([], t)
	
let rec pop_applications n t =
  let rec aux n args =
    if n > 0 then function
      | <:expr< app $t1$ $t2$ >> ->
	aux (n - 1) (t2 :: args) t1
      | f -> (f, args)
    else function f -> (f, args)
  in aux n [] t

let rec collapse_abstractions t =
  let vars, body = pop_abstractions 6 t in
    if List.length vars = 0 then uncurrify body else
      let func = List.fold_right (fun x t -> <:expr< fun $x$ -> $t$ >>) vars (collapse_abstractions body) in
	<:expr< $uid:"Lam" ^ string_of_int (List.length vars)$ $func$ >>

and collapse_applications t =
  let f, args = pop_applications 6 t in
  let f, args = uncurrify f, List.map uncurrify args in
    match args with
      | [t1] -> <:expr< app1 $f$ $t1$ >>
      | [t1; t2] -> <:expr< app2 $f$ $t1$ $t2$ >>
      | [t1; t2; t3] -> <:expr< app3 $f$ $t1$ $t2$ $t3$ >>
      | [t1; t2; t3; t4] -> <:expr< app4 $f$ $t1$ $t2$ $t3$ $t4$ >>
      | [t1; t2; t3; t4; t5] -> <:expr< app5 $f$ $t1$ $t2$ $t3$ $t4$ $t5$ >>
      | [t1; t2; t3; t4; t5; t6] -> <:expr< app6 $f$ $t1$ $t2$ $t3$ $t4$ $t5$ $t6$ >>
      | _ -> assert false

and uncurrify t = match t with
  | <:expr< Lam $_$ >> -> collapse_abstractions t
  | <:expr< app $_$ $_$ >> -> collapse_applications t
  | _ -> descend_ast uncurrify t

let lid_of_string s = print_endline ("lid: " ^ ("z" ^ s)); "z" ^ s
let uid_of_string s = "X" ^ s

let lid_of_name = function
  | Name id -> lid_of_string (string_of_id id)
  | Anonymous -> "_"

(* Redefine a bunch of functions in module Names to generate names
   acceptable to OCaml. *)
let string_of_dirpath = function
  | [] -> "_"
  | sl -> String.concat "_" (List.map string_of_id (List.rev sl))

let rec string_of_mp = function
  | MPfile sl -> string_of_dirpath (repr_dirpath sl)
  | MPbound uid -> string_of_mbid uid
  | MPdot (mp,l) -> string_of_mp mp ^ "." ^ string_of_label l

let string_of_kn kn =
  let (modpath, _dirpath, label) = repr_kn kn in
    string_of_mp modpath ^ "_" ^ string_of_label label

let string_of_con con =
  let (modpath, _dirpath, label) = repr_con con in
    string_of_mp modpath ^ "_" ^ string_of_label label

let string_of_inductive (mind, i) =
  string_of_mind mind ^ string_of_int i

let string_of_constructor (ind, i) =
  string_of_inductive ind ^ "c" ^ string_of_int i

(* First argument the index of the constructor *)
let make_constructor_pattern i args =
  let f arg = <:patt< $lid:arg$ >> in
    <:patt< Const $int:string_of_int i$ [| $list:List.map f args$ |] >>

exception NonNat
exception Pass

let rec is_nat_literal ind c =
  let rec f c = match c with
    | App (c', args) when isConstruct c' -> 1 + f (kind_of_term args.(0))
    | Construct cstr -> 0
    | _ -> raise NonNat in
    match string_of_inductive ind with
      | "Coq_Init_Datatype_nat" -> f c
      | _ -> raise NonNat

let rec gen_names = function
  | 0 -> []
  | n when n > 0 ->
      let xs = gen_names (n - 1) in
	uniq := !uniq + 1;
	("x" ^ string_of_int !uniq) :: xs
  | _ -> raise (Invalid_argument "gen_names")

exception Bound

let is_bound x env =
  try
    List.iter (fun (y, _, _) -> if x = y then raise Bound else ())
      (pre_env env).env_rel_context; false
  with Bound -> true

let rec push_value id body env =
  env_updated := true;
  let kind = lookup_named_val id (pre_env env) in
    match !kind with
    | VKvalue (v, _) -> ()
    | VKnone -> kind := VKvalue (values (translate env body), Idset.empty)

and translate_constant env c =
  let cb = lookup_constant c (pre_env env) in
    match cb.const_body with
      | Some body -> (match cb.const_body_ast with
	  | Some _ ->
	      (* <:expr< do { print_endline $str:string_of_con c$; $lid:lid_of_string (string_of_con c)$ } >> *)
	      <:expr< $lid:lid_of_string (string_of_con c)$ >>
	  | None -> print_endline ("translating: " ^ string_of_con c);
	      let ast = translate env (Declarations.force body) in
	      cb.const_body_ast <- Some (values (uncurrify (shrink [] ast)));
	      env_updated := true;
	      (* <:expr< do { print_endline $str:string_of_con c$; $lid:lid_of_string (string_of_con c)$ } >>) *)
	      <:expr< $lid:lid_of_string (string_of_con c)$ >>)
      | None -> <:expr< Con () (* $str:string_of_con c$ *) >>

(** The side-effect of translate is to add the translated construction
    to the value environment. *)
(* A simple counter is used for fresh variables. We effectively encode
   de Bruijn indices as de Bruijn levels. *)
and translate (env : Environ.env) t =
  try (
    match kind_of_term t with
    | Rel x ->
	print_endline ("rel: " ^ lid_of_name (let name, _, _ = lookup_rel x env in name));
	<:expr< $lid: lid_of_name (match lookup_rel x env with name, _, _ -> name)$ >>
    | Var id ->
	print_endline ("variable: " ^ string_of_id id);
	let v = <:expr< $lid:lid_of_string (string_of_id id)$ >> in
	  (match named_body id env with
	       (* Add compiled form of definition to environment if not already present. *)
	     | Some body -> push_value id body env; v
	     | None -> v)
    | Sort (Prop Null) -> <:expr< Prop >>
    | Sort (Prop Pos) -> <:expr< Set >>
    | Sort (Type _) -> <:expr< Type 0 >>
    | Cast (c, _, ty) ->
	translate env c
    | Prod (x, t, c) ->
	let x' = if is_bound x env then
	  (match x with
	     | Name id -> uniq := !uniq + 1; Name (id_of_string (string_of_id id ^ string_of_int !uniq))
	     | Anonymous -> Anonymous) else x in
	let newenv = Environ.push_rel (x', None, t) env in
	  <:expr< Prod $translate env t$ (fun $lid:lid_of_name x'$ -> $translate newenv c$) >>
    | Lambda (x, t, c) ->
	let x' = if is_bound x env then
	  (match x with
	     | Name id -> uniq := !uniq + 1; Name (id_of_string (string_of_id id ^ string_of_int !uniq))
	     | Anonymous -> Anonymous) else x in
	let newenv = Environ.push_rel (x', None, t) env in
	  <:expr< Lam (fun $lid:lid_of_name x'$ -> $translate newenv c$) >>
    | LetIn (x, b, t, c) ->
	let x' = if is_bound x env then
	  (match x with
	     | Name id -> uniq := !uniq + 1; Name (id_of_string (string_of_id id ^ string_of_int !uniq))
	     | Anonymous -> Anonymous) else x in
	let newenv = Environ.push_rel (x', Some b, t) env in
	  <:expr< let $lid:lid_of_name x'$ = $translate env b$ in $translate newenv c$ >>
    | App (c, args) ->
	(match kind_of_term c with
	   | Construct cstr ->
	       let f arg vs = translate env arg :: vs in
	       let i = index_of_constructor cstr in
	       let ind = inductive_of_constructor cstr in
	       let mb = lookup_mind (fst ind) (pre_env env) in
	       let ob = mb.mind_packets.(snd ind) in
	       let nparams = rel_context_length ob.mind_arity_ctxt in
	       let nargs = Array.length args in
	       let vs = drop (min nparams nargs) (Array.fold_right f args []) in
	       let missing = ob.mind_consnrealdecls.(i-1) - nargs + nparams in
	       let underscores = max (nparams - nargs) 0 in
	       let names = gen_names (missing - underscores) in
	       let pad = List.map (fun x -> <:expr< $lid:x$ >>) names in
	       print_endline ("app constructor " ^ string_of_constructor cstr ^ string_of_int i ^ string_of_int ob.mind_consnrealdecls.(i-1));
	       print_endline ("app constructor nparams " ^ string_of_int nparams);
	       let prefix =
		 let rec aux n z =
		   if n = 0
		   then z
		   else aux (n - 1) <:expr< Lam (fun _ -> $z$) >>
		 in aux underscores
		      (List.fold_left (fun e x -> <:expr< Lam (fun $lid:x$ -> $e$) >>)
		       <:expr< Const $int:string_of_int i$ [| $list: vs @ pad$ |] >>
		       names)
	       in assert (List.length vs + List.length pad = ob.mind_consnrealdecls.(i-1)); prefix
  		 (* (try let n = is_nat_literal ind t in *)
		 (*    <:expr< __from_nat $int:string_of_int n$ >> *)
		 (*  with NonNat -> <:expr< Const $int:string_of_int i$ [| $list:vs'$ |] >>) *)<commit_msg>Use util.ml's list_skipn rather than custom 'drop' function.<commit_after>(*i camlp4use: "q_MLast.cmo" i*)

open Names
open Term
open Environ
open Pre_env
open Pcaml
open Declarations
open Util


let id = ref 0
let uniq = ref 256

let loc = Ploc.dummy

(* Flag to signal whether recompilation is needed. *)
let env_updated = ref false

let descend_ast f t = match t with
  | <:expr< Con $_$ >> -> t
  | <:expr< Lam (fun $x$ -> $body$) >> -> <:expr< Lam (fun $x$ -> $f body$) >>
  | <:expr< Prod $dom$ (fun $lid:x$ -> $cod$) >> -> <:expr< Prod $f dom$ (fun $lid:x$ -> $f cod$) >>
  | <:expr< app $t1$ $t2$ >> -> <:expr< app $f t1$ $f t2$ >>
  | <:expr< Const $i$ [| $list:args$ |] >> ->
      <:expr< Const $i$ [| $list:List.fold_right (fun t ts -> f t :: ts) args []$ |] >>
  | <:expr< Set >> -> t
  | <:expr< Prop >> -> t
  | <:expr< Type $_$ >> -> t
  | <:expr< $lid:x$ >> -> t
  | <:expr< match $scrutinee$ with [ $list:branches$ ] >> ->
      let bs = List.fold_right (fun (p, w, body) bs -> (p, w, f body) :: bs) branches [] in
	<:expr< match $f scrutinee$ with [ $list:bs$ ] >>
  | <:expr< let $flag:r$ $list:defs$ in $t$ >> ->
      let defs' = List.map (fun (p, t) -> (p, f t)) defs in
	<:expr< let $flag:r$ $list:defs'$ in $f t$ >>
  | <:expr< do { $list:_$ } >> -> t
  | _ -> let s = Eprinter.apply Pcaml.pr_expr Pprintf.empty_pc t
    in print_endline s; raise (Invalid_argument "descend_ast")

let subst rho x = try List.assoc x rho with Not_found -> x

let rec shrink rho = function
  | <:expr< app $t1$ $lid:y$ >> ->
    (match shrink rho t1 with
       | <:expr< Lam (fun $lid:x$ -> $body$) >> ->
	 shrink ((x, y) :: rho) body
       | t -> <:expr< app $t$ $lid:subst rho y$ >>)
  | <:expr< $lid:x$ >> -> <:expr< $lid:subst rho x$ >>
  | t -> descend_ast (shrink rho) t

let rec pop_abstractions n =
  if n > 0 then function
    | <:expr< Lam (fun $x$ -> $t$) >> ->
      let vars, body = pop_abstractions (n - 1) t in
	(x :: vars, body)
    | t -> ([], t)
  else function t -> ([], t)
	
let rec pop_applications n t =
  let rec aux n args =
    if n > 0 then function
      | <:expr< app $t1$ $t2$ >> ->
	aux (n - 1) (t2 :: args) t1
      | f -> (f, args)
    else function f -> (f, args)
  in aux n [] t

let rec collapse_abstractions t =
  let vars, body = pop_abstractions 6 t in
    if List.length vars = 0 then uncurrify body else
      let func = List.fold_right (fun x t -> <:expr< fun $x$ -> $t$ >>) vars (collapse_abstractions body) in
	<:expr< $uid:"Lam" ^ string_of_int (List.length vars)$ $func$ >>

and collapse_applications t =
  let f, args = pop_applications 6 t in
  let f, args = uncurrify f, List.map uncurrify args in
    match args with
      | [t1] -> <:expr< app1 $f$ $t1$ >>
      | [t1; t2] -> <:expr< app2 $f$ $t1$ $t2$ >>
      | [t1; t2; t3] -> <:expr< app3 $f$ $t1$ $t2$ $t3$ >>
      | [t1; t2; t3; t4] -> <:expr< app4 $f$ $t1$ $t2$ $t3$ $t4$ >>
      | [t1; t2; t3; t4; t5] -> <:expr< app5 $f$ $t1$ $t2$ $t3$ $t4$ $t5$ >>
      | [t1; t2; t3; t4; t5; t6] -> <:expr< app6 $f$ $t1$ $t2$ $t3$ $t4$ $t5$ $t6$ >>
      | _ -> assert false

and uncurrify t = match t with
  | <:expr< Lam $_$ >> -> collapse_abstractions t
  | <:expr< app $_$ $_$ >> -> collapse_applications t
  | _ -> descend_ast uncurrify t

let lid_of_string s = print_endline ("lid: " ^ ("z" ^ s)); "z" ^ s
let uid_of_string s = "X" ^ s

let lid_of_name = function
  | Name id -> lid_of_string (string_of_id id)
  | Anonymous -> "_"

(* Redefine a bunch of functions in module Names to generate names
   acceptable to OCaml. *)
let string_of_dirpath = function
  | [] -> "_"
  | sl -> String.concat "_" (List.map string_of_id (List.rev sl))

let rec string_of_mp = function
  | MPfile sl -> string_of_dirpath (repr_dirpath sl)
  | MPbound uid -> string_of_mbid uid
  | MPdot (mp,l) -> string_of_mp mp ^ "." ^ string_of_label l

let string_of_kn kn =
  let (modpath, _dirpath, label) = repr_kn kn in
    string_of_mp modpath ^ "_" ^ string_of_label label

let string_of_con con =
  let (modpath, _dirpath, label) = repr_con con in
    string_of_mp modpath ^ "_" ^ string_of_label label

let string_of_inductive (mind, i) =
  string_of_mind mind ^ string_of_int i

let string_of_constructor (ind, i) =
  string_of_inductive ind ^ "c" ^ string_of_int i

(* First argument the index of the constructor *)
let make_constructor_pattern i args =
  let f arg = <:patt< $lid:arg$ >> in
    <:patt< Const $int:string_of_int i$ [| $list:List.map f args$ |] >>

exception NonNat
exception Pass

let rec is_nat_literal ind c =
  let rec f c = match c with
    | App (c', args) when isConstruct c' -> 1 + f (kind_of_term args.(0))
    | Construct cstr -> 0
    | _ -> raise NonNat in
    match string_of_inductive ind with
      | "Coq_Init_Datatype_nat" -> f c
      | _ -> raise NonNat

let rec gen_names = function
  | 0 -> []
  | n when n > 0 ->
      let xs = gen_names (n - 1) in
	uniq := !uniq + 1;
	("x" ^ string_of_int !uniq) :: xs
  | _ -> raise (Invalid_argument "gen_names")

exception Bound

let is_bound x env =
  try
    List.iter (fun (y, _, _) -> if x = y then raise Bound else ())
      (pre_env env).env_rel_context; false
  with Bound -> true

let rec push_value id body env =
  env_updated := true;
  let kind = lookup_named_val id (pre_env env) in
    match !kind with
    | VKvalue (v, _) -> ()
    | VKnone -> kind := VKvalue (values (translate env body), Idset.empty)

and translate_constant env c =
  let cb = lookup_constant c (pre_env env) in
    match cb.const_body with
      | Some body -> (match cb.const_body_ast with
	  | Some _ ->
	      (* <:expr< do { print_endline $str:string_of_con c$; $lid:lid_of_string (string_of_con c)$ } >> *)
	      <:expr< $lid:lid_of_string (string_of_con c)$ >>
	  | None -> print_endline ("translating: " ^ string_of_con c);
	      let ast = translate env (Declarations.force body) in
	      cb.const_body_ast <- Some (values (uncurrify (shrink [] ast)));
	      env_updated := true;
	      (* <:expr< do { print_endline $str:string_of_con c$; $lid:lid_of_string (string_of_con c)$ } >>) *)
	      <:expr< $lid:lid_of_string (string_of_con c)$ >>)
      | None -> <:expr< Con () (* $str:string_of_con c$ *) >>

(** The side-effect of translate is to add the translated construction
    to the value environment. *)
(* A simple counter is used for fresh variables. We effectively encode
   de Bruijn indices as de Bruijn levels. *)
and translate (env : Environ.env) t =
  try (
    match kind_of_term t with
    | Rel x ->
	print_endline ("rel: " ^ lid_of_name (let name, _, _ = lookup_rel x env in name));
	<:expr< $lid: lid_of_name (match lookup_rel x env with name, _, _ -> name)$ >>
    | Var id ->
	print_endline ("variable: " ^ string_of_id id);
	let v = <:expr< $lid:lid_of_string (string_of_id id)$ >> in
	  (match named_body id env with
	       (* Add compiled form of definition to environment if not already present. *)
	     | Some body -> push_value id body env; v
	     | None -> v)
    | Sort (Prop Null) -> <:expr< Prop >>
    | Sort (Prop Pos) -> <:expr< Set >>
    | Sort (Type _) -> <:expr< Type 0 >>
    | Cast (c, _, ty) ->
	translate env c
    | Prod (x, t, c) ->
	let x' = if is_bound x env then
	  (match x with
	     | Name id -> uniq := !uniq + 1; Name (id_of_string (string_of_id id ^ string_of_int !uniq))
	     | Anonymous -> Anonymous) else x in
	let newenv = Environ.push_rel (x', None, t) env in
	  <:expr< Prod $translate env t$ (fun $lid:lid_of_name x'$ -> $translate newenv c$) >>
    | Lambda (x, t, c) ->
	let x' = if is_bound x env then
	  (match x with
	     | Name id -> uniq := !uniq + 1; Name (id_of_string (string_of_id id ^ string_of_int !uniq))
	     | Anonymous -> Anonymous) else x in
	let newenv = Environ.push_rel (x', None, t) env in
	  <:expr< Lam (fun $lid:lid_of_name x'$ -> $translate newenv c$) >>
    | LetIn (x, b, t, c) ->
	let x' = if is_bound x env then
	  (match x with
	     | Name id -> uniq := !uniq + 1; Name (id_of_string (string_of_id id ^ string_of_int !uniq))
	     | Anonymous -> Anonymous) else x in
	let newenv = Environ.push_rel (x', Some b, t) env in
	  <:expr< let $lid:lid_of_name x'$ = $translate env b$ in $translate newenv c$ >>
    | App (c, args) ->
	(match kind_of_term c with
	   | Construct cstr ->
	       let f arg vs = translate env arg :: vs in
	       let i = index_of_constructor cstr in
	       let ind = inductive_of_constructor cstr in
	       let mb = lookup_mind (fst ind) (pre_env env) in
	       let ob = mb.mind_packets.(snd ind) in
	       let nparams = rel_context_length ob.mind_arity_ctxt in
	       let nargs = Array.length args in
	       let vs = list_skipn (min nparams nargs) (Array.fold_right f args []) in
	       let missing = ob.mind_consnrealdecls.(i-1) - nargs + nparams in
	       let underscores = max (nparams - nargs) 0 in
	       let names = gen_names (missing - underscores) in
	       let pad = List.map (fun x -> <:expr< $lid:x$ >>) names in
	       print_endline ("app constructor " ^ string_of_constructor cstr ^ string_of_int i ^ string_of_int ob.mind_consnrealdecls.(i-1));
	       print_endline ("app constructor nparams " ^ string_of_int nparams);
	       let prefix =
		 let rec aux n z =
		   if n = 0
		   then z
		   else aux (n - 1) <:expr< Lam (fun _ -> $z$) >>
		 in aux underscores
		      (List.fold_left (fun e x -> <:expr< Lam (fun $lid:x$ -> $e$) >>)
		       <:expr< Const $int:string_of_int i$ [| $list: vs @ pad$ |] >>
		       names)
	       in assert (List.length vs + List.length pad = ob.mind_consnrealdecls.(i-1)); prefix
  		 (* (try let n = is_nat_literal ind t in *)
		 (*    <:expr< __from_nat $int:string_of_int n$ >> *)
		 (*  with NonNat -> <:expr< Const $int:string_of_int i$ [| $list:vs'$ |] >>) *)
{'fraction_non_alphanumeric': 0.1566626041031891, 'fraction_numerical': 0.009140767824497258, 'mean_word_length': 2.7783960092095166, 'pattern_counts': {'":': 0, '<': 229, '<?xml version=': 0, '>': 435, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '76147', 'n_tokens_mistral': 7784, 'n_tokens_neox': 7517, 'n_words': 2449}