Skip to main content

kcode_k1_kmap_loader/
lib.rs

1use std::collections::{HashMap, HashSet, hash_map::Entry};
2
3use kcode_k1_kmap_format::ConnectionTier;
4pub use kcode_k1_kmap_format::{Node, NodeId};
5pub use kcode_k1_kmap_selection::DEPTH_DECAY;
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, &states)?;
98        for target in root_targets {
99            insert_preview(target, self.required(target)?, &mut outputs, &mut states);
100        }
101        let mut frontier = self.outgoing(&root, 1, &states)?;
102        let mut spent = 0_u64;
103
104        loop {
105            let mut candidates = Vec::new();
106            let mut opening_costs = HashMap::new();
107            for (index, occurrence) in frontier.iter().enumerate() {
108                if matches!(states.get(&occurrence.target), Some(NodeState::Opened)) {
109                    continue;
110                }
111                let effective = kcode_k1_kmap_selection::score(occurrence.value, occurrence.depth);
112                if !effective.is_finite() || effective <= 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.navigation_targets(node, &states)?.len();
122                            let cost =
123                                attention(preview_cost(previews)?.checked_add(NARRATIVE_TENTHS))?;
124                            let _ = opening_costs.insert(occurrence.target, cost);
125                            cost
126                        }
127                    }
128                    Some(NodeState::Opened) => continue,
129                };
130                if affordable(spent, cost, self.budget)? {
131                    candidates.push((index, effective, cost));
132                }
133            }
134            if candidates.is_empty() {
135                break;
136            }
137
138            let choice =
139                kcode_k1_kmap_selection::choose(&candidates, self.temperature, &mut self.random)?;
140            let (selected, _, cost) = candidates[choice];
141            let occurrence = frontier[selected].clone();
142            if let Some(NodeState::Previewed(node)) = states.get(&occurrence.target).cloned() {
143                let _ = states.insert(occurrence.target, NodeState::Opened);
144                let guarantees = self.navigation_targets(&node, &states)?;
145                outputs
146                    .iter_mut()
147                    .find(|node| node.node_id == occurrence.target)
148                    .ok_or_else(|| "previewed Kmap node had no output".to_owned())?
149                    .narrative = Some(node.narrative.clone());
150                frontier.retain(|entry| entry.target != occurrence.target);
151                for target in guarantees {
152                    insert_preview(target, self.required(target)?, &mut outputs, &mut states);
153                }
154                frontier.extend(self.outgoing(
155                    &node,
156                    increment_depth(occurrence.depth)?,
157                    &states,
158                )?);
159            } else {
160                insert_preview(
161                    occurrence.target,
162                    self.required(occurrence.target)?,
163                    &mut outputs,
164                    &mut states,
165                );
166            }
167            spent = attention(spent.checked_add(cost))?;
168        }
169        Ok(result(outputs, spent))
170    }
171
172    fn navigation_only(&mut self, node_id: NodeId, root: Node) -> Result<OpenResult, String> {
173        let mut outputs = vec![loaded(node_id, &root, false)];
174        let mut states = HashMap::from([(node_id, NodeState::Opened)]);
175        let root_targets = self.navigation_targets(&root, &states)?;
176        let root_cost = preview_cost(root_targets.len())?;
177        if !affordable(0, root_cost, self.budget)? {
178            return Ok(result(outputs, 0));
179        }
180        let mut root_nodes = Vec::with_capacity(root_targets.len());
181        for target in root_targets {
182            let node = self.required(target)?;
183            outputs.push(loaded(target, &node, false));
184            let _ = states.insert(target, NodeState::Opened);
185            root_nodes.push(node);
186        }
187
188        let mut spent = root_cost;
189        let mut frontier = self.outgoing(&root, 1, &states)?;
190        for node in &root_nodes {
191            frontier.extend(self.outgoing(node, 2, &states)?);
192        }
193        loop {
194            let candidates: Vec<(usize, f64, u64)> = frontier
195                .iter()
196                .enumerate()
197                .filter(|(_, occurrence)| !states.contains_key(&occurrence.target))
198                .filter_map(|(index, occurrence)| {
199                    let effective =
200                        kcode_k1_kmap_selection::score(occurrence.value, occurrence.depth);
201                    (effective.is_finite() && effective > 0.0).then_some((index, effective, 0))
202                })
203                .collect();
204            if candidates.is_empty() {
205                break;
206            }
207
208            let choice =
209                kcode_k1_kmap_selection::choose(&candidates, self.temperature, &mut self.random)?;
210            let occurrence = frontier[candidates[choice].0].clone();
211            let node = self.required(occurrence.target)?;
212            let _ = states.insert(occurrence.target, NodeState::Opened);
213            let children = self.navigation_targets(&node, &states)?;
214            let cost = preview_cost(attention(children.len().checked_add(1))?)?;
215            if !affordable(spent, cost, self.budget)? {
216                break;
217            }
218            outputs.push(loaded(occurrence.target, &node, false));
219            let mut child_nodes = Vec::with_capacity(children.len());
220            for target in children {
221                let child = self.required(target)?;
222                let _ = states.insert(target, NodeState::Opened);
223                outputs.push(loaded(target, &child, false));
224                child_nodes.push(child);
225            }
226            let next_depth = increment_depth(occurrence.depth)?;
227            let mut additions = self.outgoing(&node, next_depth, &states)?;
228            for child in &child_nodes {
229                additions.extend(self.outgoing(child, increment_depth(next_depth)?, &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        states: &HashMap<NodeId, NodeState>,
281    ) -> Result<Vec<NodeId>, String> {
282        self.resolve_candidates(node)?;
283        Ok(node
284            .connections
285            .iter()
286            .filter(|connection| {
287                !states.contains_key(&connection.target)
288                    && self.decisions.get(&connection.target) == Some(&true)
289                    && connection.tier == ConnectionTier::Navigation
290            })
291            .map(|connection| connection.target)
292            .collect())
293    }
294
295    fn outgoing(
296        &mut self,
297        node: &Node,
298        depth: usize,
299        states: &HashMap<NodeId, NodeState>,
300    ) -> Result<Vec<Occurrence>, String> {
301        self.resolve_candidates(node)?;
302        Ok(node
303            .connections
304            .iter()
305            .filter(|connection| {
306                !matches!(states.get(&connection.target), Some(NodeState::Opened))
307                    && self.decisions.get(&connection.target) == Some(&true)
308            })
309            .map(|connection| Occurrence {
310                target: connection.target,
311                value: connection.weight.value,
312                depth,
313            })
314            .collect())
315    }
316}
317
318#[derive(Clone)]
319enum NodeState {
320    Previewed(Node),
321    Opened,
322}
323
324#[derive(Clone)]
325struct Occurrence {
326    target: NodeId,
327    value: f64,
328    depth: usize,
329}
330
331fn loaded(node_id: NodeId, node: &Node, opened: bool) -> LoadedNode {
332    LoadedNode {
333        node_id,
334        title: node.title.clone(),
335        navigation_hint: node.navigation_hint.clone(),
336        narrative: opened.then(|| node.narrative.clone()),
337    }
338}
339
340fn insert_preview(
341    node_id: NodeId,
342    node: Node,
343    outputs: &mut Vec<LoadedNode>,
344    states: &mut HashMap<NodeId, NodeState>,
345) {
346    if let Entry::Vacant(entry) = states.entry(node_id) {
347        outputs.push(loaded(node_id, &node, false));
348        entry.insert(NodeState::Previewed(node));
349    }
350}
351
352fn result(nodes: Vec<LoadedNode>, spent: u64) -> OpenResult {
353    OpenResult {
354        nodes,
355        automatic_attention_spent: spent as f64 / 10.0,
356    }
357}
358
359fn increment_depth(depth: usize) -> Result<usize, String> {
360    depth
361        .checked_add(1)
362        .ok_or_else(|| "Kmap traversal depth overflow".to_owned())
363}
364
365fn attention<T>(value: Option<T>) -> Result<T, String> {
366    value.ok_or_else(|| "Kmap attention cost overflow".to_owned())
367}
368
369fn preview_cost(count: usize) -> Result<u64, String> {
370    let count = attention(u64::try_from(count).ok())?;
371    attention(count.checked_mul(PREVIEW_TENTHS))
372}
373
374fn affordable(spent: u64, cost: u64, budget: f64) -> Result<bool, String> {
375    let total = attention(spent.checked_add(cost))? as f64 / 10.0;
376    let tolerance = 8.0 * f64::EPSILON * total.abs().max(budget.abs()).max(1.0);
377    Ok(total <= budget || total - budget <= tolerance)
378}
379
380#[cfg(test)]
381mod tests {
382    use std::cell::{Cell, RefCell};
383
384    use kcode_k1_kmap_format::{Connection, ConnectionTier};
385
386    use super::{Node, NodeId, OpenMode, OpenResult, open_node_with_random};
387
388    fn id(value: u64) -> NodeId {
389        let mut bytes = [0; 12];
390        bytes[..8].copy_from_slice(&value.to_le_bytes());
391        NodeId(bytes)
392    }
393
394    fn node(connections: Vec<Connection>) -> Node {
395        Node {
396            title: String::new(),
397            navigation_hint: String::new(),
398            narrative: String::new(),
399            connections,
400        }
401    }
402
403    fn run_filter(returned: Result<Vec<NodeId>, String>) -> Result<OpenResult, String> {
404        let root = id(0);
405        let root_node = node(vec![Connection::new(id(1), ConnectionTier::Automated)]);
406        open_node_with_random(
407            root,
408            1.0,
409            1.0,
410            OpenMode::NavigationOnly,
411            |node_id| Ok((node_id == root).then(|| root_node.clone())),
412            move |_| returned.clone(),
413            || panic!("RNG invoked for rejected or denied candidates"),
414        )
415    }
416
417    #[test]
418    fn batches_large_denial_before_reads_and_randomness() {
419        let root = id(0);
420        let root_node = node(
421            (1..=10_000)
422                .map(|value| Connection::new(id(value), ConnectionTier::Automated))
423                .collect(),
424        );
425        let filters = Cell::new(0);
426        let result = open_node_with_random(
427            root,
428            10_000.0,
429            1.0,
430            OpenMode::NavigationOnly,
431            |node_id| Ok((node_id == root).then(|| root_node.clone())),
432            |targets| {
433                filters.set(filters.get() + 1);
434                assert_eq!(targets.len(), 10_000);
435                assert_eq!((targets[0], targets[9_999]), (id(1), id(10_000)));
436                Ok(Vec::new())
437            },
438            || panic!("RNG invoked after every candidate was denied"),
439        )
440        .unwrap();
441        assert_eq!(result.nodes.len(), 1);
442        assert_eq!(filters.get(), 1);
443    }
444
445    #[test]
446    fn propagates_and_validates_filter_results() {
447        assert_eq!(
448            run_filter(Err("Access unavailable".to_owned())).unwrap_err(),
449            "Kmap candidate filter failed: Access unavailable"
450        );
451        assert!(
452            run_filter(Ok(vec![id(1), id(1)]))
453                .unwrap_err()
454                .contains("returned duplicate node")
455        );
456        assert!(
457            run_filter(Ok(vec![id(2)]))
458                .unwrap_err()
459                .contains("returned unrequested node")
460        );
461        assert_eq!(run_filter(Ok(Vec::new())).unwrap().nodes.len(), 1);
462    }
463
464    #[test]
465    fn memoizes_visibility_across_repeated_paths() {
466        let (root, a, b, target) = (id(0), id(1), id(2), id(3));
467        let root_node = node(vec![
468            Connection::new(a, ConnectionTier::Navigation),
469            Connection::new(b, ConnectionTier::Navigation),
470        ]);
471        let branch = node(vec![Connection::new(target, ConnectionTier::Automated)]);
472        let target_node = node(Vec::new());
473        let batches = RefCell::new(Vec::new());
474        open_node_with_random(
475            root,
476            0.6,
477            1.0,
478            OpenMode::NavigationOnly,
479            |node_id| {
480                Ok(match node_id {
481                    value if value == root => Some(root_node.clone()),
482                    value if value == a || value == b => Some(branch.clone()),
483                    value if value == target => Some(target_node.clone()),
484                    _ => None,
485                })
486            },
487            |targets| {
488                batches.borrow_mut().push(targets.to_vec());
489                Ok(targets.to_vec())
490            },
491            || Ok(0.0),
492        )
493        .unwrap();
494        assert_eq!(batches.into_inner(), vec![vec![a, b], vec![target]]);
495    }
496}