use elara_core::NodeId;
use std::collections::{HashMap, HashSet};
#[derive(Debug, Clone, Copy)]
pub struct PropagationEdge {
pub from: NodeId,
pub to: NodeId,
pub latency_ms: u32,
pub bandwidth: u8,
pub active: bool,
}
impl PropagationEdge {
pub fn new(from: NodeId, to: NodeId) -> Self {
Self {
from,
to,
latency_ms: 50,
bandwidth: 100,
active: true,
}
}
pub fn with_latency(mut self, latency_ms: u32) -> Self {
self.latency_ms = latency_ms;
self
}
pub fn with_bandwidth(mut self, bandwidth: u8) -> Self {
self.bandwidth = bandwidth;
self
}
}
#[derive(Debug, Clone, Default)]
pub struct PropagationTopology {
nodes: HashSet<NodeId>,
edges: HashMap<NodeId, Vec<PropagationEdge>>,
reverse_edges: HashMap<NodeId, HashSet<NodeId>>,
}
impl PropagationTopology {
pub fn new() -> Self {
Self::default()
}
pub fn add_node(&mut self, node: NodeId) {
self.nodes.insert(node);
}
pub fn remove_node(&mut self, node: NodeId) {
self.nodes.remove(&node);
self.edges.remove(&node);
for edges in self.edges.values_mut() {
edges.retain(|e| e.to != node);
}
self.reverse_edges.remove(&node);
for sources in self.reverse_edges.values_mut() {
sources.remove(&node);
}
}
pub fn add_edge(&mut self, edge: PropagationEdge) {
self.nodes.insert(edge.from);
self.nodes.insert(edge.to);
self.edges.entry(edge.from).or_default().push(edge);
self.reverse_edges
.entry(edge.to)
.or_default()
.insert(edge.from);
}
pub fn edges_from(&self, node: NodeId) -> &[PropagationEdge] {
self.edges.get(&node).map(|v| v.as_slice()).unwrap_or(&[])
}
pub fn downstream(&self, node: NodeId) -> Vec<NodeId> {
self.edges_from(node).iter().map(|e| e.to).collect()
}
pub fn upstream(&self, node: NodeId) -> Vec<NodeId> {
self.reverse_edges
.get(&node)
.map(|s| s.iter().copied().collect())
.unwrap_or_default()
}
pub fn node_count(&self) -> usize {
self.nodes.len()
}
pub fn has_node(&self, node: NodeId) -> bool {
self.nodes.contains(&node)
}
}
#[derive(Debug, Clone)]
pub struct StarTopology {
pub center: NodeId,
pub leaves: HashSet<NodeId>,
pub topology: PropagationTopology,
}
impl StarTopology {
pub fn new(center: NodeId) -> Self {
let mut topology = PropagationTopology::new();
topology.add_node(center);
Self {
center,
leaves: HashSet::new(),
topology,
}
}
pub fn add_leaf(&mut self, leaf: NodeId) {
self.leaves.insert(leaf);
self.topology
.add_edge(PropagationEdge::new(self.center, leaf));
}
pub fn remove_leaf(&mut self, leaf: NodeId) {
self.leaves.remove(&leaf);
self.topology.remove_node(leaf);
}
pub fn leaf_count(&self) -> usize {
self.leaves.len()
}
}
#[derive(Debug, Clone)]
pub struct TreeTopology {
pub root: NodeId,
parents: HashMap<NodeId, NodeId>,
children: HashMap<NodeId, HashSet<NodeId>>,
pub max_fanout: usize,
pub topology: PropagationTopology,
}
impl TreeTopology {
pub fn new(root: NodeId, max_fanout: usize) -> Self {
let mut topology = PropagationTopology::new();
topology.add_node(root);
Self {
root,
parents: HashMap::new(),
children: HashMap::new(),
max_fanout,
topology,
}
}
pub fn add_node(&mut self, node: NodeId) -> NodeId {
let parent = self.find_parent();
self.parents.insert(node, parent);
self.children.entry(parent).or_default().insert(node);
self.topology.add_edge(PropagationEdge::new(parent, node));
parent
}
fn find_parent(&self) -> NodeId {
let root_children = self.children.get(&self.root).map(|c| c.len()).unwrap_or(0);
if root_children < self.max_fanout {
return self.root;
}
let mut queue: Vec<NodeId> = self
.children
.get(&self.root)
.map(|c| c.iter().copied().collect())
.unwrap_or_default();
while let Some(node) = queue.pop() {
let child_count = self.children.get(&node).map(|c| c.len()).unwrap_or(0);
if child_count < self.max_fanout {
return node;
}
if let Some(children) = self.children.get(&node) {
queue.extend(children.iter().copied());
}
}
self.root
}
pub fn remove_node(&mut self, node: NodeId) {
if node == self.root {
return; }
let parent = self.parents.remove(&node);
let children = self.children.remove(&node).unwrap_or_default();
if let Some(p) = parent {
if let Some(siblings) = self.children.get_mut(&p) {
siblings.remove(&node);
}
}
for child in children {
if let Some(p) = parent {
self.parents.insert(child, p);
self.children.entry(p).or_default().insert(child);
self.topology.add_edge(PropagationEdge::new(p, child));
}
}
self.topology.remove_node(node);
}
pub fn depth(&self, node: NodeId) -> usize {
let mut depth = 0;
let mut current = node;
while let Some(&parent) = self.parents.get(¤t) {
depth += 1;
current = parent;
}
depth
}
pub fn node_count(&self) -> usize {
self.topology.node_count()
}
}
#[derive(Debug, Clone)]
pub struct MeshTopology {
nodes: HashSet<NodeId>,
pub topology: PropagationTopology,
}
impl MeshTopology {
pub fn new() -> Self {
Self {
nodes: HashSet::new(),
topology: PropagationTopology::new(),
}
}
pub fn add_node(&mut self, node: NodeId) {
for &existing in &self.nodes {
self.topology.add_edge(PropagationEdge::new(existing, node));
self.topology.add_edge(PropagationEdge::new(node, existing));
}
self.nodes.insert(node);
self.topology.add_node(node);
}
pub fn remove_node(&mut self, node: NodeId) {
self.nodes.remove(&node);
self.topology.remove_node(node);
}
pub fn node_count(&self) -> usize {
self.nodes.len()
}
}
impl Default for MeshTopology {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_star_topology() {
let broadcaster = NodeId::new(1);
let mut star = StarTopology::new(broadcaster);
star.add_leaf(NodeId::new(2));
star.add_leaf(NodeId::new(3));
star.add_leaf(NodeId::new(4));
assert_eq!(star.leaf_count(), 3);
assert_eq!(star.topology.downstream(broadcaster).len(), 3);
}
#[test]
fn test_tree_topology() {
let root = NodeId::new(1);
let mut tree = TreeTopology::new(root, 2);
for i in 2..=7 {
tree.add_node(NodeId::new(i));
}
assert_eq!(tree.node_count(), 7);
assert_eq!(tree.topology.downstream(root).len(), 2);
}
#[test]
fn test_mesh_topology() {
let mut mesh = MeshTopology::new();
mesh.add_node(NodeId::new(1));
mesh.add_node(NodeId::new(2));
mesh.add_node(NodeId::new(3));
assert_eq!(mesh.node_count(), 3);
assert_eq!(mesh.topology.downstream(NodeId::new(1)).len(), 2);
}
}