Skip to main content

kcode_k1_kmap_loader/
lib.rs

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