fix: preserve known cardinality when restoring stored sets

lib/persistent_sorted_set.ml

@@ -1889,14 +1889,18 @@ let store set =
       let address = storage.store_node (Leaf []) in
       (address, stored_set address)
 
-let restore ?(cmp = default_cmp) ?(settings = default_settings) storage address
-    =
+let restore ?count ?(cmp = default_cmp) ?(settings = default_settings) storage
+    address =
+  (match count with
+  | Some count when count < 0 ->
+      invalid_arg "restore count must be non-negative"
+  | Some _ | None -> ());
   let settings = validate_settings settings in
   Some
     {
       cmp = normalize_cmp cmp;
       set_settings = settings;
       set_storage = Some (cache_storage settings storage);
       data = Deferred { address };
-      count_cache = None;
+      count_cache = count;
     }

lib/persistent_sorted_set.mli

@@ -63,11 +63,18 @@ val seek : 'a -> 'a seq -> 'a seq
 val store : 'a t -> string * 'a t
 
 val restore :
+  ?count:int ->
   ?cmp:'a comparator ->
   ?settings:settings ->
   'a storage ->
   string ->
   'a t option
+(** [restore ~count storage address] restores a set lazily with its known
+    cardinality. [count] must be the exact number of elements in the snapshot at
+    [address], before any later edits. The caller is responsible for this
+    consistency: restore does not read nodes to verify it. Negative counts are
+    rejected. Without [count], cardinality is computed by traversing the set
+    when [count set] is called, preserving the existing restore behavior. *)
 
 (* DEBUG vals *)
 val validate_invariants : 'a t -> unit

test/js_smoke.ml

@@ -11,4 +11,25 @@ let () =
   if settings set <> { default_settings with branching_factor = 4 } then
     failwith "js smoke settings failed";
   if to_list (of_sorted_array [| 1; 1; 2; 3 |]) <> [ 1; 2; 3 ] then
-    failwith "js smoke sorted array failed"
+    failwith "js smoke sorted array failed";
+  let memory = Hashtbl.create 8 in
+  let next = ref 0 and reads = ref 0 in
+  let storage =
+    {
+      store_node =
+        (fun node ->
+          incr next;
+          let address = string_of_int !next in
+          Hashtbl.replace memory address node;
+          address);
+      restore_node =
+        (fun address ->
+          incr reads;
+          Hashtbl.find_opt memory address);
+      accessed = (fun _ -> ());
+    }
+  in
+  let root, _ = store (of_list_by ~storage [ 1; 2; 3 ]) in
+  let restored = Option.get (restore ~count:3 storage root) in
+  if count restored <> 3 || count restored <> 3 || !reads <> 0 then
+    failwith "known js_of_ocaml count loaded nodes"

test/melange_storage_smoke.ml

@@ -11,6 +11,7 @@ let require_equal_list label actual expected =
 let () =
   let memory = Hashtbl.create 8 in
   let next_id = ref 0 in
+  let reads = ref 0 in
   let storage =
     {
       PSet.store_node =
@@ -19,14 +20,20 @@ let () =
           let address = "node-" ^ string_of_int !next_id in
           Hashtbl.replace memory address node;
           address);
-      restore_node = (fun address -> Hashtbl.find_opt memory address);
+      restore_node =
+        (fun address ->
+          incr reads;
+          Hashtbl.find_opt memory address);
       accessed = (fun _address -> ());
     }
   in
   let set = PSet.of_list_by ~storage [ 3; 1; 2 ] in
   let root, _stored = PSet.store set in
-  match PSet.restore storage root with
+  match PSet.restore ~count:3 storage root with
   | None -> failwith "expected stored set to restore"
   | Some restored ->
-      require_equal_list "restored Melange stored set"
-        (PSet.to_list restored) [ 1; 2; 3 ]
+      if PSet.count restored <> 3 || PSet.count restored <> 3 then
+        failwith "restored Melange count differs";
+      if !reads <> 0 then failwith "known Melange count loaded nodes";
+      require_equal_list "restored Melange stored set" (PSet.to_list restored)
+        [ 1; 2; 3 ]

test/test_persistent_sorted_set.ml

@@ -1282,6 +1282,81 @@ let test_stored_count_uses_cached_size () =
   assert_equal_list "edited stored count should not mark stored nodes accessed"
     [] !accessed
 
