linked_markov/
step.rs

1use rand::prelude::*;
2use std::{
3    collections::HashMap,
4    error::Error,
5    fmt::{Debug, Formatter, Result as FmtResult},
6    hash::{Hash, Hasher},
7    sync::{Arc, RwLock},
8};
9
10/// A reference-counted pointer to a `Step` in the Markov chain.
11pub type ToStep<T> = Arc<Step<T>>;
12
13/// A node in the Markov chain, holding a state and weighted transitions to other steps.
14#[derive(Default)]
15pub struct Step<T: Eq + Copy + Hash + Debug + Send + Sync> {
16    /// The state value for this step.
17    pub state: T,
18    /// Outgoing transitions and their weights.
19    pub transitions: RwLock<HashMap<ToStep<T>, usize>>,
20}
21
22impl<T> Clone for Step<T>
23where
24    T: Eq + Copy + Hash + Debug + Send + Sync,
25{
26    fn clone(&self) -> Self {
27        #[allow(clippy::mutable_key_type)]
28        let transitions = self.transitions.read().unwrap().clone();
29        Step {
30            state: self.state,
31            transitions: RwLock::new(transitions),
32        }
33    }
34}
35
36impl<T> Debug for Step<T>
37where
38    T: Eq + Copy + Hash + Debug + Send + Sync,
39{
40    fn fmt(&self, f: &mut Formatter) -> FmtResult {
41        write!(f, "Step {{ state: {:?} }}", self.state)
42    }
43}
44
45impl<T> Hash for Step<T>
46where
47    T: Eq + Copy + Hash + Debug + Send + Sync,
48{
49    fn hash<H: Hasher>(&self, hasher: &mut H) {
50        self.state.hash(hasher);
51    }
52}
53
54impl<T> PartialEq for Step<T>
55where
56    T: Eq + Copy + Hash + Debug + Send + Sync,
57{
58    fn eq(&self, other: &Self) -> bool {
59        self.state == other.state
60    }
61}
62
63impl<T> Eq for Step<T> where T: Eq + Copy + Hash + Debug + Send + Sync {}
64
65impl<T> Step<T>
66where
67    T: Eq + Copy + Hash + Debug + Send + Sync,
68{
69    /// Create a new `Step` with the given state.
70    pub fn new(state: T) -> Self {
71        Step {
72            state,
73            transitions: RwLock::new(HashMap::new()),
74        }
75    }
76
77    /// Add or update a transition to another step with a given weight.
78    pub fn insert_transition(&self, to_step: ToStep<T>, weight: usize) {
79        self.transitions.write().unwrap().insert(to_step, weight);
80    }
81
82    /// Randomly select the next step based on transition weights.
83    pub fn next(&self) -> Option<ToStep<T>> {
84        let mut rng = rand::rng();
85        let transitions = self.transitions.read().unwrap();
86        if transitions.is_empty() {
87            return None;
88        }
89        let total: usize = transitions.values().sum();
90        if total == 0 {
91            return None;
92        }
93        let roll = rng.random_range(0..total);
94        let mut cumulative = 0;
95        transitions.iter().find_map(|(to_step, &weight)| {
96            cumulative += weight;
97            if roll < cumulative {
98                Some(Arc::clone(to_step))
99            } else {
100                None
101            }
102        })
103    }
104}
105
106/// Walk the Markov chain for a fixed number of steps, returning the visited states.
107pub fn walk<T>(start: ToStep<T>, steps: usize) -> Vec<T>
108where
109    T: Eq + Copy + Hash + Debug + Send + Sync,
110{
111    let mut current = start;
112    let mut path = vec![current.state];
113    for _ in 1..steps {
114        if let Some(next) = current.next() {
115            path.push(next.state);
116            current = next;
117        } else {
118            break;
119        }
120    }
121    path
122}
123
124/// Walk the Markov chain for a fixed number of steps, applying a function to each transition.
125///
126/// The `apply` function is called with the current and next step, and can mutate the chain or collect data.
127/// # Examples:
128/// ```
129/// use linked_markov::{Step, ToStep, mut_walk};
130/// use std::sync::Arc;
131///
132/// let step_false: ToStep<bool> = Arc::new(Step::new(false));
133/// let step_true: ToStep<bool> = Arc::new(Step::new(true));
134/// step_false.insert_transition(step_true.clone(), 3);
135/// step_false.insert_transition(step_false.clone(), 1);
136/// step_true.insert_transition(step_false.clone(), 3);
137/// step_true.insert_transition(step_true.clone(), 1);
138/// let path = mut_walk(step_false, 100, |current, next| {
139///     current
140///         .transitions
141///         .write()
142///         .unwrap()
143///         .entry(next)
144///         .and_modify(|e| *e += 1)
145///         .or_insert(1);
146///     Ok(())
147/// })
148/// .unwrap();
149/// ```
150pub fn mut_walk<T, F>(start: ToStep<T>, steps: usize, apply: F) -> Result<Vec<T>, Box<dyn Error>>
151where
152    T: Eq + Copy + Hash + Debug + Send + Sync,
153    F: Fn(ToStep<T>, ToStep<T>) -> Result<(), Box<dyn Error>>,
154{
155    let mut current = start;
156    let mut path = vec![current.state];
157    for _ in 1..steps {
158        if let Some(next) = current.next() {
159            apply(current.clone(), next.clone())?;
160            path.push(current.state);
161            current = next;
162        } else {
163            break;
164        }
165    }
166    Ok(path)
167}
linked_markov/step.rs

linked_markov/
step.rs
