Optimize pretty-printing of notations

The printer now does a eager filtering of notations (only
notations that share the same head symbol as the printed
formula are considered)

The does not change the behavior of the pretty-printer,
but greatly improves performance.

Author: strub

src/ecEnv.ml

@@ -173,9 +173,7 @@ end = struct
     }
 end
 
-
 (* -------------------------------------------------------------------- *)
-
 type preenv = {
   env_top      : EcPath.path option;
   env_gstate   : EcGState.gstate;
@@ -191,7 +189,7 @@ type preenv = {
   env_rwbase   : Sp.t Mip.t;
   env_atbase   : (path list Mint.t) Msym.t;
   env_redbase  : mredinfo;
-  env_ntbase   : (path * env_notation) list;
+  env_ntbase   : ntbase Mop.t;
   env_modlcs   : Sid.t;                 (* declared modules *)
   env_item     : theory_item list;      (* in reverse order *)
   env_norm     : env_norm ref;
@@ -222,6 +220,8 @@ and mredinfo = redinfo Mrd.t
 
 and env_notation = ty_params * EcDecl.notation
 
+and ntbase = (path * env_notation) list
+
 (* -------------------------------------------------------------------- *)
 type env = preenv
 
@@ -309,7 +309,7 @@ let empty gstate =
     env_rwbase   = Mip.empty;
     env_atbase   = Msym.empty;
     env_redbase  = Mrd.empty;
-    env_ntbase   = [];
+    env_ntbase   = Mop.empty;
     env_modlcs   = Sid.empty;
     env_item     = [];
     env_norm     = ref empty_norm_cache; }
@@ -2837,19 +2837,33 @@ module Op = struct
   let lookup_path name env =
     fst (lookup name env)
 
-  let bind ?(import = import0) name op env =
-    let env = if import.im_immediate then MC.bind_operator name op env else env in
-    let op  = NormMp.norm_op env op in
+  let update_ntbase path (name, op) base =
     let nt  =
       match op.op_kind with
-      | OB_nott nt ->
-         Some (EcPath.pqname (root env) name, (op.op_tparams, nt))
+      | OB_nott nt -> begin
+        let head =
+          match nt.ont_body.e_node with
+          | Eapp ({ e_node = Eop (p, _)}, _) | Eop (p, _) -> Some p
+          | _ -> None
+        in
+        Some (head, (EcPath.pqname path name, (op.op_tparams, nt)))
+      end
       | _ -> None
     in
 
+    ofold
+      (fun (hd, nt) nts ->
+        Mop.change (fun nts -> Some (nt :: odfl [] nts)) hd nts)
+      base nt
+
+  let bind ?(import = import0) name op env =
+    let env = if import.im_immediate then MC.bind_operator name op env else env in
+    let op  = NormMp.norm_op env op in
+    let env_ntbase = update_ntbase (root env) (name, op) env.env_ntbase in
+
     { env with
-        env_ntbase = ofold List.cons env.env_ntbase nt;
-        env_item   = mkitem import (Th_operator (name, op)) :: env.env_item; }
+        env_ntbase;
+        env_item = mkitem import (Th_operator (name, op)) :: env.env_item; }
 
   let rebind name op env =
     MC.bind_operator name op env
@@ -2931,8 +2945,8 @@ module Op = struct
 
   type notation = env_notation
 
-  let get_notations env =
-    env.env_ntbase
+  let get_notations ~(head : path option) (env : env) =
+    Mop.find_def [] head env.env_ntbase
 
   let iter ?name f (env : env) =
     gen_iter (fun mc -> mc.mc_operators) MC.lookup_operators ?name f env
@@ -3214,8 +3228,8 @@ module Theory = struct
   (* ------------------------------------------------------------------ *)
   let bind_nt_th =
     let for1 path base = function
-      | Th_operator (x, ({ op_kind = OB_nott nt } as op)) ->
-         Some ((EcPath.pqname path x, (op.op_tparams, nt)) :: base)
+      | Th_operator (x, ({ op_kind = OB_nott _ } as op)) ->
+        Some (Op.update_ntbase path (x, op) base)
       | _ -> None
 
     in bind_base_th for1

src/ecEnv.mli

@@ -341,7 +341,7 @@ module Op : sig
 
   type notation = ty_params * EcDecl.notation
 
-  val get_notations : env -> (path * notation) list
+  val get_notations : head:path option -> env -> (path * notation) list
 
   val iter : ?name:qsymbol -> (path -> t -> unit) -> env -> unit
   val all  : ?check:(path -> t -> bool) -> ?name:qsymbol -> env -> (path * t) list

src/ecPath.ml

@@ -70,6 +70,14 @@ module Sp = struct
     List.ksort ~key:identity ~cmp:p_ntr_compare (elements s)
 end
 
+(* -------------------------------------------------------------------- *)
+module Mop = Map.Make(struct
+  type t = path option
+
+  let compare (p1 : path option) (p2 : path option) =
+    ocompare p_compare p1 p2
+end)
+
 (* -------------------------------------------------------------------- *)
 let mk_path node =
   Hspath.hashcons { p_node = node; p_tag = -1; }

src/ecPath.mli

@@ -47,6 +47,9 @@ module Sp : sig
   val ntr_elements : t -> elt list
 end
 
+(* -------------------------------------------------------------------- *)
+module Mop : Map.S with type key = path option
+
 (* -------------------------------------------------------------------- *)
 type mpath = private {
   m_top  : mpath_top;

src/ecPrinting.ml

@@ -1618,7 +1618,13 @@ and try_pp_notations (ppe : PPEnv.t) outer fmt f =
      not nt.ont_ponly && try_notation args
   in
 
-  let nts = EcEnv.Op.get_notations ppe.PPEnv.ppe_env in
+  let head =
+    try
+      Some (fst (destr_op (fst (destr_app f))))
+    with DestrError _ -> None
+  in
+
+  let nts = EcEnv.Op.get_notations ~head ppe.PPEnv.ppe_env in
 
   List.exists try_notation nts