1use std::fmt;
4
5pub type NodeId = usize;
7
8pub type PlayerId = usize;
10
11pub type Payoff = f64;
13
14#[derive(Clone, Debug, PartialEq)]
16pub enum NodeType {
17 Decision { player: PlayerId },
19 Chance { probabilities: Vec<f64> },
21 Terminal { payoffs: Vec<Payoff> },
23}
24
25#[derive(Clone, Debug)]
27pub struct NodeData {
28 pub id: NodeId,
30 pub node_type: NodeType,
32 pub label: String,
34 pub actions: Vec<String>,
36 pub children: Vec<NodeId>,
38 pub parent: Option<NodeId>,
40 pub incoming_action: Option<String>,
42}
43
44impl NodeData {
45 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 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 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 pub fn is_terminal(&self) -> bool {
86 matches!(self.node_type, NodeType::Terminal { .. })
87 }
88
89 pub fn is_decision(&self) -> bool {
91 matches!(self.node_type, NodeType::Decision { .. })
92 }
93
94 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 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 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}