more informative trace of context allocations;

--- a/src/Pure/Concurrent/unsynchronized.ML	Mon May 08 21:11:01 2023 +0200
+++ b/src/Pure/Concurrent/unsynchronized.ML	Mon May 08 21:50:21 2023 +0200
@@ -7,6 +7,7 @@
 signature UNSYNCHRONIZED =
 sig
   datatype ref = datatype ref
+  type 'a weak_ref = 'a ref option ref
   val := : 'a ref * 'a -> unit
   val ! : 'a ref -> 'a
   val change: 'a ref -> ('a -> 'a) -> unit
@@ -22,6 +23,8 @@
 
 datatype ref = datatype ref;
 
+type 'a weak_ref = 'a ref option ref;
+
 val op := = op :=;
 val ! = !;
 

--- a/src/Pure/context.ML	Mon May 08 21:11:01 2023 +0200
+++ b/src/Pure/context.ML	Mon May 08 21:50:21 2023 +0200
@@ -53,8 +53,6 @@
   val proper_subthy: theory * theory -> bool
   val subthy_id: theory_id * theory_id -> bool
   val subthy: theory * theory -> bool
-  val trace_theories: bool Unsynchronized.ref
-  val theories_trace: unit -> {active_positions: Position.T list, active: int, total: int}
   val join_thys: theory list -> theory
   val begin_thy: string -> theory list -> theory
   val finish_thy: theory -> theory
@@ -69,6 +67,16 @@
   val join_certificate: certificate * certificate -> certificate
   (*generic context*)
   datatype generic = Theory of theory | Proof of Proof.context
+  val trace_theories: bool Unsynchronized.ref
+  val trace_proofs: bool Unsynchronized.ref
+  val allocations_trace: unit ->
+   {contexts: generic list,
+    active_contexts: int,
+    active_theories: int,
+    active_proofs: int,
+    total_contexts: int,
+    total_theories: int,
+    total_proofs: int}
   val cases: (theory -> 'a) -> (Proof.context -> 'a) -> generic -> 'a
   val mapping: (theory -> theory) -> (Proof.context -> Proof.context) -> generic -> generic
   val mapping_result: (theory -> 'a * theory) -> (Proof.context -> 'a * Proof.context) ->
@@ -156,24 +164,89 @@
 (* theory and proof context *)
 
 datatype theory =
-  Thy of
-   (*allocation state*)
-   state *
-   (*identity*)
-   {theory_id: theory_id,
-    token: Position.T Unsynchronized.ref} *
-   (*ancestry*)
-   {parents: theory list,         (*immediate predecessors*)
-    ancestors: theory list} *     (*all predecessors -- canonical reverse order*)
-   (*data*)
-   Any.T Datatab.table;           (*body content*)
+  Thy_Undef
+| Thy of
+    (*allocation state*)
+    state *
+    (*identity*)
+    {theory_id: theory_id,
+     theory_token: theory Unsynchronized.ref} *
+    (*ancestry*)
+    {parents: theory list,         (*immediate predecessors*)
+     ancestors: theory list} *     (*all predecessors -- canonical reverse order*)
+    (*data*)
+    Any.T Datatab.table;           (*body content*)
 
-datatype proof = Prf of Any.T Datatab.table * theory;
+datatype proof =
+  Prf_Undef
+| Prf of
+    (*identity*)
+    proof Unsynchronized.ref *
+    theory *
+    (*data*)
+    Any.T Datatab.table;
+
 structure Proof = struct type context = proof end;
 
 datatype generic = Theory of theory | Proof of Proof.context;
 
 
