Alpha/Refactor: simplify code of `Raw_context.prepare_first_block`

This patch simplifies the protocol migration code.
It does not change the semantics.

Co-authored-by: Marco Stronati <marco@stronati.org>

src/proto_alpha/lib_protocol/src/raw_context.ml

@@ -460,40 +460,43 @@ let prepare ~level ~timestamp ~fitness ctxt =
     internal_nonces_used = Int_set.empty ;
   }
 
-type 'a previous_protocol =
-  | Genesis of 'a
+type previous_protocol =
+  | Genesis of Parameters_repr.t
   | Alpha_previous
 
-let check_first_block ctxt =
-  Context.get ctxt version_key >>= function
-  | None ->
-      failwith "Internal error: un-initialized context in check_first_block."
-  | Some bytes ->
-      let s = MBytes.to_string bytes in
-      if Compare.String.(s = version_value) then
-        failwith "Internal error: previously initialized context."
-      else if Compare.String.(s = "genesis") then
-        return (Genesis ())
-      else if Compare.String.(s = "alpha_previous") then
-        return Alpha_previous
-      else
-        storage_error (Incompatible_protocol_version s)
+let check_and_update_protocol_version ctxt =
+  begin
+    Context.get ctxt version_key >>= function
+    | None ->
+        failwith "Internal error: un-initialized context in check_first_block."
+    | Some bytes ->
+        let s = MBytes.to_string bytes in
+        if Compare.String.(s = version_value) then
+          failwith "Internal error: previously initialized context."
+        else if Compare.String.(s = "genesis") then
+          get_proto_param ctxt >>=? fun (param, ctxt) ->
+          return (Genesis param, ctxt)
+        else if Compare.String.(s = "alpha_previous") then
+          return (Alpha_previous, ctxt)
+        else
+          storage_error (Incompatible_protocol_version s)
+  end >>=? fun (previous_proto, ctxt) ->
+  Context.set ctxt version_key
+    (MBytes.of_string version_value) >>= fun ctxt ->
+  return (previous_proto, ctxt)
 
 let prepare_first_block ~level ~timestamp ~fitness ctxt =
-  check_first_block ctxt >>=? fun previous_protocol ->
+  check_and_update_protocol_version ctxt >>=? fun (previous_proto, ctxt) ->
   begin
-    match previous_protocol with
-    | Genesis () ->
+    match previous_proto with
+    | Genesis param ->
         Lwt.return (Raw_level_repr.of_int32 level) >>=? fun first_level ->
-        get_proto_param ctxt >>=? fun (param, ctxt) ->
         set_first_level ctxt first_level >>=? fun ctxt ->
         set_constants ctxt param.constants >>= fun ctxt ->
-        return (Genesis param, ctxt)
+        return ctxt
     | Alpha_previous ->
-        return (Alpha_previous, ctxt)
-  end >>=? fun (previous_proto, ctxt) ->
-  Context.set ctxt version_key
-    (MBytes.of_string version_value) >>= fun ctxt ->
+        return ctxt
+  end >>=? fun ctxt ->
   prepare ctxt ~level ~timestamp ~fitness >>=? fun ctxt ->
   return (previous_proto, ctxt)
 

src/proto_alpha/lib_protocol/src/raw_context.mli

@@ -58,15 +58,15 @@ val prepare:
   fitness: Fitness.t ->
   Context.t -> context tzresult Lwt.t
 
-type 'a previous_protocol =
-  | Genesis of 'a
+type previous_protocol =
+  | Genesis of Parameters_repr.t
   | Alpha_previous
 
 val prepare_first_block:
   level:int32 ->
   timestamp:Time.t ->
   fitness:Fitness.t ->
-  Context.t -> (Parameters_repr.t previous_protocol * context) tzresult Lwt.t
+  Context.t -> (previous_protocol * context) tzresult Lwt.t
 
 val activate: context -> Protocol_hash.t -> t Lwt.t
 val fork_test_chain: context -> Protocol_hash.t -> Time.t -> t Lwt.t