+let count_restore_fixture values =
+  let nodes = Hashtbl.create 64 in
+  let writes = ref 0 and reads = ref 0 in
+  let storage =
+    {
+      store_node =
+        (fun node ->
+          incr writes;
+          let address = string_of_int !writes in
+          Hashtbl.replace nodes address node;
+          address);
+      restore_node =
+        (fun address ->
+          incr reads;
+          Hashtbl.find_opt nodes address);
+      accessed = (fun _ -> ());
+    }
+  in
+  let settings = { branching_factor = 4; ref_type = Strong } in
+  let root, _ = Pss.store (Pss.of_list_by ~settings ~storage values) in
+  (settings, storage, root, reads)
+
+let test_restore_known_count_is_lazy_and_tracks_edits () =
+  let settings, storage, root, reads =
+    count_restore_fixture (List.init 4096 Fun.id)
+  in
+  let restored = Option.get (Pss.restore ~count:4096 ~settings storage root) in
+  assert_equal_int "known-count restore is lazy" 0 !reads;
+  List.iter
+    (fun _ ->
+      assert_equal_int "known restored count" 4096 (Pss.count restored);
+      assert_equal_int "repeated count does not read nodes" 0 !reads)
+    [ 1; 2 ];
+  let added = Pss.add 4096 restored in
+  let after_edit = !reads in
+  assert_equal_int "cached add count" 4097 (Pss.count added);
+  assert_equal_int "count after add does not read nodes" after_edit !reads;
+  let duplicate = Pss.add 4096 added in
+  let missing = Pss.remove (-1) duplicate in
+  let removed = Pss.remove 10 missing in
+  let after_edit = !reads in
+  assert_equal_int "duplicate and missing deletion preserve count" 4097
+    (Pss.count missing);
+  assert_equal_int "successful deletion decrements count" 4096
+    (Pss.count removed);
+  assert_equal_int "edit counts do not read nodes" after_edit !reads;
+  assert_equal_int "persistent original count" 4096 (Pss.count restored)
+
+let test_restore_known_count_empty_and_invalid () =
+  let settings, storage, root, reads = count_restore_fixture [] in
+  let restored = Option.get (Pss.restore ~count:0 ~settings storage root) in
+  assert_equal_int "empty restored count" 0 (Pss.count restored);
+  assert_equal_int "empty count does not read nodes" 0 !reads;
+  let absent = Pss.remove 1 restored in
+  assert_equal_int "delete from empty count" 0 (Pss.count absent);
+  let singleton = Pss.add 1 absent in
+  assert_equal_int "add to restored empty count" 1 (Pss.count singleton);
+  assert_equal_int "remove singleton count" 0
+    (Pss.count (Pss.remove 1 singleton));
+  let before = !reads in
+  assert_raises_invalid_arg "negative known count" (fun () ->
+      ignore (Pss.restore ~count:(-1) ~settings storage root));
+  assert_equal_int "invalid count validation does not read nodes" before !reads
+
+let test_restore_without_count_remains_compatible () =
+  let settings, storage, root, reads =
+    count_restore_fixture (List.init 100 Fun.id)
+  in
+  let restored = Option.get (Pss.restore ~settings storage root) in
+  assert_equal_int "legacy restore remains lazy" 0 !reads;
+  assert_equal_int "legacy restore computes actual count" 100
+    (Pss.count restored);
+  if !reads = 0 then failwith "legacy count should load stored nodes";
+  assert_equal_int "legacy edited count" 99 (Pss.count (Pss.remove 10 restored))
+
 let test_storage_uses_leaf_and_branch_nodes_for_large_sets () =
   let memory = Hashtbl.create 16 in
   let writes = ref 0 in
@@ -2442,6 +2517,9 @@ let () =
   test_upstream_overflow_batched_insert_smoke ();
   test_storage_round_trip_and_stable_addresses ();
   test_stored_count_uses_cached_size ();
+  test_restore_known_count_is_lazy_and_tracks_edits ();
+  test_restore_known_count_empty_and_invalid ();
+  test_restore_without_count_remains_compatible ();
   test_storage_uses_leaf_and_branch_nodes_for_large_sets ();
   test_storage_remove_preserves_unchanged_leaf_addresses ();
   test_storage_remove_borrows_from_right_sibling ();