Skip to main content

extensive_form/
node.rs

1//! Node types for game trees.
2
3use std::fmt;
4
5/// Unique node identifier.
6pub type NodeId = usize;
7
8/// Player identifier (0, 1, ...). Chance/terminal use special values.
9pub type PlayerId = usize;
10
11/// Payoff for a player at a terminal node.
12pub type Payoff = f64;
13
14/// The type of a game tree node.
15#[derive(Clone, Debug, PartialEq)]
16pub enum NodeType {
17    /// A decision node where a player chooses an action.
18    Decision { player: PlayerId },
19    /// A chance (nature) node with probability distribution over actions.
20    Chance { probabilities: Vec<f64> },
21    /// A terminal (leaf) node with payoffs for each player.
22    Terminal { payoffs: Vec<Payoff> },
23}
24
25/// Data associated with a node in the game tree.
26#[derive(Clone, Debug)]
27pub struct NodeData {
28    /// Unique identifier for this node.
29    pub id: NodeId,
30    /// The type of this node.
31    pub node_type: NodeType,
32    /// Label for this node (e.g., description of the game state).
33    pub label: String,
34    /// Actions available at this node (labels).
35    pub actions: Vec<String>,
36    /// Children node IDs, one per action.
37    pub children: Vec<NodeId>,
38    /// Parent node ID (None for root).
39    pub parent: Option<NodeId>,
40    /// The action label that led to this node from its parent.
41    pub incoming_action: Option<String>,
42}
43
44impl NodeData {
45    /// Create a new decision node.
46    pub fn decision(id: NodeId, player: PlayerId, label: &str) -> Self {
47        Self {
48            id,
49            node_type: NodeType::Decision { player },
50            label: label.to_string(),
51            actions: Vec::new(),
52            children: Vec::new(),
53            parent: None,
54            incoming_action: None,
55        }
56    }
57
58    /// Create a chance node.
59    pub fn chance(id: NodeId, probabilities: Vec<f64>, label: &str) -> Self {
60        Self {
61            id,
62            node_type: NodeType::Chance { probabilities },
63            label: label.to_string(),
64            actions: Vec::new(),
65            children: Vec::new(),
66            parent: None,
67            incoming_action: None,
68        }
69    }
70
71    /// Create a terminal node.
72    pub fn terminal(id: NodeId, payoffs: Vec<Payoff>, label: &str) -> Self {
73        Self {
74            id,
75            node_type: NodeType::Terminal { payoffs },
76            label: label.to_string(),
77            actions: Vec::new(),
78            children: Vec::new(),
79            parent: None,
80            incoming_action: None,
81        }
82    }
83
84    /// Is this a terminal node?
85    pub fn is_terminal(&self) -> bool {
86        matches!(self.node_type, NodeType::Terminal { .. })
87    }
88
89    /// Is this a decision node?
90    pub fn is_decision(&self) -> bool {
91        matches!(self.node_type, NodeType::Decision { .. })
92    }
93
94    /// Get the player at this decision node (panics if not a decision node).
95    pub fn player(&self) -> PlayerId {
96        match &self.node_type {
97            NodeType::Decision { player } => *player,
98            _ => panic!("Not a decision node"),
99        }
100    }
101
102    /// Get payoffs at this terminal node (panics if not terminal).
103    pub fn payoffs(&self) -> &[Payoff] {
104        match &self.node_type {
105            NodeType::Terminal { payoffs } => payoffs,
106            _ => panic!("Not a terminal node"),
107        }
108    }
109
110    /// Get the number of actions at this node.
111    pub fn num_actions(&self) -> usize {
112        self.actions.len()
113    }
114}
115
116impl fmt::Display for NodeData {
117    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118        match &self.node_type {
119            NodeType::Decision { player } => {
120                write!(f, "Node {} [P{}: {}]", self.id, player, self.label)
121            }
122            NodeType::Chance { .. } => {
123                write!(f, "Node {} [Chance: {}]", self.id, self.label)
124            }
125            NodeType::Terminal { payoffs } => {
126                write!(f, "Node {} [Terminal: {:?}]", self.id, payoffs)
127            }
128        }
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn test_decision_node() {
138        let n = NodeData::decision(0, 0, "Root");
139        assert!(n.is_decision());
140        assert!(!n.is_terminal());
141        assert_eq!(n.player(), 0);
142    }
143
144    #[test]
145    fn test_terminal_node() {
146        let n = NodeData::terminal(3, vec![2.0, 1.0], "End");
147        assert!(n.is_terminal());
148        assert_eq!(n.payoffs(), &[2.0, 1.0]);
149    }
150
151    #[test]
152    fn test_chance_node() {
153        let n = NodeData::chance(1, vec![0.5, 0.5], "Nature");
154        assert!(matches!(n.node_type, NodeType::Chance { .. }));
155    }
156
157    #[test]
158    fn test_node_display() {
159        let n = NodeData::decision(0, 1, "Choose");
160        let s = format!("{}", n);
161        assert!(s.contains("P1"));
162    }
163}