gwr_components/
router.rs

1// Copyright (c) 2023 Graphcore Ltd. All rights reserved.
2
3//! Perform routing between an input interface and a number number of outputs.
4//!
5//! The [Router] is passed an algorithm that makes the decision about which
6//! egress port to send each routed object to.
7//!
8//! # Ports
9//!
10//! This component has the following ports:
11//!  - One [input port](gwr_engine::port::InPort): `rx`
12//!  - N [output ports](gwr_engine::port::OutPort): `tx[i]` for `i in [0, N-1]`
13
14//! # Function
15//!
16//! The [Router] will take objects from the single input and send them to the
17//! correct output. A simplified summary of its functionality is:
18//!
19//! ```rust
20//! # use std::rc::Rc;
21//! # use async_trait::async_trait;
22//! # use gwr_components::router::Route;
23//! # use gwr_engine::port::{InPort, OutPort};
24//! # use gwr_engine::sim_error;
25//! # use gwr_engine::time::clock::Clock;
26//! # use gwr_engine::traits::{Routable, SimObject};
27//! # use gwr_engine::types::SimResult;
28//! # use gwr_track::entity::Entity;
29//! #
30//! # async fn run<T>(
31//! #     tx: Vec<OutPort<T>>,
32//! #     rx: InPort<T>,
33//! #     routing_algorithm: Box<dyn Route<T>>
34//! # ) -> SimResult
35//! # where
36//! #     T: SimObject + Routable
37//! # {
38//! loop {
39//!     let value = rx.get()?.await;
40//!     let tx_index = routing_algorithm.route(&value)?;
41//!
42//!     match tx.get(tx_index) {
43//!         None => {
44//!             // Report error
45//!         }
46//!         Some(tx) => {
47//!             tx.put(value)?.await;
48//!         }
49//!     }
50//! }
51//! # }
52//! ```
53
54use 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
69/// Trait required for routing algorithms to implement.
70pub trait Route<T>
71where
72    T: Routable,
73{
74    /// Given an object, return the index of the egress port to map the object
75    /// to.
76    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    /// Determine route by taking the object destination as an index.
86    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}