npc_engine_core/
node.rs

1/*
2 *  SPDX-License-Identifier: Apache-2.0 OR MIT
3 *  © 2020-2022 ETH Zurich and other contributors, see AUTHORS.txt for details
4 */
5
6use std::{
7    collections::{BTreeMap, BTreeSet},
8    fmt,
9    hash::{Hash, Hasher},
10    mem,
11    sync::{Arc, Weak},
12};
13
14use crate::{
15    active_task::{ActiveTask, ActiveTasks},
16    get_task_for_agent, AgentId, AgentValue, Domain, StateDiffRef, Task,
17};
18
19/// Strong atomic reference counted node.
20pub type Node<D> = Arc<NodeInner<D>>;
21
22/// Weak atomic reference counted node.
23pub type WeakNode<D> = Weak<NodeInner<D>>;
24
25/// The data associated to a node that form its key.
26pub struct NodeInner<D: Domain> {
27    pub(crate) diff: D::Diff,
28    pub(crate) active_agent: AgentId,
29    pub(crate) tick: u64,
30    pub(crate) tasks: ActiveTasks<D>,
31    current_values: BTreeMap<AgentId, AgentValue>, // pre-computed current values
32}
33
34impl<D: Domain> fmt::Debug for NodeInner<D> {
35    fn fmt(&self, f: &'_ mut fmt::Formatter) -> fmt::Result {
36        f.debug_struct("NodeInner")
37            .field("diff", &self.diff)
38            .field("agent", &self.active_agent)
39            .field("tick", &self.tick)
40            .field("tasks", &self.tasks)
41            .field("current_values", &self.current_values)
42            .finish()
43    }
44}
45
46impl<D: Domain> NodeInner<D> {
47    /// Create a new node, check for visible agents, and re-assign current tasks to the matching ones.
48    /// Return None if no active agent is not visible, and Some(node) otherwise.
49    pub fn new(
50        initial_state: &D::State,
51        start_tick: u64,
52        diff: D::Diff,
53        active_agent: AgentId,
54        tick: u64,
55        tasks: BTreeSet<ActiveTask<D>>,
56    ) -> Self {
57        let state_diff = StateDiffRef::new(initial_state, &diff);
58        // Get list of agents we consider in planning
59        let mut agents = tasks.iter().map(|task| task.agent).collect();
60        D::update_visible_agents(start_tick, tick, state_diff, active_agent, &mut agents);
61
62        // Assign idle tasks to agents without a task
63        let (tasks, current_values): (ActiveTasks<D>, _) = agents
64            .into_iter()
65            .map(|agent| {
66                get_task_for_agent(&tasks, agent).map_or_else(
67                    || ActiveTask::new_idle(tick, agent, active_agent),
68                    |task| task.clone(),
69                )
70            })
71            // Set child current values
72            .map(|task| {
73                let agent = task.agent;
74                (task, (agent, D::get_current_value(tick, state_diff, agent)))
75            })
76            .unzip();
77
78        NodeInner {
79            active_agent,
80            diff,
81            tick,
82            tasks,
83            current_values,
84        }
85    }
86
87    /// Returns the agent who owns the node
88    pub fn agent(&self) -> AgentId {
89        self.active_agent
90    }
91
92    /// Returns the tick
93    pub fn tick(&self) -> u64 {
94        self.tick
95    }
96
97    /// Returns all agents that are in considered by this node
98    pub fn agents(&self) -> BTreeSet<AgentId> {
99        self.tasks.iter().map(|task| task.agent).collect()
100    }
101
102    /// Returns the diff of current node.
103    pub fn diff(&self) -> &D::Diff {
104        &self.diff
105    }
106
107    /// Returns the current value from an agent, panic if not present in the node
108    pub fn current_value(&self, agent: AgentId) -> AgentValue {
109        self.current_values
110            .get(&agent)
111            .copied()
112            .unwrap_or_else(|| AgentValue::new(0.0).unwrap())
113    }
114
115    /// Returns the current value from an agent, compute if not present in the node
116    pub fn current_value_or_compute(&self, agent: AgentId, initial_state: &D::State) -> AgentValue {
117        self.current_values.get(&agent).copied().unwrap_or_else(|| {
118            D::get_current_value(
119                self.tick,
120                StateDiffRef::new(initial_state, &self.diff),
121                agent,
122            )
123        })
124    }
125
126    /// Returns all current values
127    pub fn current_values(&self) -> &BTreeMap<AgentId, AgentValue> {
128        &self.current_values
129    }
130
131    // Returns the size in bytes
132    pub fn size(&self, task_size: fn(&dyn Task<D>) -> usize) -> usize {
133        let mut size = 0;
134
135        size += mem::size_of::<Self>();
136        size += self.current_values.len() * mem::size_of::<(AgentId, f32)>();
137
138        for task in &self.tasks {
139            size += task.size(task_size);
140        }
141
142        size
143    }
144}
145
146impl<D: Domain> Hash for NodeInner<D> {
147    fn hash<H: Hasher>(&self, hasher: &mut H) {
148        self.active_agent.hash(hasher);
149        self.diff.hash(hasher);
150        self.tasks.hash(hasher);
151        self.tick.hash(hasher);
152    }
153}
154
155impl<D: Domain> PartialEq for NodeInner<D> {
156    fn eq(&self, other: &Self) -> bool {
157        self.active_agent.eq(&other.active_agent)
158            && self.diff.eq(&other.diff)
159            && self.tasks.eq(&other.tasks)
160            && self.tick.eq(&other.tick)
161    }
162}
163
164impl<D: Domain> Eq for NodeInner<D> {}