diff --git a/lib/persistent_sorted_set.ml b/lib/persistent_sorted_set.ml index a849dd3..0122ddd 100644 --- a/lib/persistent_sorted_set.ml +++ b/lib/persistent_sorted_set.ml @@ -1889,8 +1889,12 @@ 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 { @@ -1898,5 +1902,5 @@ let restore ?(cmp = default_cmp) ?(settings = default_settings) storage address set_settings = settings; set_storage = Some (cache_storage settings storage); data = Deferred { address }; - count_cache = None; + count_cache = count; } diff --git a/lib/persistent_sorted_set.mli b/lib/persistent_sorted_set.mli index f1f5102..d17c408 100644 --- a/lib/persistent_sorted_set.mli +++ b/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 diff --git a/test/js_smoke.ml b/test/js_smoke.ml index e4eda93..00fe475 100644 --- a/test/js_smoke.ml +++ b/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" diff --git a/test/melange_storage_smoke.ml b/test/melange_storage_smoke.ml index ee14e6f..326289c 100644 --- a/test/melange_storage_smoke.ml +++ b/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 ] diff --git a/test/test_persistent_sorted_set.ml b/test/test_persistent_sorted_set.ml index 9429fef..3e13a43 100644 --- a/test/test_persistent_sorted_set.ml +++ b/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 ();