Shell: Memoize locator computations

lib_node_shell/block_locator.ml

@@ -7,11 +7,12 @@
 (*                                                                        *)
 (**************************************************************************)
 
-open State
+type t = raw
 
-type t = Block_header.t * Block_hash.t list
+(** Non private version of Block_store_locator.t for coercions *)
+and raw = Block_header.t * Block_hash.t list
 
-type error += Invalid_locator of P2p.Peer_id.t * t
+let raw x = x
 
 let encoding =
   let open Data_encoding in
@@ -20,108 +21,64 @@ let encoding =
      (req "current_head" (dynamic_size Block_header.encoding))
      (req "history" (dynamic_size (list Block_hash.encoding))))
 
-let compute (b: Block.t) sz =
+let predecessor (store : Store.Block.store) (b: Block_hash.t) =
-  let rec loop acc sz step cpt b =
+  Store.Block.Contents.read_exn (store, b) >>= fun contents ->
+  let predecessor = contents.header.shell.predecessor in
+  if Block_hash.equal b predecessor then
+    Lwt.return_none
+  else
+    Lwt.return_some predecessor
+
+let compute (store : Store.Block.store) (b: Block_hash.t) sz =
+  let rec loop acc ~sz step cpt b =
     if sz = 0 then
       Lwt.return (List.rev acc)
     else
-      Block.predecessor b >>= function
+      predecessor store b >>= function
       | None ->
-          Lwt.return (List.rev (Block.hash b :: acc))
+          Lwt.return (List.rev (b :: acc))
       | Some predecessor ->
           if cpt = 0 then
-            loop (Block.hash b :: acc) (sz - 1)
+            loop (b :: acc) ~sz:(sz - 1)
               (step * 2) (step * 20 - 1) predecessor
           else if cpt mod step = 0 then
-            loop (Block.hash b :: acc) (sz - 1)
+            loop (b :: acc) ~sz:(sz - 1)
               step (cpt - 1) predecessor
           else
-            loop acc sz step (cpt - 1) predecessor in
+            loop acc ~sz step (cpt - 1) predecessor in
-  Block.predecessor b >>= function
+  Store.Block.Contents.read_exn (store, b) >>= fun { header } ->
-  | None -> Lwt.return (State.Block.header b, [])
+  predecessor store b >>= function
+  | None -> Lwt.return (header, [])
   | Some p ->
-      loop [] sz 1 9 p >>= fun hist ->
+      loop [] ~sz 1 9 p >>= fun hist ->
-      Lwt.return (State.Block.header b, hist)
+      Lwt.return (header, hist)
 
-let estimated_length (_head, hist) =
+type validity =
-  let rec loop acc step cpt = function
+  | Unknown
-    | [] -> acc
+  | Known_valid
-    | _ :: hist ->
+  | Known_invalid
-        if cpt = 0 then
+
-          loop (acc+step) (step*2) 9 hist
+let unknown_prefix cond (head, hist) =
-        else
+  let rec loop hist acc =
-          loop (acc+step) step (cpt-1) hist
+    match hist with
+    | [] -> Lwt.return_none
+    | h :: t ->
+        cond h >>= function
+        | Known_valid ->
+            Lwt.return_some (h, (List.rev (h :: acc)))
+        | Known_invalid ->
+            Lwt.return_none
+        | Unknown ->
+            loop t (h :: acc)
   in
-  loop 1 1 9 hist
+  cond (Block_header.hash head) >>= function
-
+  | Known_valid ->
-let fold ~f acc (head, hist) =
+      Lwt.return_some (Block_header.hash head, (head, []))
-  let rec loop step cpt acc = function
+  | Known_invalid ->
-    | [] | [_] -> acc
+      Lwt.return_none
-    | block :: (pred :: rem as hist) ->
+  | Unknown ->
-        let step, cpt =
+      loop hist [] >>= function
-          if cpt = 0 then
-            2 * step, 9
-          else
-            step, cpt - 1 in
-        let acc = f acc ~block ~pred ~step ~strict_step:(rem <> []) in
-        loop step cpt acc hist
-  in
-  loop 1 10 acc (Block_header.hash head :: hist)
-
-type step = {
-  block: Block_hash.t ;
-  predecessor: Block_hash.t ;
-  step: int ;
-  strict_step: bool ;
-}
-
-let to_steps locator =
-  fold
-    ~f:begin fun acc ~block ~pred ~step ~strict_step -> {
-        block ; predecessor = pred ; step ; strict_step ;
-      } :: acc
-    end
-    [] locator
-
-let rec known_ancestor net_state acc hist =
-  match hist with
-  | [] -> Lwt.return_none
-  | h :: hist ->
-      Block.read_opt net_state h >>= function
-      | Some block -> Lwt.return (Some (block, List.rev (h :: acc)))
       | None ->
