use std::collections::HashMap;
use parking_lot::Mutex;
use super::topology::{InnovationId, NodeId};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct NodeSplit {
pub new_node: NodeId,
pub in_innov: InnovationId,
pub out_innov: InnovationId,
}
#[derive(Debug)]
pub struct InnovationRegistry {
inner: Mutex<RegistryInner>,
}
#[derive(Debug)]
struct RegistryInner {
next_innovation: InnovationId,
next_node: NodeId,
conn_cache: HashMap<(NodeId, NodeId), InnovationId>,
node_cache: HashMap<InnovationId, NodeSplit>,
}
impl InnovationRegistry {
#[must_use]
pub fn new(initial_node_count: usize, initial_innovation_count: usize) -> Self {
Self {
inner: Mutex::new(RegistryInner {
next_innovation: InnovationId::new(initial_innovation_count as u64),
next_node: NodeId::new(initial_node_count as u64),
conn_cache: HashMap::new(),
node_cache: HashMap::new(),
}),
}
}
#[must_use]
pub fn register_connection(&self, source: NodeId, target: NodeId) -> InnovationId {
let mut inner = self.inner.lock();
if let Some(&id) = inner.conn_cache.get(&(source, target)) {
return id;
}
let id = inner.next_innovation;
inner.next_innovation = id.succ();
inner.conn_cache.insert((source, target), id);
id
}
#[must_use]
pub fn register_node_split(&self, split: InnovationId) -> NodeSplit {
let mut inner = self.inner.lock();
if let Some(&existing) = inner.node_cache.get(&split) {
return existing;
}
let new_node = inner.next_node;
inner.next_node = new_node.succ();
let in_innov = inner.next_innovation;
let out_innov = in_innov.succ();
inner.next_innovation = out_innov.succ();
let result = NodeSplit {
new_node,
in_innov,
out_innov,
};
inner.node_cache.insert(split, result);
result
}
#[must_use]
pub fn next_innovation(&self) -> InnovationId {
self.inner.lock().next_innovation
}
#[must_use]
pub fn next_node_id(&self) -> NodeId {
self.inner.lock().next_node
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_registry_starts_after_seed() {
let registry = InnovationRegistry::new(3, 2);
assert_eq!(
registry.next_node_id().get(),
3,
"node ids start after I+O seed nodes"
);
assert_eq!(
registry.next_innovation().get(),
2,
"innovations start after the I*O seed connections"
);
}
#[test]
fn test_register_connection_caches_id() {
let registry = InnovationRegistry::new(3, 2);
let a = registry.register_connection(NodeId::new(0), NodeId::new(2));
let b = registry.register_connection(NodeId::new(1), NodeId::new(2));
assert_eq!(a.get(), 2);
assert_eq!(b.get(), 3);
assert_eq!(
registry.register_connection(NodeId::new(0), NodeId::new(2)),
a
);
assert_eq!(registry.next_innovation().get(), 4);
}
#[test]
fn test_register_node_split_caches_and_allocates_one_node_two_innovs() {
let registry = InnovationRegistry::new(3, 2);
let s = registry.register_node_split(InnovationId::new(0));
assert_eq!(
s.new_node.get(),
3,
"first hidden node id is after the seed nodes"
);
assert_eq!(s.in_innov.get(), 2);
assert_eq!(s.out_innov.get(), 3);
assert_eq!(registry.next_node_id().get(), 4);
assert_eq!(registry.next_innovation().get(), 4);
assert_eq!(registry.register_node_split(InnovationId::new(0)), s);
let t = registry.register_node_split(InnovationId::new(1));
assert_eq!(t.new_node.get(), 4);
assert_eq!(t.in_innov.get(), 4);
assert_eq!(t.out_innov.get(), 5);
}
#[test]
fn test_independent_registries_replay_identical_ids() {
fn run() -> (Vec<InnovationId>, Vec<NodeId>) {
let reg = InnovationRegistry::new(3, 2);
let mut innovs = Vec::new();
let mut nodes = Vec::new();
innovs.push(reg.register_connection(NodeId::new(0), NodeId::new(2)));
let s = reg.register_node_split(InnovationId::new(0));
nodes.push(s.new_node);
innovs.push(s.in_innov);
innovs.push(s.out_innov);
innovs.push(reg.register_connection(NodeId::new(1), s.new_node));
(innovs, nodes)
}
let (i1, n1) = run();
let (i2, n2) = run();
assert_eq!(i1, i2, "innovation sequence is reproducible across runs");
assert_eq!(n1, n2, "node id sequence is reproducible across runs");
}
#[test]
fn test_shared_registry_aligns_same_split_across_genomes() {
let registry = InnovationRegistry::new(3, 2);
let split_in_genome_a = registry.register_node_split(InnovationId::new(0));
let split_in_genome_b = registry.register_node_split(InnovationId::new(0));
assert_eq!(split_in_genome_a, split_in_genome_b);
}
}