1use std::cell::RefCell;
55use std::rc::Rc;
56
57use async_trait::async_trait;
58use gwr_engine::engine::Engine;
59use gwr_engine::port::{InPort, OutPort, PortStateResult};
60use gwr_engine::sim_error;
61use gwr_engine::traits::{Routable, Runnable, SimObject};
62use gwr_engine::types::{SimError, SimResult};
63use gwr_model_builder::EntityDisplay;
64use gwr_track::entity::Entity;
65use gwr_track::{enter, exit, trace};
66
67use crate::take_option;
68
69pub trait Route<T>
71where
72 T: Routable,
73{
74 fn route(&self, object: &T) -> Result<usize, SimError>;
77}
78
79pub struct DefaultAlgorithm {}
80
81impl<T> Route<T> for DefaultAlgorithm
82where
83 T: Routable,
84{
85 fn route(&self, obj_to_route: &T) -> Result<usize, SimError> {
87 Ok(obj_to_route.destination() as usize)
88 }
89}
90
91#[derive(EntityDisplay)]
92pub struct Router<T>
93where
94 T: SimObject + Routable,
95{
96 pub entity: Rc<Entity>,
97 rx: RefCell<Option<InPort<T>>>,
98 tx: RefCell<Vec<OutPort<T>>>,
99 algorithm: Box<dyn Route<T>>,
100}
101
102impl<T> Router<T>
103where
104 T: SimObject + Routable,
105{
106 pub fn new_and_register(
107 engine: &Engine,
108 parent: &Rc<Entity>,
109 name: &str,
110 num_egress: usize,
111 algorithm: Box<dyn Route<T>>,
112 ) -> Result<Rc<Self>, SimError> {
113 let entity = Rc::new(Entity::new(parent, name));
114 let rx = InPort::new(&entity, "rx");
115 let mut tx = Vec::with_capacity(num_egress);
116 for i in 0..num_egress {
117 tx.push(OutPort::new(&entity, format!("tx{i}").as_str()));
118 }
119 let rc_self = Rc::new(Self {
120 entity,
121 rx: RefCell::new(Some(rx)),
122 tx: RefCell::new(tx),
123 algorithm,
124 });
125 engine.register(rc_self.clone());
126 Ok(rc_self)
127 }
128
129 pub fn connect_port_tx_i(&self, i: usize, port_state: PortStateResult<T>) -> SimResult {
130 match self.tx.borrow_mut().get_mut(i) {
131 None => {
132 sim_error!(format!("{self}: no tx port {i}"))
133 }
134 Some(tx) => tx.connect(port_state),
135 }
136 }
137
138 pub fn port_rx(&self) -> PortStateResult<T> {
139 self.rx.borrow().as_ref().unwrap().state()
140 }
141}
142
143#[async_trait(?Send)]
144impl<T> Runnable for Router<T>
145where
146 T: SimObject + Routable,
147{
148 async fn run(&self) -> SimResult {
149 let tx: Vec<OutPort<T>> = self.tx.borrow_mut().drain(..).collect();
150 let rx = take_option!(self.rx);
151 let algorithm = &self.algorithm;
152
153 loop {
154 let value = rx.get()?.await;
155 enter!(self.entity ; value.id());
156
157 let tx_index = algorithm.route(&value)?;
158 trace!(self.entity ; "Route {} to {}", value, tx_index);
159
160 match tx.get(tx_index) {
161 None => {
162 return sim_error!(format!(
163 "{self}: {value:?} selected invalid egress index {tx_index}"
164 ));
165 }
166 Some(tx) => {
167 exit!(self.entity ; value.id());
168 tx.put(value)?.await;
169 }
170 }
171 }
172 }
173}