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::time::clock::Clock;
62use gwr_engine::traits::{Routable, Runnable, SimObject};
63use gwr_engine::types::{SimError, SimResult};
64use gwr_model_builder::{EntityDisplay, EntityGet};
65use gwr_track::entity::Entity;
66use gwr_track::tracker::aka::Aka;
67use gwr_track::{enter, exit, trace};
68
69use crate::take_option;
70
71/// Trait required for routing algorithms to implement.
72pub trait Route<T>
73where
74    T: Routable,
75{
76    /// Given an object, return the index of the egress port to map the object
77    /// to.
78    fn route(&self, object: &T) -> Result<usize, SimError>;
79}
80
81pub struct DefaultAlgorithm {}
82
83impl<T> Route<T> for DefaultAlgorithm
84where
85    T: Routable,
86{
87    /// Determine route by taking the object destination as an index.
88    fn route(&self, obj_to_route: &T) -> Result<usize, SimError> {
89        Ok(obj_to_route.destination() as usize)
90    }
91}
92
93#[derive(EntityGet, EntityDisplay)]
94pub struct Router<T>
95where
96    T: SimObject + Routable,
97{
98    entity: Rc<Entity>,
99    rx: RefCell<Option<InPort<T>>>,
100    tx: RefCell<Vec<OutPort<T>>>,
101    algorithm: Box<dyn Route<T>>,
102}
103
104impl<T> Router<T>
105where
106    T: SimObject + Routable,
107{
108    pub fn new_and_register_with_renames(
109        engine: &Engine,
110        clock: &Clock,
111        parent: &Rc<Entity>,
112        name: &str,
113        aka: Option<&Aka>,
114        num_egress: usize,
115        algorithm: Box<dyn Route<T>>,
116    ) -> Result<Rc<Self>, SimError> {
117        let entity = Rc::new(Entity::new(parent, name));
118        let rx = InPort::new_with_renames(engine, clock, &entity, "rx", aka);
119        let mut tx = Vec::with_capacity(num_egress);
120        for i in 0..num_egress {
121            tx.push(OutPort::new_with_renames(&entity, &format!("tx_{i}"), aka));
122        }
123        let rc_self = Rc::new(Self {
124            entity,
125            rx: RefCell::new(Some(rx)),
126            tx: RefCell::new(tx),
127            algorithm,
128        });
129        engine.register(rc_self.clone());
130        Ok(rc_self)
131    }
132
133    pub fn new_and_register(
134        engine: &Engine,
135        clock: &Clock,
136        parent: &Rc<Entity>,
137        name: &str,
138        num_egress: usize,
139        algorithm: Box<dyn Route<T>>,
140    ) -> Result<Rc<Self>, SimError> {
141        Self::new_and_register_with_renames(
142            engine, clock, parent, name, None, num_egress, algorithm,
143        )
144    }
145
146    pub fn connect_port_tx_i(&self, i: usize, port_state: PortStateResult<T>) -> SimResult {
147        match self.tx.borrow_mut().get_mut(i) {
148            None => {
149                sim_error!(format!("{self}: no tx port {i}"))
150            }
151            Some(tx) => tx.connect(port_state),
152        }
153    }
154
155    pub fn port_rx(&self) -> PortStateResult<T> {
156        self.rx.borrow().as_ref().unwrap().state()
157    }
158}
159
160#[async_trait(?Send)]
161impl<T> Runnable for Router<T>
162where
163    T: SimObject + Routable,
164{
165    async fn run(&self) -> SimResult {
166        let tx: Vec<OutPort<T>> = self.tx.borrow_mut().drain(..).collect();
167        let rx = take_option!(self.rx);
168        let algorithm = &self.algorithm;
169
170        loop {
171            let value = rx.get()?.await;
172            enter!(self.entity ; value.id());
173
174            let tx_index = algorithm.route(&value)?;
175            trace!(self.entity ; "Route {} to {}", value, tx_index);
176
177            match tx.get(tx_index) {
178                None => {
179                    return sim_error!(format!(
180                        "{self}: {value:?} selected invalid egress index {tx_index}"
181                    ));
182                }
183                Some(tx) => {
184                    exit!(self.entity ; value.id());
185                    tx.put(value)?.await;
186                }
187            }
188        }
189    }
190}