-          Block.known_invalid net_state h >>= function
+          Lwt.return_none
-          | true -> Lwt.return_none
+      | Some (tail, hist) ->
-          | false -> known_ancestor net_state (h :: acc) hist
+          Lwt.return_some (tail, (head, hist))
-
-let known_ancestor net_state (head, hist) =
-  let hash = Block_header.hash head in
-  if Block_hash.equal hash (State.Net.faked_genesis_hash net_state) then
-    State.Block.read_exn
-      net_state (State.Net.genesis net_state).block >>= fun genesis ->
-    Lwt.return_some (genesis, (head, []))
-  else
-    State.Block.read_opt net_state hash >>= function
-    | Some ancestor -> Lwt.return_some (ancestor, (head, []))
-    | None ->
-        known_ancestor net_state [] hist >>= function
-        | None -> Lwt.return_none
-        | Some (ancestor, prefix) ->
-            Lwt.return_some (ancestor, (head, prefix))
-
-let find_new net_state locator sz =
-  let rec path sz acc h =
-    if sz <= 0 then Lwt.return (List.rev acc)
-    else
-      read_chain_store net_state begin fun chain_store _data ->
-        Store.Chain.In_chain.read_opt (chain_store, h)
-      end >>= function
-      | None -> Lwt.return (List.rev acc)
-      | Some s -> path (sz-1) (s :: acc) s in
-  known_ancestor net_state locator >>= function
-  | None -> Lwt.return_nil
-  | Some (known, _) ->
-      Chain.head net_state >>= fun head ->
-      Chain_traversal.common_ancestor known head >>= fun ancestor ->
-      path sz [] (Block.hash ancestor)
-

lib_node_shell/block_locator.mli

@@ -7,61 +7,27 @@
 (*                                                                        *)
 (**************************************************************************)
 
-open State
+type t = private raw
-
-type t = private (Block_header.t * Block_hash.t list)
 (** A type for sparse block locator (/à la/ Bitcoin) *)
 
+and raw = Block_header.t * Block_hash.t list
+(** Non private version of Block_store_locator.t for coercions *)
+
+val raw: t -> raw
+
 val encoding: t Data_encoding.t
 
-type error += Invalid_locator of P2p.Peer_id.t * t
+val compute: Store.Block.store -> Block_hash.t -> int -> t Lwt.t
-
-val compute: Block.t -> int -> t Lwt.t
 (** [compute block max_length] compute the sparse block locator for
     the [block]. The locator contains at most [max_length] elements. *)
 
