Source file stake_storage.ml

(*****************************************************************************)
(*                                                                           *)
(* Open Source License                                                       *)
(* Copyright (c) 2021 Nomadic Labs, <contact@nomadic-labs.com>               *)
(*                                                                           *)
(* Permission is hereby granted, free of charge, to any person obtaining a   *)
(* copy of this software and associated documentation files (the "Software"),*)
(* to deal in the Software without restriction, including without limitation *)
(* the rights to use, copy, modify, merge, publish, distribute, sublicense,  *)
(* and/or sell copies of the Software, and to permit persons to whom the     *)
(* Software is furnished to do so, subject to the following conditions:      *)
(*                                                                           *)
(* The above copyright notice and this permission notice shall be included   *)
(* in all copies or substantial portions of the Software.                    *)
(*                                                                           *)
(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*)
(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  *)
(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL   *)
(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*)
(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING   *)
(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER       *)
(* DEALINGS IN THE SOFTWARE.                                                 *)
(*                                                                           *)
(*****************************************************************************)

open Misc

module Selected_distribution_for_cycle = struct
  module Cache_client = struct
    type cached_value = (Signature.Public_key_hash.t * Tez_repr.t) list

    let namespace = Cache_repr.create_namespace "stake_distribution"

    let cache_index = 1

    let value_of_identifier ctxt identifier =
      let cycle = Cycle_repr.of_string_exn identifier in
      Storage.Stake.Selected_distribution_for_cycle.get ctxt cycle
  end

  module Cache = (val Cache_repr.register_exn (module Cache_client))

  let identifier_of_cycle cycle = Format.asprintf "%a" Cycle_repr.pp cycle

  let init ctxt cycle stakes =
    let id = identifier_of_cycle cycle in
    Storage.Stake.Selected_distribution_for_cycle.init ctxt cycle stakes
    >>=? fun ctxt ->
    let size = Constants_repr.stake_distribution_size in
    Cache.update ctxt id (Some (stakes, size)) >>?= fun ctxt -> return ctxt

  let get ctxt cycle =
    let id = identifier_of_cycle cycle in
    Cache.find ctxt id >>=? function
    | None -> Storage.Stake.Selected_distribution_for_cycle.get ctxt cycle
    | Some v -> return v

  let remove_existing ctxt cycle =
    let id = identifier_of_cycle cycle in
    (Cache.find ctxt id >>=? function
     | None -> return ctxt
     | Some _ -> Cache.update ctxt id None |> Lwt.return)
    >>=? fun ctxt ->
    Storage.Stake.Selected_distribution_for_cycle.remove_existing ctxt cycle
end

module Delegate_sampler_state = struct
  module Cache_client = struct
    type cached_value =
      (Signature.Public_key.t * Signature.Public_key_hash.t) Sampler.t

    let namespace = Cache_repr.create_namespace "sampler_state"

    let cache_index = 2

    let value_of_identifier ctxt identifier =
      let cycle = Cycle_repr.of_string_exn identifier in
      Storage.Delegate_sampler_state.get ctxt cycle
  end

  module Cache = (val Cache_repr.register_exn (module Cache_client))

  let identifier_of_cycle cycle = Format.asprintf "%a" Cycle_repr.pp cycle

  let init ctxt cycle sampler_state =
    let id = identifier_of_cycle cycle in
    Storage.Delegate_sampler_state.init ctxt cycle sampler_state
    >>=? fun ctxt ->
    let size = Constants_repr.sampler_state_size in
    Cache.update ctxt id (Some (sampler_state, size)) >>?= fun ctxt ->
    return ctxt

  let get ctxt cycle =
    let id = identifier_of_cycle cycle in
    Cache.find ctxt id >>=? function
    | None -> Storage.Delegate_sampler_state.get ctxt cycle
    | Some v -> return v

  let remove_existing ctxt cycle =
    let id = identifier_of_cycle cycle in
    (Cache.find ctxt id >>=? function
     | None -> return ctxt
     | Some _ -> Cache.update ctxt id None |> Lwt.return)
    >>=? fun ctxt -> Storage.Delegate_sampler_state.remove_existing ctxt cycle
end

let get_staking_balance = Storage.Stake.Staking_balance.get

let ensure_stake_inited ctxt delegate =
  Storage.Stake.Staking_balance.mem ctxt delegate >>= function
  | true -> return ctxt
  | false ->
      Frozen_deposits_storage.init ctxt delegate >>=? fun ctxt ->
      Storage.Stake.Staking_balance.init ctxt delegate Tez_repr.zero