+(* heap allocations *)
+
+val trace_theories = Unsynchronized.ref false;
+val trace_proofs = Unsynchronized.ref false;
+
+local
+
+fun make_token guard var token0 =
+  if ! guard then
+    let
+      val token = Unsynchronized.ref (! token0);
+      val _ = Synchronized.change var (cons (Weak.weak (SOME token)));
+    in (token, fn res => (token := res; res)) end
+  else (token0, I);
+
+val theory_tokens = Synchronized.var "theory_tokens" ([]: theory Unsynchronized.weak_ref list);
+val proof_tokens = Synchronized.var "proof_tokens" ([]: Proof.context Unsynchronized.weak_ref list);
+
+val theory_token0 = Unsynchronized.ref Thy_Undef;
+val proof_token0 = Unsynchronized.ref Prf_Undef;
+
+in
+
+fun theory_token () = make_token trace_theories theory_tokens theory_token0;
+fun proof_token () = make_token trace_proofs proof_tokens proof_token0;
+
+fun allocations_trace () =
+  let
+    val _ = ML_Heap.full_gc ();
+    val trace1 = Synchronized.value theory_tokens;
+    val trace2 = Synchronized.value proof_tokens;
+
+    fun cons1 (Unsynchronized.ref (SOME (Unsynchronized.ref (thy as Thy _)))) = cons (Theory thy)
+      | cons1 _ = I;
+    fun cons2 (Unsynchronized.ref (SOME (Unsynchronized.ref (ctxt as Prf _)))) = cons (Proof ctxt)
+      | cons2 _ = I;
+
+    val contexts = build (fold cons1 trace1 #> fold cons2 trace2);
+    val active_theories = fold (fn Theory _ => Integer.add 1 | _ => I) contexts 0;
+    val active_proofs = fold (fn Proof _ => Integer.add 1 | _ => I) contexts 0;
+
+    val total_theories = length trace1;
+    val total_proofs = length trace2;
+  in
+    {contexts = contexts,
+     active_contexts = active_theories + active_proofs,
+     active_theories = active_theories,
+     active_proofs = active_proofs,
+     total_contexts = total_theories + total_proofs,
+     total_theories = total_theories,
+     total_proofs = total_proofs}
+  end;
+
+end;
+
+
 
 (*** theory operations ***)
 
@@ -348,45 +421,12 @@
 
 (* create theory *)
 
-val trace_theories = Unsynchronized.ref false;
-
-local
-
-val theories =
-  Synchronized.var "theory_tokens"
-    ([]: Position.T Unsynchronized.ref option Unsynchronized.ref list);
-
-val dummy_token = Unsynchronized.ref Position.none;
-
-fun make_token () =
-  if ! trace_theories then
-    let
-      val token = Unsynchronized.ref (Position.thread_data ());
-      val _ = Synchronized.change theories (cons (Weak.weak (SOME token)));
-    in token end
-  else dummy_token;
-
-in
-
-fun theories_trace () =
-  let
-    val trace = Synchronized.value theories;
-    val _ = ML_Heap.full_gc ();
-    val active_positions =
-      fold (fn Unsynchronized.ref (SOME pos) => cons (! pos) | _ => I) trace [];
-  in
-    {active_positions = active_positions,
-     active = length active_positions,
-     total = length trace}
-  end;
-
 fun create_thy state ids name stage ancestry data =
   let
     val theory_id = make_theory_id {id = new_id (), ids = ids, name = name, stage = stage};
-    val identity = {theory_id = theory_id, token = make_token ()};
-  in Thy (state, identity, ancestry, data) end;
-
-end;
+    val (token, assign) = theory_token ();
+    val identity = {theory_id = theory_id, theory_token = token};
+  in assign (Thy (state, identity, ancestry, data)) end;
 
 
 (* primitives *)
@@ -515,16 +555,21 @@
 
 in
 
-fun raw_transfer thy' (Prf (data, thy)) =
-  let
-    val _ = subthy (thy, thy') orelse error "Cannot transfer proof context: not a super theory";
-    val data' = init_new_data thy' data;
-  in Prf (data', thy') end;
+fun raw_transfer thy' (ctxt as Prf (_, thy, data)) =
+  if eq_thy (thy, thy') then ctxt
+  else if proper_subthy (thy, thy') then
+    let
+      val (token', assign) = proof_token ();
+      val data' = init_new_data thy' data;
+    in assign (Prf (token', thy', data')) end
+  else error "Cannot transfer proof context: not a super theory";
 
 structure Proof_Context =
 struct
-  fun theory_of (Prf (_, thy)) = thy;
-  fun init_global thy = Prf (init_data thy, thy);
+  fun theory_of (Prf (_, thy, _)) = thy;
+  fun init_global thy =
+    let val (token, assign) = proof_token ()
+    in assign (Prf (token, thy, init_data thy)) end;
   fun get_global long thy name = init_global (get_theory long thy name);
 end;
 
@@ -537,13 +582,14 @@
     val _ = Synchronized.change kinds (Datatab.update (k, init));
   in k end;
 
-fun get k dest (Prf (data, thy)) =
+fun get k dest (Prf (_, thy, data)) =
   (case Datatab.lookup data k of
     SOME x => x
   | NONE => init_fallback k thy) |> dest;
 
-fun put k make x (Prf (data, thy)) =
-  Prf (Datatab.update (k, make x) data, thy);
+fun put k make x (Prf (token, thy, data)) =
+  let val (token, assign) = proof_token ()
+  in assign (Prf (token, thy, Datatab.update (k, make x) data)) end;
 
 end;