1use std::collections::{HashMap, HashSet, hash_map::Entry};
2
3use kcode_k1_kmap_format::ConnectionTier;
4pub use kcode_k1_kmap_format::{Node, NodeId};
5use kcode_k1_kmap_selection::score;
6
7pub const PREVIEW_COST: f64 = 0.3;
8pub const NARRATIVE_COST: f64 = 1.0;
9
10const PREVIEW_TENTHS: u64 = 3;
11const NARRATIVE_TENTHS: u64 = 10;
12
13type CandidateFilter<'a> = dyn FnMut(&[NodeId]) -> Result<Vec<NodeId>, String> + 'a;
14
15#[derive(Clone, Debug, PartialEq)]
16pub struct LoadedNode {
17 pub node_id: NodeId,
18 pub title: String,
19 pub navigation_hint: String,
20 pub narrative: Option<String>,
21}
22
23#[derive(Clone, Debug, PartialEq)]
24pub struct OpenResult {
25 pub nodes: Vec<LoadedNode>,
26 pub automatic_attention_spent: f64,
27}
28
29#[derive(Clone, Copy, Debug, Eq, PartialEq)]
30pub enum OpenMode {
31 Full,
32 NavigationOnly,
33}
34
35pub fn open_node(
36 node_id: NodeId,
37 budget: f64,
38 temperature: f64,
39 mode: OpenMode,
40 load_node: impl FnMut(NodeId) -> Result<Option<Node>, String>,
41 candidate_filter: impl FnMut(&[NodeId]) -> Result<Vec<NodeId>, String>,
42) -> Result<OpenResult, String> {
43 open_node_with_random(
44 node_id,
45 budget,
46 temperature,
47 mode,
48 load_node,
49 candidate_filter,
50 kcode_k1_kmap_selection::os_random_unit,
51 )
52}
53
54fn open_node_with_random(
55 node_id: NodeId,
56 budget: f64,
57 temperature: f64,
58 mode: OpenMode,
59 mut load_node: impl FnMut(NodeId) -> Result<Option<Node>, String>,
60 mut candidate_filter: impl FnMut(&[NodeId]) -> Result<Vec<NodeId>, String>,
61 mut random: impl FnMut() -> Result<f64, String>,
62) -> Result<OpenResult, String> {
63 if !budget.is_finite() || budget < 0.0 {
64 return Err("budget must be finite and nonnegative".to_owned());
65 }
66 if !temperature.is_finite() || temperature < 0.0 {
67 return Err("temperature must be finite and nonnegative".to_owned());
68 }
69 let mut engine = Engine {
70 budget,
71 temperature,
72 load_node: &mut load_node,
73 candidate_filter: &mut candidate_filter,
74 random: &mut random,
75 decisions: HashMap::from([(node_id, true)]),
76 };
77 let root = engine.required(node_id)?;
78 match mode {
79 OpenMode::Full => engine.full(node_id, root),
80 OpenMode::NavigationOnly => engine.navigation_only(node_id, root),
81 }
82}
83
84struct Engine<'a> {
85 budget: f64,
86 temperature: f64,
87 load_node: &'a mut dyn FnMut(NodeId) -> Result<Option<Node>, String>,
88 candidate_filter: &'a mut CandidateFilter<'a>,
89 random: &'a mut dyn FnMut() -> Result<f64, String>,
90 decisions: HashMap<NodeId, bool>,
91}
92
93impl Engine<'_> {
94 fn full(&mut self, node_id: NodeId, root: Node) -> Result<OpenResult, String> {
95 let mut outputs = vec![loaded(node_id, &root, true)];
96 let mut states = HashMap::from([(node_id, NodeState::Opened)]);
97 let root_targets = self.navigation_targets(&root, 1.0, &states)?;
98 for occurrence in root_targets {
99 let target = occurrence.target;
100 insert_preview(target, self.required(target)?, &mut outputs, &mut states);
101 }
102 let mut frontier = self.outgoing(&root, 1.0, &states)?;
103 let mut spent = 0_u64;
104
105 loop {
106 let mut candidates = Vec::new();
107 let mut opening_costs = HashMap::new();
108 for (index, occurrence) in frontier.iter().enumerate() {
109 if matches!(states.get(&occurrence.target), Some(NodeState::Opened)) {
110 continue;
111 }
112 if !occurrence.strength.is_finite() || occurrence.strength <= 0.0 {
113 continue;
114 }
115 let cost = match states.get(&occurrence.target) {
116 None => PREVIEW_TENTHS,
117 Some(NodeState::Previewed(node)) => {
118 if let Some(cost) = opening_costs.get(&occurrence.target) {
119 *cost
120 } else {
121 let previews = self
122 .navigation_targets(node, occurrence.strength, &states)?
123 .len();
124 let cost =
125 attention(preview_cost(previews)?.checked_add(NARRATIVE_TENTHS))?;
126 let _ = opening_costs.insert(occurrence.target, cost);
127 cost
128 }
129 }
130 Some(NodeState::Opened) => continue,
131 };
132 if affordable(spent, cost, self.budget)? {
133 candidates.push((index, occurrence.strength, cost));
134 }
135 }
136 if candidates.is_empty() {
137 break;
138 }
139
140 let choice =
141 kcode_k1_kmap_selection::choose(&candidates, self.temperature, &mut self.random)?;
142 let (selected, _, cost) = candidates[choice];
143 let occurrence = frontier[selected].clone();
144 if let Some(NodeState::Previewed(node)) = states.get(&occurrence.target).cloned() {
145 let _ = states.insert(occurrence.target, NodeState::Opened);
146 let guarantees = self.navigation_targets(&node, occurrence.strength, &states)?;
147 outputs
148 .iter_mut()
149 .find(|node| node.node_id == occurrence.target)
150 .ok_or_else(|| "previewed Kmap node had no output".to_owned())?
151 .narrative = Some(node.narrative.clone());
152 frontier.retain(|entry| entry.target != occurrence.target);
153 for guarantee in guarantees {
154 let target = guarantee.target;
155 insert_preview(target, self.required(target)?, &mut outputs, &mut states);
156 }
157 frontier.extend(self.outgoing(&node, occurrence.strength, &states)?);
158 } else {
159 insert_preview(
160 occurrence.target,
161 self.required(occurrence.target)?,
162 &mut outputs,
163 &mut states,
164 );
165 }
166 spent = attention(spent.checked_add(cost))?;
167 }
168 Ok(result(outputs, spent))
169 }
170
171 fn navigation_only(&mut self, node_id: NodeId, root: Node) -> Result<OpenResult, String> {
172 let mut outputs = vec![loaded(node_id, &root, false)];
173 let mut states = HashMap::from([(node_id, NodeState::Opened)]);
174 let root_targets = self.navigation_targets(&root, 1.0, &states)?;
175 let root_cost = preview_cost(root_targets.len())?;
176 if !affordable(0, root_cost, self.budget)? {
177 return Ok(result(outputs, 0));
178 }
179 let mut root_nodes = Vec::with_capacity(root_targets.len());
180 for occurrence in root_targets {
181 let node = self.required(occurrence.target)?;
182 outputs.push(loaded(occurrence.target, &node, false));
183 let _ = states.insert(occurrence.target, NodeState::Opened);
184 root_nodes.push((node, occurrence.strength));
185 }
186
187 let mut spent = root_cost;
188 let mut frontier = self.outgoing(&root, 1.0, &states)?;
189 for (node, strength) in &root_nodes {
190 frontier.extend(self.outgoing(node, *strength, &states)?);
191 }
192 loop {
193 let candidates: Vec<(usize, f64, u64)> = frontier
194 .iter()
195 .enumerate()
196 .filter(|(_, occurrence)| !states.contains_key(&occurrence.target))
197 .filter_map(|(index, occurrence)| {
198 (occurrence.strength.is_finite() && occurrence.strength > 0.0).then_some((
199 index,
200 occurrence.strength,
201 0,
202 ))
203 })
204 .collect();
205 if candidates.is_empty() {
206 break;
207 }
208
209 let choice =
210 kcode_k1_kmap_selection::choose(&candidates, self.temperature, &mut self.random)?;
211 let occurrence = frontier[candidates[choice].0].clone();
212 let node = self.required(occurrence.target)?;
213 let _ = states.insert(occurrence.target, NodeState::Opened);
214 let children = self.navigation_targets(&node, occurrence.strength, &states)?;
215 let cost = preview_cost(attention(children.len().checked_add(1))?)?;
216 if !affordable(spent, cost, self.budget)? {
217 break;
218 }
219 outputs.push(loaded(occurrence.target, &node, false));
220 let mut child_nodes = Vec::with_capacity(children.len());
221 for child in children {
222 let node = self.required(child.target)?;
223 let _ = states.insert(child.target, NodeState::Opened);
224 outputs.push(loaded(child.target, &node, false));
225 child_nodes.push((node, child.strength));
226 }
227 let mut additions = self.outgoing(&node, occurrence.strength, &states)?;
228 for (child, strength) in &child_nodes {
229 additions.extend(self.outgoing(child, *strength, &states)?);
230 }
231 frontier.retain(|entry| !states.contains_key(&entry.target));
232 frontier.extend(additions);
233 spent = attention(spent.checked_add(cost))?;
234 }
235 Ok(result(outputs, spent))
236 }
237
238 fn required(&mut self, node_id: NodeId) -> Result<Node, String> {
239 (self.load_node)(node_id)?
240 .ok_or_else(|| format!("authorized Kmap target {node_id:?} is missing"))
241 }
242
243 fn resolve_candidates(&mut self, node: &Node) -> Result<(), String> {
244 let mut expected = HashSet::with_capacity(node.connections.len());
245 let requested: Vec<NodeId> = node
246 .connections
247 .iter()
248 .map(|connection| connection.target)
249 .filter(|target| !self.decisions.contains_key(target) && expected.insert(*target))
250 .collect();
251 if requested.is_empty() {
252 return Ok(());
253 }
254 let returned = (self.candidate_filter)(&requested)
255 .map_err(|error| format!("Kmap candidate filter failed: {error}"))?;
256 let mut allowed = HashSet::with_capacity(returned.len());
257 for target in returned {
258 if !expected.contains(&target) {
259 return Err(format!(
260 "Kmap candidate filter returned unrequested node {target:?}"
261 ));
262 }
263 if !allowed.insert(target) {
264 return Err(format!(
265 "Kmap candidate filter returned duplicate node {target:?}"
266 ));
267 }
268 }
269 self.decisions.extend(
270 requested
271 .into_iter()
272 .map(|target| (target, allowed.contains(&target))),
273 );
274 Ok(())
275 }
276
277 fn navigation_targets(
278 &mut self,
279 node: &Node,
280 inherited_strength: f64,
281 states: &HashMap<NodeId, NodeState>,
282 ) -> Result<Vec<Occurrence>, String> {
283 self.resolve_candidates(node)?;
284 Ok(node
285 .connections
286 .iter()
287 .filter(|connection| {
288 !states.contains_key(&connection.target)
289 && self.decisions.get(&connection.target) == Some(&true)
290 && connection.tier == ConnectionTier::Navigation
291 })
292 .map(|connection| Occurrence {
293 target: connection.target,
294 strength: score(connection.weight.value, inherited_strength),
295 })
296 .collect())
297 }
298
299 fn outgoing(
300 &mut self,
301 node: &Node,
302 inherited_strength: f64,
303 states: &HashMap<NodeId, NodeState>,
304 ) -> Result<Vec<Occurrence>, String> {
305 self.resolve_candidates(node)?;
306 Ok(node
307 .connections
308 .iter()
309 .filter(|connection| {
310 !matches!(states.get(&connection.target), Some(NodeState::Opened))
311 && self.decisions.get(&connection.target) == Some(&true)
312 })
313 .map(|connection| Occurrence {
314 target: connection.target,
315 strength: score(connection.weight.value, inherited_strength),
316 })
317 .collect())
318 }
319}
320
321#[derive(Clone)]
322enum NodeState {
323 Previewed(Node),
324 Opened,
325}
326
327#[derive(Clone)]
328struct Occurrence {
329 target: NodeId,
330 strength: f64,
331}
332
333fn loaded(node_id: NodeId, node: &Node, opened: bool) -> LoadedNode {
334 LoadedNode {
335 node_id,
336 title: node.title.clone(),
337 navigation_hint: node.navigation_hint.clone(),
338 narrative: opened.then(|| node.narrative.clone()),
339 }
340}
341
342fn insert_preview(
343 node_id: NodeId,
344 node: Node,
345 outputs: &mut Vec<LoadedNode>,
346 states: &mut HashMap<NodeId, NodeState>,
347) {
348 if let Entry::Vacant(entry) = states.entry(node_id) {
349 outputs.push(loaded(node_id, &node, false));
350 entry.insert(NodeState::Previewed(node));
351 }
352}
353
354fn result(nodes: Vec<LoadedNode>, spent: u64) -> OpenResult {
355 OpenResult {
356 nodes,
357 automatic_attention_spent: spent as f64 / 10.0,
358 }
359}
360
361fn attention<T>(value: Option<T>) -> Result<T, String> {
362 value.ok_or_else(|| "Kmap attention cost overflow".to_owned())
363}
364
365fn preview_cost(count: usize) -> Result<u64, String> {
366 let count = attention(u64::try_from(count).ok())?;
367 attention(count.checked_mul(PREVIEW_TENTHS))
368}
369
370fn affordable(spent: u64, cost: u64, budget: f64) -> Result<bool, String> {
371 let total = attention(spent.checked_add(cost))? as f64 / 10.0;
372 let tolerance = 8.0 * f64::EPSILON * total.abs().max(budget.abs()).max(1.0);
373 Ok(total <= budget || total - budget <= tolerance)
374}
375
376#[cfg(test)]
377mod tests {
378 use std::cell::{Cell, RefCell};
379
380 use kcode_k1_kmap_format::{Connection, ConnectionTier};
381
382 use super::{Node, NodeId, OpenMode, OpenResult, open_node_with_random};
383
384 fn id(value: u64) -> NodeId {
385 let mut bytes = [0; 12];
386 bytes[..8].copy_from_slice(&value.to_le_bytes());
387 NodeId(bytes)
388 }
389
390 fn node(connections: Vec<Connection>) -> Node {
391 Node {
392 title: String::new(),
393 navigation_hint: String::new(),
394 narrative: String::new(),
395 connections,
396 }
397 }
398
399 fn edge(target: NodeId, tier: ConnectionTier, weight: f64) -> Connection {
400 let mut connection = Connection::new(target, tier);
401 connection.weight.value = weight;
402 connection
403 }
404
405 fn run_filter(returned: Result<Vec<NodeId>, String>) -> Result<OpenResult, String> {
406 let root = id(0);
407 let root_node = node(vec![Connection::new(id(1), ConnectionTier::Automated)]);
408 open_node_with_random(
409 root,
410 1.0,
411 1.0,
412 OpenMode::NavigationOnly,
413 |node_id| Ok((node_id == root).then(|| root_node.clone())),
414 move |_| returned.clone(),
415 || panic!("RNG invoked for rejected or denied candidates"),
416 )
417 }
418
419 #[test]
420 fn batches_large_denial_before_reads_and_randomness() {
421 let root = id(0);
422 let root_node = node(
423 (1..=10_000)
424 .map(|value| Connection::new(id(value), ConnectionTier::Automated))
425 .collect(),
426 );
427 let filters = Cell::new(0);
428 let result = open_node_with_random(
429 root,
430 10_000.0,
431 1.0,
432 OpenMode::NavigationOnly,
433 |node_id| Ok((node_id == root).then(|| root_node.clone())),
434 |targets| {
435 filters.set(filters.get() + 1);
436 assert_eq!(targets.len(), 10_000);
437 assert_eq!((targets[0], targets[9_999]), (id(1), id(10_000)));
438 Ok(Vec::new())
439 },
440 || panic!("RNG invoked after every candidate was denied"),
441 )
442 .unwrap();
443 assert_eq!(result.nodes.len(), 1);
444 assert_eq!(filters.get(), 1);
445 }
446
447 #[test]
448 fn propagates_and_validates_filter_results() {
449 assert_eq!(
450 run_filter(Err("Access unavailable".to_owned())).unwrap_err(),
451 "Kmap candidate filter failed: Access unavailable"
452 );
453 assert!(
454 run_filter(Ok(vec![id(1), id(1)]))
455 .unwrap_err()
456 .contains("returned duplicate node")
457 );
458 assert!(
459 run_filter(Ok(vec![id(2)]))
460 .unwrap_err()
461 .contains("returned unrequested node")
462 );
463 assert_eq!(run_filter(Ok(Vec::new())).unwrap().nodes.len(), 1);
464 }
465
466 #[test]
467 fn memoizes_visibility_across_repeated_paths() {
468 let (root, a, f, x, b, c, e, d) = (id(0), id(1), id(2), id(3), id(4), id(5), id(6), id(7));
469 let root_node = node(vec![
470 edge(a, ConnectionTier::Navigation, 0.5),
471 edge(f, ConnectionTier::Navigation, 0.1),
472 edge(x, ConnectionTier::Automated, 0.15),
473 ]);
474 let a_node = node(vec![edge(b, ConnectionTier::Automated, 0.4)]);
475 let f_node = node(vec![edge(b, ConnectionTier::Automated, 0.9)]);
476 let b_node = node(vec![
477 edge(e, ConnectionTier::Automated, 0.8),
478 edge(c, ConnectionTier::Navigation, 0.5),
479 ]);
480 let c_node = node(vec![edge(d, ConnectionTier::Automated, 1.0)]);
481 let leaf = node(Vec::new());
482 let batches = RefCell::new(Vec::new());
483 let result = open_node_with_random(
484 root,
485 2.1,
486 0.0,
487 OpenMode::NavigationOnly,
488 |node_id| {
489 Ok(match node_id {
490 value if value == root => Some(root_node.clone()),
491 value if value == a => Some(a_node.clone()),
492 value if value == f => Some(f_node.clone()),
493 value if value == b => Some(b_node.clone()),
494 value if value == c => Some(c_node.clone()),
495 value if [x, e, d].contains(&value) => Some(leaf.clone()),
496 _ => None,
497 })
498 },
499 |targets| {
500 batches.borrow_mut().push(targets.to_vec());
501 Ok(targets.to_vec())
502 },
503 || Ok(0.0),
504 )
505 .unwrap();
506 let ids = result
507 .nodes
508 .iter()
509 .map(|node| node.node_id)
510 .collect::<Vec<_>>();
511 assert_eq!(ids, vec![root, a, f, b, c, e, x, d]);
512 assert_eq!(
513 batches.into_inner(),
514 vec![vec![a, f, x], vec![b], vec![e, c], vec![d]]
515 );
516 }
517}