let remove_stake ctxt delegate amount =
  ensure_stake_inited ctxt delegate >>=? fun ctxt ->
  let tokens_per_roll = Constants_storage.tokens_per_roll ctxt in
  get_staking_balance ctxt delegate >>=? fun staking_balance_before ->
  Tez_repr.(staking_balance_before -? amount) >>?= fun staking_balance ->
  Storage.Stake.Staking_balance.update ctxt delegate staking_balance
  >>=? fun ctxt ->
  Delegate_activation_storage.is_inactive ctxt delegate >>=? fun inactive ->
  if (not inactive) && Tez_repr.(staking_balance_before >= tokens_per_roll) then
    if Tez_repr.(staking_balance < tokens_per_roll) then
      Storage.Stake.Active_delegate_with_one_roll.remove ctxt delegate
      >>= fun ctxt -> return ctxt
    else return ctxt
  else
    (* The delegate was not in Stake.Active_delegate_with_one_roll,
       either because it was inactive, or because it did not have a
       roll, in which case it still does not have a roll. *)
    return ctxt

let add_stake ctxt delegate amount =
  ensure_stake_inited ctxt delegate >>=? fun ctxt ->
  let tokens_per_roll = Constants_storage.tokens_per_roll ctxt in
  get_staking_balance ctxt delegate >>=? fun staking_balance_before ->
  Tez_repr.(amount +? staking_balance_before) >>?= fun staking_balance ->
  Storage.Stake.Staking_balance.update ctxt delegate staking_balance
  >>=? fun ctxt ->
  if Tez_repr.(staking_balance >= tokens_per_roll) then
    Delegate_activation_storage.is_inactive ctxt delegate >>=? fun inactive ->
    if inactive || Tez_repr.(staking_balance_before >= tokens_per_roll) then
      return ctxt
    else
      Storage.Stake.Active_delegate_with_one_roll.add ctxt delegate ()
      >>= fun ctxt -> return ctxt
  else
    (* The delegate was not in Stake.Active_delegate_with_one_roll,
       because it did not have a roll (as otherwise it would have a
       roll now). *)
    return ctxt

let deactivate_only_call_from_delegate_storage ctxt delegate =
  Storage.Stake.Active_delegate_with_one_roll.remove ctxt delegate

let activate_only_call_from_delegate_storage ctxt delegate =
  ensure_stake_inited ctxt delegate >>=? fun ctxt ->
  get_staking_balance ctxt delegate >>=? fun staking_balance ->
  let tokens_per_roll = Constants_storage.tokens_per_roll ctxt in
  if Tez_repr.(staking_balance >= tokens_per_roll) then
    Storage.Stake.Active_delegate_with_one_roll.add ctxt delegate ()
    >>= fun ctxt -> return ctxt
  else return ctxt

let snapshot ctxt =
  Storage.Stake.Last_snapshot.get ctxt >>=? fun index ->
  Storage.Stake.Last_snapshot.update ctxt (index + 1) >>=? fun ctxt ->
  Storage.Stake.Staking_balance.snapshot ctxt index >>=? fun ctxt ->
  Storage.Stake.Active_delegate_with_one_roll.snapshot ctxt index