-val fold:
+type validity =
-  f:('a ->
+  | Unknown
-     block:Block_hash.t -> pred:Block_hash.t ->
+  | Known_valid
-     step:int -> strict_step:bool -> 'a) ->
+  | Known_invalid
-  'a -> t -> 'a
-(** [map f l] applies [f] to each block of the locator, the last one
-    excepted. The function also receives the expected predecessor
-    [pred] of the [block] after [step] steps, i.e. the next block in
-    the locator. When [strict_step] is [true], then [step] is the
-    exact number of predecessor to be followed before to found
-    [pred]. Otherwise, it is only an upper bound. *)
-
-type step = {
-  block: Block_hash.t ;
-  predecessor: Block_hash.t ;
-  step: int ;
-  strict_step: bool ;
-}
-(** A 'step' in a locator is a couple of consecutives hashes in the
-    locator, and the expected difference of level the two blocks (or
-    an upper bounds when [strict_step = false]). *)
-
-val to_steps: t -> step list
-(** Build all the 'steps' composing the locator, starting with the
-    oldest one (typically the predecessor of the first step will be
-    `genesis`). All steps contains [strict_step = true], except the
-    first one. *)
-
-val estimated_length: t -> int
-(** [estimated_length locator] estimate the length of the chain
-    represented by [locator]. *)
-
-val known_ancestor: State.Net.t -> t -> (Block.t * t) option Lwt.t
-(** [known_ancestor net_state locator] computes the first block of
-    [locator] that is known to be a valid block. It also computes the
-    'prefix' of [locator] with end at the first valid block.  The
-    function returns [None] when no block in the locator are known or
-    if the first known block is invalid. *)
-
-val find_new:
-  State.Net.t -> t -> int -> Block_hash.t list Lwt.t
-(** [find_new net locator max_length] returns the blocks from our
-    current branch that would be unknown to a peer that sends us the
-    [locator]. *)
 
+val unknown_prefix:
+  (Block_hash.t -> validity Lwt.t) -> t -> (Block_hash.t * t) option Lwt.t
+(** [unknown_prefix validity locator] keeps only the unknown part of
+    the locator up to the first valid block. If there is no known valid
+    block or there is a known invalid one, None is returned. *)

lib_node_shell/block_locator_iterator.ml

@@ -0,0 +1,93 @@
+(**************************************************************************)
+(*                                                                        *)
+(*    Copyright (c) 2014 - 2017.                                          *)
+(*    Dynamic Ledger Solutions, Inc. <contact@tezos.com>                  *)
+(*                                                                        *)
+(*    All rights reserved. No warranty, explicit or implicit, provided.   *)
+(*                                                                        *)
+(**************************************************************************)
+
+let estimated_length locator =
+  let (_head, hist) = Block_locator.raw locator in
+  let rec loop acc step cpt = function
+    | [] -> acc
+    | _ :: hist ->
+        if cpt = 0 then
+          loop (acc+step) (step*2) 9 hist
+        else
+          loop (acc+step) step (cpt-1) hist
+  in
+  loop 1 1 9 hist
+
+let fold ~f acc locator =
+  let (head, hist) = Block_locator.raw locator in
+  let rec loop step cpt acc = function
+    | [] | [_] -> acc
+    | block :: (pred :: rem as hist) ->
+        let step, cpt =
+          if cpt = 0 then
+            2 * step, 9
+          else
+            step, cpt - 1 in
+        let acc = f acc ~block ~pred ~step ~strict_step:(rem <> []) in
+        loop step cpt acc hist
+  in
+  loop 1 10 acc (Block_header.hash head :: hist)
+
+type step = {
+  block: Block_hash.t ;
+  predecessor: Block_hash.t ;
+  step: int ;
+  strict_step: bool ;
+}
+
+let to_steps locator =
+  fold
+    ~f:begin fun acc ~block ~pred ~step ~strict_step -> {
+        block ; predecessor = pred ; step ; strict_step ;
+      } :: acc
+    end
+    [] locator
+
+let block_validity net_state block : Block_locator.validity Lwt.t =
+  State.Block.known net_state block >>= function
+  | false ->
+      if Block_hash.equal block (State.Net.faked_genesis_hash net_state) then
+        Lwt.return Block_locator.Known_valid
+      else
+        Lwt.return Block_locator.Unknown
+  | true ->
+      State.Block.known_invalid net_state block >>= function
+      | true ->
+          Lwt.return Block_locator.Known_invalid
+      | false ->
+          Lwt.return Block_locator.Known_valid
+
+let known_ancestor net_state locator =
+  Block_locator.unknown_prefix (block_validity net_state) locator >>= function
+  | None -> Lwt.return_none
+  | Some (tail, locator) ->
+      if Block_hash.equal tail (State.Net.faked_genesis_hash net_state) then
+        State.Block.read_exn
+          net_state (State.Net.genesis net_state).block >>= fun genesis ->
+        Lwt.return_some (genesis, locator)
+      else
+        State.Block.read_exn net_state tail >>= fun block ->
+        Lwt.return_some (block, locator)
+
+let find_new net_state locator sz =
+  let rec path sz acc h =
+    if sz <= 0 then Lwt.return (List.rev acc)
+    else
+      State.read_chain_store net_state begin fun chain_store _data ->
+        Store.Chain.In_chain.read_opt (chain_store, h)
+      end >>= function
+      | None -> Lwt.return (List.rev acc)
+      | Some s -> path (sz-1) (s :: acc) s in
+  known_ancestor net_state locator >>= function
+  | None -> Lwt.return_nil
+  | Some (known, _) ->
+      Chain.head net_state >>= fun head ->
+      Chain_traversal.common_ancestor known head >>= fun ancestor ->
+      path sz [] (State.Block.hash ancestor)
+

lib_node_shell/block_locator_iterator.mli

@@ -0,0 +1,55 @@
+(**************************************************************************)
+(*                                                                        *)
+(*    Copyright (c) 2014 - 2017.                                          *)
+(*    Dynamic Ledger Solutions, Inc. <contact@tezos.com>                  *)
+(*                                                                        *)
+(*    All rights reserved. No warranty, explicit or implicit, provided.   *)
+(*                                                                        *)
+(**************************************************************************)
+
+val fold:
+  f:('a ->
+     block:Block_hash.t -> pred:Block_hash.t ->
+     step:int -> strict_step:bool -> 'a) ->
+  'a -> Block_locator.t -> 'a
+(** [map f l] applies [f] to each block of the locator, the last one
+    excepted. The function also receives the expected predecessor
+    [pred] of the [block] after [step] steps, i.e. the next block in
+    the locator. When [strict_step] is [true], then [step] is the
+    exact number of predecessor to be followed before to found
+    [pred]. Otherwise, it is only an upper bound. *)
+
+type step = {
+  block: Block_hash.t ;
+  predecessor: Block_hash.t ;
+  step: int ;
+  strict_step: bool ;
+}
+(** A 'step' in a locator is a couple of consecutives hashes in the
+    locator, and the expected difference of level the two blocks (or
+    an upper bounds when [strict_step = false]). *)
+
+val to_steps: Block_locator.t -> step list
+(** Build all the 'steps' composing the locator, starting with the
+    oldest one (typically the predecessor of the first step will be
+    `genesis`). All steps contains [strict_step = true], except the
+    first one. *)
+
+val estimated_length: Block_locator.t -> int
+(** [estimated_length locator] estimate the length of the chain
+    represented by [locator]. *)
+
+val known_ancestor:
+  State.Net.t -> Block_locator.t -> (State.Block.t * Block_locator.t) option Lwt.t
+(** [known_ancestor net_state locator] computes the first block of
+    [locator] that is known to be a valid block. It also computes the
+    'prefix' of [locator] with end at the first valid block.  The
+    function returns [None] when no block in the locator are known or
+    if the first known block is invalid. *)
+
+val find_new:
+  State.Net.t -> Block_locator.t -> int -> Block_hash.t list Lwt.t
+(** [find_new net locator max_length] returns the blocks from our
+    current branch that would be unknown to a peer that sends us the
+    [locator]. *)
+

lib_node_shell/bootstrap_pipeline.ml

@@ -9,6 +9,8 @@
 
 include Logging.Make(struct let name = "node.validator.bootstrap_pipeline" end)
 
+type error += Invalid_locator of P2p.Peer_id.t * Block_locator.t
+
 type t = {
   canceler: Lwt_canceler.t ;
   block_header_timeout: float ;
@@ -29,7 +31,7 @@ type t = {
   mutable errors: Error_monad.error list ;
 }
 
-let fetch_step pipeline (step : Block_locator.step)  =
+let fetch_step pipeline (step : Block_locator_iterator.step)  =
   lwt_log_info "fetching step %a -> %a (%d%s) from peer %a."
     Block_hash.pp_short step.block
     Block_hash.pp_short step.predecessor
@@ -41,14 +43,12 @@ let fetch_step pipeline (step : Block_locator.step)  =
     if cpt < 0 then
       lwt_log_info "invalid step from peer %a (too long)."
         P2p.Peer_id.pp_short pipeline.peer_id >>= fun () ->
-      fail (Block_locator.Invalid_locator
+      fail (Invalid_locator (pipeline.peer_id, pipeline.locator))
-              (pipeline.peer_id, pipeline.locator))
     else if Block_hash.equal hash step.predecessor then
       if step.strict_step && cpt <> 0 then
         lwt_log_info "invalid step from peer %a (too short)."
           P2p.Peer_id.pp_short pipeline.peer_id >>= fun () ->
-        fail (Block_locator.Invalid_locator
+        fail (Invalid_locator (pipeline.peer_id, pipeline.locator))
-                (pipeline.peer_id, pipeline.locator))
       else
         return acc
     else
@@ -78,7 +78,7 @@ let fetch_step pipeline (step : Block_locator.step)  =
 
 let headers_fetch_worker_loop pipeline =
   begin
-    let steps = Block_locator.to_steps pipeline.locator in
+    let steps = Block_locator_iterator.to_steps pipeline.locator in
     iter_s (fetch_step pipeline) steps >>=? fun () ->
     return ()
   end >>= function

lib_node_shell/bootstrap_pipeline.mli

@@ -9,6 +9,8 @@
 
 type t
 
+type error += Invalid_locator of P2p.Peer_id.t * Block_locator.t
+
 val create:
   ?notify_new_block: (State.Block.t -> unit) ->
   block_header_timeout:float ->

lib_node_shell/chain.ml

@@ -40,6 +40,7 @@ type data = State.chain_data = {
   current_mempool: mempool ;
   live_blocks: Block_hash.Set.t ;
   live_operations: Operation_hash.Set.t ;
+  locator: Block_locator.t Lwt.t lazy_t ;
 }
 
 let data net_state =
@@ -47,7 +48,12 @@ let data net_state =
     Lwt.return data
   end
 
-let locked_set_head chain_store data block =
+let locator net_state =
+  data net_state >>= begin fun data ->
+    Lazy.force data.locator
+  end
+
+let locked_set_head net_state chain_store data block =
   let rec pop_blocks ancestor block =
     let hash = Block.hash block in
     if Block_hash.equal hash ancestor then
@@ -83,11 +89,12 @@ let locked_set_head chain_store data block =
                current_mempool = State.empty_mempool ;
                live_blocks ;
                live_operations ;
+               locator = lazy (State.compute_locator net_state block) ;
              }
 
 let set_head net_state block =
   update_chain_store net_state begin fun chain_store data ->
-    locked_set_head chain_store data block >>= fun new_chain_data ->
+    locked_set_head net_state chain_store data block >>= fun new_chain_data ->
     Lwt.return (Some new_chain_data,
                 data.current_head)
   end
@@ -97,7 +104,7 @@ let test_and_set_head net_state ~old block =
     if not (Block.equal data.current_head old) then
       Lwt.return (None, false)
     else
-      locked_set_head chain_store data block >>= fun new_chain_data ->
+      locked_set_head net_state chain_store data block >>= fun new_chain_data ->
       Lwt.return (Some new_chain_data, true)
   end
 

lib_node_shell/chain.mli

@@ -17,6 +17,7 @@ val genesis: Net.t -> Block.t Lwt.t
 
 (** The current head of the network's blockchain. *)
 val head: Net.t -> Block.t Lwt.t
+val locator: Net.t -> Block_locator.t Lwt.t
 
 (** All the available chain data. *)
 type data = {
@@ -24,11 +25,13 @@ type data = {
   current_mempool: mempool ;
   live_blocks: Block_hash.Set.t ;
   live_operations: Operation_hash.Set.t ;
+  locator: Block_locator.t Lwt.t lazy_t ;
 }
 
 (** Reading atomically all the chain data. *)
 val data: Net.t -> data Lwt.t
 
+
 (** The current head and all the known (valid) alternate heads. *)
 val known_heads: Net.t -> Block.t list Lwt.t
 

lib_node_shell/distributed_db.ml

@@ -466,8 +466,7 @@ module P2p_reader = struct
           ignore
           @@ P2p.try_send global_db.p2p state.conn
           @@ Get_current_branch net_id ;
-        Chain.head net_db.net_state >>= fun head ->
+        Chain.locator net_db.net_state >>= fun locator ->
-        Block_locator.compute head 200 >>= fun locator ->
         ignore
         @@ P2p.try_send global_db.p2p state.conn
         @@ Current_branch (net_id, locator) ;
@@ -951,10 +950,8 @@ module Advertise = struct
     send net_db ?peer @@
     Current_head (net_id, State.Block.header head, mempool)
 
-  let current_branch net_db ?peer head =
+  let current_branch net_db ?peer locator =
     let net_id = State.Net.id net_db.net_state in
-    assert (Net_id.equal net_id (State.Block.net_id head)) ;
-    Block_locator.compute head 200 >>= fun locator ->
     send net_db ?peer @@ Current_branch (net_id, locator) ;
     Lwt.return_unit
 

lib_node_shell/distributed_db.mli

@@ -84,7 +84,7 @@ module Advertise : sig
       network, of a new head and its sparse history. *)
   val current_branch:
     net_db -> ?peer:P2p.Peer_id.t ->
-    State.Block.t -> unit Lwt.t
+    Block_locator.t -> unit Lwt.t
 
 end
 

lib_node_shell/net_validator.ml

@@ -114,7 +114,8 @@ let broadcast_head nv ~previous block =
       Distributed_db.Advertise.current_head nv.net_db block ;
       Lwt.return_unit
     end else begin
-      Distributed_db.Advertise.current_branch nv.net_db block
+      Chain.locator (Distributed_db.net_state nv.net_db) >>= fun locator ->
+      Distributed_db.Advertise.current_branch nv.net_db locator
     end
   end
 

lib_node_shell/peer_validator.ml

@@ -51,7 +51,7 @@ let set_bootstrapped pv =
   end
 
 let bootstrap_new_branch pv _ancestor _head unknown_prefix =
-  let len = Block_locator.estimated_length unknown_prefix in
+  let len = Block_locator_iterator.estimated_length unknown_prefix in
   lwt_log_info
     "validating new branch from peer %a (approx. %d blocks)"
     P2p.Peer_id.pp_short pv.peer_id len >>= fun () ->
@@ -154,7 +154,7 @@ let may_validate_new_branch pv distant_hash locator =
     return ()
   end else begin
     let net_state = Distributed_db.net_state pv.net_db in
-    Block_locator.known_ancestor net_state locator >>= function
+    Block_locator_iterator.known_ancestor net_state locator >>= function
     | None ->
         lwt_log_info
           "ignoring branch %a without common ancestor from peer: %a."
@@ -192,7 +192,7 @@ let rec worker_loop pv =
   | Ok () ->
       worker_loop pv
   | Error ((( Unknown_ancestor
-           | Block_locator.Invalid_locator _
+           | Bootstrap_pipeline.Invalid_locator _
            | Block_validator.Invalid_block _ ) :: _) as errors ) ->
       (* TODO ban the peer_id... *)
       lwt_log_info "Terminating the validation worker for peer %a (kickban)."

lib_node_shell/state.ml

@@ -90,6 +90,7 @@ and chain_data = {
   current_mempool: mempool ;
   live_blocks: Block_hash.Set.t ;
   live_operations: Operation_hash.Set.t ;
+  locator: Block_locator.t Lwt.t lazy_t ;
 }
 
 and mempool = {
@@ -137,6 +138,14 @@ let update_chain_store { net_id ; context_index ; chain_state } f =
     Lwt.return res
   end
 
+let compute_locator_from_hash (net : net_state) ?(size = 200) head =
+  Shared.use net.block_store begin fun block_store ->
+    Block_locator.compute block_store head size
+  end
+
+let compute_locator net ?size head =
+  compute_locator_from_hash net ?size head.hash
+
 type t = global_state
 
 module Locked_block = struct
@@ -198,6 +207,7 @@ module Net = struct
         current_mempool = empty_mempool ;
         live_blocks = Block_hash.Set.singleton genesis.block ;
         live_operations = Operation_hash.Set.empty ;
+        locator = lazy (compute_locator_from_hash net_state current_head) ;
       } ;
       chain_store ;
     }

lib_node_shell/state.mli

@@ -160,6 +160,8 @@ val read_block:
 val read_block_exn:
   global_state -> Block_hash.t -> Block.t Lwt.t
 
+val compute_locator: Net.t -> ?size:int -> Block.t -> Block_locator.t Lwt.t
+
 val fork_testnet:
   Block.t -> Protocol_hash.t -> Time.t -> Net.t tzresult Lwt.t
 
@@ -176,6 +178,7 @@ type chain_data = {
   current_mempool: mempool ;
   live_blocks: Block_hash.Set.t ;
   live_operations: Operation_hash.Set.t ;
+  locator: Block_locator.t Lwt.t lazy_t ;
 }
 
 val read_chain_store:

test/shell/test_state.ml

@@ -271,8 +271,8 @@ let test_ancestor s =
 
 let test_locator s =
   let check_locator h1 expected =
-    Block_locator.compute
+    State.compute_locator s.net
-      (vblock s h1) (List.length expected) >>= fun l ->
+      ~size:(List.length expected) (vblock s h1) >>= fun l ->
     let _, l = (l : Block_locator.t :> _ * _) in
     if List.length l <> List.length expected then
       Assert.fail_msg
@@ -415,8 +415,9 @@ let test_new_blocks s =
 
 let test_find_new s =
   let test s h expected =
-    Block_locator.compute (vblock s h) 50 >>= fun loc ->
+    State.compute_locator s.net ~size:50 (vblock s h) >>= fun loc ->
-    Block_locator.find_new s.net loc (List.length expected) >>= fun blocks ->
+    Block_locator_iterator.find_new
+      s.net loc (List.length expected) >>= fun blocks ->
     if List.length blocks <> List.length expected then
       Assert.fail_msg
         "Invalid find new length %s (found: %d, expected: %d)"