diff --git a/carnot/tree_overlay.py b/carnot/tree_overlay.py index 2d9eac0..1dc6af3 100644 --- a/carnot/tree_overlay.py +++ b/carnot/tree_overlay.py @@ -8,15 +8,20 @@ import random class CarnotTree: def __init__(self, nodes: List[Id], number_of_committees: int): self.number_of_committees = number_of_committees - self.committee_size, self.inner_committees, self.membership_committees = ( + # inner_commitees: list of tree nodes (int index) matching hashed external committee id + self.inner_committees: List[Id] + # membership committees: matching external (hashed) id to the set of members of a committee + self.membership_committees: Dict[Id, Committee] + self.inner_committees, self.membership_committees = ( CarnotTree.build_committee_from_nodes_with_size( nodes, self.number_of_committees ) ) - self.committees = {k: v for v, k in self.inner_committees.items()} - self.nodes = CarnotTree.build_nodes_index(nodes, self.committee_size) - self.committees_by_member = { - member: self.inner_committees[committee] + # committee match between tree nodes and external hashed ids + self.committees: Dict[Id, int] = {c: i for i, c in enumerate(self.inner_committees)} + # id (int index) of committee membership by member id + self.committees_by_member: Dict[Id, int] = { + member: committee for committee, v in self.membership_committees.items() for member in v } @@ -25,7 +30,7 @@ class CarnotTree: def build_committee_from_nodes_with_size( nodes: List[Id], number_of_committees: int, - ) -> Tuple[int, Dict[int, Id], Dict[int, Set[Id]]]: + ) -> Tuple[List[Id], Dict[int, Committee]]: committee_size, remainder = divmod(len(nodes), number_of_committees) committees = [ set(nodes[n*committee_size:(n+1)*committee_size]) @@ -40,16 +45,10 @@ class CarnotTree: committees = [frozenset(s) for s in committees] # TODO: This hash method should be specific to what we would want to use for the protocol hashes = [hash(s) for s in committees] - return committee_size, dict(enumerate(hashes)), dict(enumerate(committees)) - - @staticmethod - def build_nodes_index(nodes: List[Id], committee_size: int) -> Dict[Id, int]: - return { - _id: i // committee_size for i, _id in enumerate(nodes) - } + return hashes, dict(enumerate(committees)) def parent_committee(self, committee_id: Id) -> Optional[Id]: - return self.inner_committees.get(self.committees[committee_id] // 2 - 1) + return self.inner_committees[max(self.committees[committee_id] // 2 - 1, 0)] def child_committees(self, committee_id: Id) -> Tuple[Optional[Id], Optional[Id]]: base = self.committees[committee_id] * 2 @@ -67,11 +66,18 @@ class CarnotTree: def root_committee(self) -> Committee: return self.membership_committees[0] - def committee_by_committee_id(self, committee_id: Id) -> Optional[Committee]: - return self.membership_committees.get(self.inner_committees[committee_id]) + def committee_by_committee_idx(self, committee_id: int) -> Optional[Committee]: + return self.membership_committees.get(committee_id) - def committee_by_member_id(self, member_id: Id) -> Id: - return self.committees_by_member[member_id] + def committee_idx_by_member_id(self, member_id: Id) -> Optional[int]: + return self.committees_by_member.get(member_id) + + def committee_id_by_member_id(self, member_id: Id) -> Id: + return self.inner_committees[self.committees_by_member.get(member_id)] + + def committee_by_member_id(self, member_id: Id) -> Optional[Committee]: + if (committee_idx := self.committee_idx_by_member_id(member_id)) is not None: + return self.committee_by_committee_idx(committee_idx) class CarnotOverlay(EntropyOverlay): @@ -104,19 +110,15 @@ class CarnotOverlay(EntropyOverlay): return _id in self.carnot_tree.root_committee() def is_member_of_child_committee(self, parent: Id, child: Id) -> bool: - l, r = self.carnot_tree.child_committees(parent) - l = self.carnot_tree.committee_by_committee_id(l) if l is not None else set() or set() - r = self.carnot_tree.committee_by_committee_id(r) if r is not None else set() or set() - return child in l.join(r) + child_parent = self.parent_committee(child) + parent = self.carnot_tree.committee_by_member_id(parent) + return child_parent is parent def parent_committee(self, _id: Id) -> Optional[Committee]: if (parent_id := self.carnot_tree.parent_committee( - self.carnot_tree.committee_by_member_id(_id) - )) is None: - return None - return self.carnot_tree.committee_by_committee_id( - parent_id - ) + self.carnot_tree.committee_id_by_member_id(_id) + )) is not None: + return self.carnot_tree.committee_by_committee_idx(self.carnot_tree.committees[parent_id]) def leaf_committees(self) -> Set[Committee]: return set(self.carnot_tree.leaf_committees().values()) @@ -125,13 +127,15 @@ class CarnotOverlay(EntropyOverlay): return self.carnot_tree.root_committee() def is_child_of_root_committee(self, _id: Id) -> bool: - return _id in self.root_committee() + return self.parent_committee(_id) is self.root_committee() def leader_super_majority_threshold(self, _id: Id) -> int: - return (self.carnot_tree.committee_size * 2 // 3) + 1 + committee_size = len(self.carnot_tree.committee_by_member_id(_id)) + return (committee_size * 2 // 3) + 1 def super_majority_threshold(self, _id: Id) -> int: if self.is_member_of_leaf_committee(_id): return 0 - return (self.carnot_tree.committee_size * 2 // 3) + 1 + committee_size = len(self.carnot_tree.committee_by_member_id(_id)) + return (committee_size * 2 // 3) + 1