From b63b333b1af59f746d87e13af3ee714b47e80524 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gr=C3=A9goire=20Henry?= <gregoire.henry@tezos.com>
Date: Fri, 23 Feb 2018 10:29:12 -0500
Subject: [PATCH] Alpha: allow to iter on delegates

---
 src/proto_alpha/lib_protocol/src/alpha_context.mli    | 6 ++++++
 src/proto_alpha/lib_protocol/src/delegate_storage.ml  | 9 +++++++++
 src/proto_alpha/lib_protocol/src/delegate_storage.mli | 7 +++++++
 src/proto_alpha/lib_protocol/src/storage.ml           | 5 +++++
 src/proto_alpha/lib_protocol/src/storage.mli          | 4 ++++
 5 files changed, 31 insertions(+)

diff --git a/src/proto_alpha/lib_protocol/src/alpha_context.mli b/src/proto_alpha/lib_protocol/src/alpha_context.mli
index e954e341a..04efc2212 100644
--- a/src/proto_alpha/lib_protocol/src/alpha_context.mli
+++ b/src/proto_alpha/lib_protocol/src/alpha_context.mli
@@ -471,6 +471,12 @@ module Delegate : sig
   val set:
     context -> Contract.t -> public_key_hash option -> context tzresult Lwt.t
 
+  val fold:
+    context ->
+    init:'a -> f:(public_key_hash -> 'a -> 'a Lwt.t) -> 'a Lwt.t
+
+  val list: context -> public_key_hash list Lwt.t
+
 end
 
 module Vote : sig
diff --git a/src/proto_alpha/lib_protocol/src/delegate_storage.ml b/src/proto_alpha/lib_protocol/src/delegate_storage.ml
index 7a5691d87..bb1432726 100644
--- a/src/proto_alpha/lib_protocol/src/delegate_storage.ml
+++ b/src/proto_alpha/lib_protocol/src/delegate_storage.ml
@@ -111,8 +111,17 @@ let set c contract delegate =
         unlink c contract balance >>=? fun c ->
         Storage.Contract.Delegate.init_set c contract delegate >>= fun c ->
         link c contract delegate balance >>=? fun c ->
+        begin
+          if self_delegation then
+            Storage.Delegate.add c delegate
+          else
+            Lwt.return c
+        end >>= fun c ->
         return c
 
 let remove ctxt contract =
   Storage.Contract.Balance.get ctxt contract >>=? fun balance ->
   unlink ctxt contract balance
+
+let fold = Storage.Delegate.fold
+let list = Storage.Delegate.elements
diff --git a/src/proto_alpha/lib_protocol/src/delegate_storage.mli b/src/proto_alpha/lib_protocol/src/delegate_storage.mli
index 670611769..f3a68c425 100644
--- a/src/proto_alpha/lib_protocol/src/delegate_storage.mli
+++ b/src/proto_alpha/lib_protocol/src/delegate_storage.mli
@@ -26,3 +26,10 @@ val set:
   Raw_context.t tzresult Lwt.t
 
 val remove: Raw_context.t -> Contract_repr.t -> Raw_context.t tzresult Lwt.t
+
+val fold:
+  Raw_context.t ->
+  init:'a ->
+  f:(Ed25519.Public_key_hash.t -> 'a -> 'a Lwt.t) -> 'a Lwt.t
+
+val list: Raw_context.t -> Ed25519.Public_key_hash.t list Lwt.t
diff --git a/src/proto_alpha/lib_protocol/src/storage.ml b/src/proto_alpha/lib_protocol/src/storage.ml
index 7ea5d4a0a..6f1438158 100644
--- a/src/proto_alpha/lib_protocol/src/storage.ml
+++ b/src/proto_alpha/lib_protocol/src/storage.ml
@@ -129,6 +129,11 @@ module Contract = struct
 
 end
 
+module Delegate =
+  Make_data_set_storage
+    (Make_subcontext(Raw_context)(struct let name = ["delegates"] end))
+    (Ed25519.Public_key_hash)
+
 (** Rolls *)
 
 module Cycle = struct
diff --git a/src/proto_alpha/lib_protocol/src/storage.mli b/src/proto_alpha/lib_protocol/src/storage.mli
index 7f7d5ae62..53b21d862 100644
--- a/src/proto_alpha/lib_protocol/src/storage.mli
+++ b/src/proto_alpha/lib_protocol/src/storage.mli
@@ -154,6 +154,10 @@ module Contract : sig
 
 end
 
+module Delegate : Data_set_storage
+  with type t := Raw_context.t
+   and type elt = Ed25519.Public_key_hash.t
+
 (** Votes *)
 
 module Vote : sig