let select_distribution_for_cycle ctxt cycle pubkey =
  Storage.Stake.Last_snapshot.get ctxt >>=? fun max_index ->
  Storage.Seed.For_cycle.get ctxt cycle >>=? fun seed ->
  let rd = Seed_repr.initialize_new seed [Bytes.of_string "stake_snapshot"] in
  let seq = Seed_repr.sequence rd 0l in
  let selected_index =
    Seed_repr.take_int32 seq (Int32.of_int max_index) |> fst |> Int32.to_int
  in
  List.fold_left_es
    (fun ctxt index ->
      (if Compare.Int.(index = selected_index) then
       Storage.Stake.Active_delegate_with_one_roll.fold_snapshot
         ctxt
         index
         ~order:`Sorted
         ~init:([], Tez_repr.zero)
         ~f:(fun delegate () (acc, total_stake) ->
           Storage.Stake.Staking_balance.Snapshot.get ctxt (index, delegate)
           >>=? fun staking_balance ->
           let delegate_contract = Contract_repr.implicit_contract delegate in
           Storage.Contract.Frozen_deposits_limit.find ctxt delegate_contract
           >>=? fun frozen_deposits_limit ->
           Storage.Contract.Balance.get ctxt delegate_contract
           >>=? fun balance ->
           Frozen_deposits_storage.get ctxt delegate_contract
           >>=? fun frozen_deposits ->
           Tez_repr.(balance +? frozen_deposits.current_amount)
           >>?= fun total_balance ->
           let frozen_deposits_percentage =
             Constants_storage.frozen_deposits_percentage ctxt
           in
           let stake_to_consider =
             match frozen_deposits_limit with
             | Some frozen_deposits_limit -> (
                 try
                   let max_mutez = Tez_repr.of_mutez_exn Int64.max_int in
                   let frozen_stake_limit =
                     if Tez_repr.(frozen_deposits_limit > div_exn max_mutez 100)
                     then max_mutez
                     else
                       Tez_repr.(
                         div_exn
                           (mul_exn frozen_deposits_limit 100)
                           frozen_deposits_percentage)
                   in
                   Tez_repr.min staking_balance frozen_stake_limit
                 with _ -> staking_balance)
             | None -> staking_balance
           in
           let max_staking_capacity =
             Tez_repr.(
               div_exn (mul_exn total_balance 100) frozen_deposits_percentage)
           in
           let stake_for_cycle =
             Tez_repr.min stake_to_consider max_staking_capacity
           in
           Tez_repr.(total_stake +? stake_for_cycle) >>?= fun total_stake ->
           return ((delegate, stake_for_cycle) :: acc, total_stake))
       >>=? fun (stakes, total_stake) ->
       let stakes =
         List.sort (fun (_, x) (_, y) -> Tez_repr.compare y x) stakes
       in
       Selected_distribution_for_cycle.init ctxt cycle stakes >>=? fun ctxt ->
       Storage.Total_active_stake.add ctxt cycle total_stake >>= fun ctxt ->
       List.fold_left_es
         (fun acc (pkh, stake) ->
           pubkey ctxt pkh >>=? fun pk ->
           return (((pk, pkh), Tez_repr.to_mutez stake) :: acc))
         []
         stakes
       >>=? fun stakes_pk ->
       let state = Sampler.create stakes_pk in
       Delegate_sampler_state.init ctxt cycle state
      else return ctxt)
      >>=? fun ctxt ->
      Storage.Stake.Staking_balance.delete_snapshot ctxt index >>= fun ctxt ->
      Storage.Stake.Active_delegate_with_one_roll.delete_snapshot ctxt index
      >>= fun ctxt -> return ctxt)
    ctxt
    Misc.(0 --> (max_index - 1))
  >>=? fun ctxt -> Storage.Stake.Last_snapshot.update ctxt 0

let select_distribution_for_cycle_do_not_call_except_for_migration =
  select_distribution_for_cycle

let clear_cycle ctxt cycle =
  Storage.Total_active_stake.remove_existing ctxt cycle >>=? fun ctxt ->
  Selected_distribution_for_cycle.remove_existing ctxt cycle >>=? fun ctxt ->
  Delegate_sampler_state.remove_existing ctxt cycle >>=? fun ctxt ->
  Storage.Seed.For_cycle.remove_existing ctxt cycle

let init_first_cycles ctxt pubkey =
  let preserved = Constants_storage.preserved_cycles ctxt in
  List.fold_left_es
    (fun ctxt c ->
      let cycle = Cycle_repr.of_int32_exn (Int32.of_int c) in
      snapshot ctxt >>=? fun ctxt ->
      select_distribution_for_cycle ctxt cycle pubkey)
    ctxt
    (0 --> preserved)
  >>=? fun ctxt ->
  (* Precompute a snapshot for cycle (preserved_cycles + 1) *)
  snapshot ctxt

let fold ctxt ~f ~order init =
  Storage.Stake.Active_delegate_with_one_roll.fold
    ctxt
    ~order
    ~init:(Ok init)
    ~f:(fun delegate () acc ->
      acc >>?= fun acc ->
      get_staking_balance ctxt delegate >>=? fun stake ->
      f (delegate, stake) acc)

let select_new_distribution_at_cycle_end ctxt ~new_cycle =
  let preserved = Constants_storage.preserved_cycles ctxt in
  let for_cycle = Cycle_repr.add new_cycle preserved in
  select_distribution_for_cycle ctxt for_cycle

let clear_at_cycle_end ctxt ~new_cycle =
  let max_slashing_period = Constants_storage.max_slashing_period ctxt in
  match Cycle_repr.sub new_cycle max_slashing_period with
  | None -> return ctxt
  | Some cycle_to_clear -> clear_cycle ctxt cycle_to_clear

let get ctxt delegate =
  Storage.Stake.Active_delegate_with_one_roll.mem ctxt delegate >>= function
  | true -> get_staking_balance ctxt delegate
  | false -> return Tez_repr.zero

let fold_on_active_delegates_with_rolls =
  Storage.Stake.Active_delegate_with_one_roll.fold

let get_selected_distribution = Selected_distribution_for_cycle.get

let find_selected_distribution =
  Storage.Stake.Selected_distribution_for_cycle.find

let prepare_stake_distribution ctxt =
  let level = Level_storage.current ctxt in
  Selected_distribution_for_cycle.get ctxt level.cycle >>=? fun stakes ->
  let stake_distribution =
    List.fold_left
      (fun map (pkh, stake) -> Signature.Public_key_hash.Map.add pkh stake map)
      Signature.Public_key_hash.Map.empty
      stakes
  in
  return
    (Raw_context.init_stake_distribution_for_current_cycle
       ctxt
       stake_distribution)

let get_total_active_stake = Storage.Total_active_stake.get