kcode_k1_kmap_loader/

lib.rs

use std::collections::HashMap;

use kcode_k1_kmap_format::ConnectionTier;
pub use kcode_k1_kmap_format::{Node, NodeId};

pub const PREVIEW_COST: f64 = 0.3;
pub const NARRATIVE_COST: f64 = 1.0;
pub const DEPTH_DECAY: f64 = 0.7;

const PREVIEW_TENTHS: u64 = 3;
const NARRATIVE_TENTHS: u64 = 10;

#[derive(Clone, Debug, PartialEq)]
pub struct LoadedNode {
    pub node_id: NodeId,
    pub title: String,
    pub navigation_hint: String,
    pub narrative: Option<String>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct OpenResult {
    pub nodes: Vec<LoadedNode>,
    pub automatic_attention_spent: f64,
}

pub fn open_node<L, A>(
    node_id: NodeId,
    budget: f64,
    temperature: f64,
    load_node: L,
    access_filter: A,
) -> Result<OpenResult, String>
where
    L: FnMut(NodeId) -> Result<Option<Node>, String>,
    A: FnMut(NodeId) -> bool,
{
    open_node_with_random(
        node_id,
        budget,
        temperature,
        load_node,
        access_filter,
        os_random_unit,
    )
}

fn open_node_with_random<L, A, R>(
    node_id: NodeId,
    budget: f64,
    temperature: f64,
    mut load_node: L,
    mut access_filter: A,
    mut random: R,
) -> Result<OpenResult, String>
where
    L: FnMut(NodeId) -> Result<Option<Node>, String>,
    A: FnMut(NodeId) -> bool,
    R: FnMut() -> Result<f64, String>,
{
    if !budget.is_finite() || budget < 0.0 {
        return Err("budget must be finite and nonnegative".to_owned());
    }
    if !temperature.is_finite() || temperature < 0.0 {
        return Err("temperature must be finite and nonnegative".to_owned());
    }

    let root = required_node(node_id, &mut load_node)?;
    let mut outputs = vec![loaded(node_id, &root, true)];
    let mut output_indices = HashMap::from([(node_id, 0)]);
    let mut states = HashMap::from([(node_id, NodeState::Opened)]);
    let mut cache = HashMap::new();
    let mut decisions = HashMap::new();
    let mut frontier = Vec::new();
    let mut spent = 0_u64;

    for connection in &root.connections {
        let target = connection.target;
        if states.get(&target) == Some(&NodeState::Opened) {
            continue;
        }
        if !allowed(target, &mut decisions, &mut access_filter) {
            continue;
        }
        if connection.tier == ConnectionTier::Navigation && !states.contains_key(&target) {
            let node = required_node(target, &mut load_node)?;
            insert_preview(
                target,
                node,
                &mut outputs,
                &mut output_indices,
                &mut states,
                &mut cache,
            )?;
        }
        frontier.push(Occurrence {
            target,
            value: connection.weight.value,
            depth: 1,
        });
    }

    loop {
        let mut candidates = Vec::new();
        let mut opening_costs = HashMap::new();
        for (index, occurrence) in frontier.iter().enumerate() {
            if states.get(&occurrence.target) == Some(&NodeState::Opened) {
                continue;
            }
            let effective = occurrence.value * DEPTH_DECAY.powf(occurrence.depth as f64);
            if !effective.is_finite() || effective <= 0.0 {
                continue;
            }
            let cost = match states.get(&occurrence.target) {
                None => PREVIEW_TENTHS,
                Some(NodeState::Previewed) => {
                    if let Some(cost) = opening_costs.get(&occurrence.target) {
                        *cost
                    } else {
                        let node = cache
                            .get(&occurrence.target)
                            .ok_or_else(|| "previewed Kmap node was not cached".to_owned())?;
                        let cost = opening_cost(node, &states, &mut decisions, &mut access_filter)?;
                        let _ = opening_costs.insert(occurrence.target, cost);
                        cost
                    }
                }
                Some(NodeState::Opened) => continue,
            };
            if affordable(spent, cost, budget)? {
                candidates.push((index, effective, cost));
            }
        }
        if candidates.is_empty() {
            break;
        }

        let choice = choose(&candidates, temperature, &mut random)?;
        let (selected, _, cost) = candidates[choice];
        let occurrence = frontier[selected].clone();
        match states.get(&occurrence.target).copied() {
            None => {
                let node = required_node(occurrence.target, &mut load_node)?;
                insert_preview(
                    occurrence.target,
                    node,
                    &mut outputs,
                    &mut output_indices,
                    &mut states,
                    &mut cache,
                )?;
            }
            Some(NodeState::Previewed) => {
                let node = cache
                    .get(&occurrence.target)
                    .cloned()
                    .ok_or_else(|| "previewed Kmap node was not cached".to_owned())?;
                let guarantees =
                    navigation_guarantees(&node, &states, &mut decisions, &mut access_filter);
                let _ = states.insert(occurrence.target, NodeState::Opened);
                let output = output_indices
                    .get(&occurrence.target)
                    .copied()
                    .ok_or_else(|| "previewed Kmap node had no output".to_owned())?;
                outputs[output].narrative = Some(node.narrative.clone());
                frontier.retain(|entry| entry.target != occurrence.target);

                for target in guarantees {
                    let guaranteed = required_node(target, &mut load_node)?;
                    insert_preview(
                        target,
                        guaranteed,
                        &mut outputs,
                        &mut output_indices,
                        &mut states,
                        &mut cache,
                    )?;
                }

                let depth = occurrence
                    .depth
                    .checked_add(1)
                    .ok_or_else(|| "Kmap traversal depth overflow".to_owned())?;
                for connection in &node.connections {
                    let target = connection.target;
                    if states.get(&target) == Some(&NodeState::Opened) {
                        continue;
                    }
                    if !allowed(target, &mut decisions, &mut access_filter) {
                        continue;
                    }
                    frontier.push(Occurrence {
                        target,
                        value: connection.weight.value,
                        depth,
                    });
                }
            }
            Some(NodeState::Opened) => {
                return Err("opened Kmap node remained selectable".to_owned());
            }
        }
        spent = spent
            .checked_add(cost)
            .ok_or_else(|| "Kmap attention cost overflow".to_owned())?;
    }

    Ok(OpenResult {
        nodes: outputs,
        automatic_attention_spent: spent as f64 / 10.0,
    })
}

#[derive(Clone, Copy, Eq, PartialEq)]
enum NodeState {
    Previewed,
    Opened,
}

#[derive(Clone)]
struct Occurrence {
    target: NodeId,
    value: f64,
    depth: usize,
}

fn required_node<L>(node_id: NodeId, load_node: &mut L) -> Result<Node, String>
where
    L: FnMut(NodeId) -> Result<Option<Node>, String>,
{
    load_node(node_id)?.ok_or_else(|| format!("authorized Kmap target {node_id:?} is missing"))
}

fn loaded(node_id: NodeId, node: &Node, opened: bool) -> LoadedNode {
    LoadedNode {
        node_id,
        title: node.title.clone(),
        navigation_hint: node.navigation_hint.clone(),
        narrative: opened.then(|| node.narrative.clone()),
    }
}

fn insert_preview(
    node_id: NodeId,
    node: Node,
    outputs: &mut Vec<LoadedNode>,
    output_indices: &mut HashMap<NodeId, usize>,
    states: &mut HashMap<NodeId, NodeState>,
    cache: &mut HashMap<NodeId, Node>,
) -> Result<(), String> {
    if states.contains_key(&node_id) {
        return Ok(());
    }
    if output_indices.contains_key(&node_id) {
        return Err("Kmap output existed without traversal state".to_owned());
    }
    let index = outputs.len();
    outputs.push(loaded(node_id, &node, false));
    let _ = output_indices.insert(node_id, index);
    let _ = states.insert(node_id, NodeState::Previewed);
    let _ = cache.insert(node_id, node);
    Ok(())
}

fn allowed<A>(node_id: NodeId, decisions: &mut HashMap<NodeId, bool>, access_filter: &mut A) -> bool
where
    A: FnMut(NodeId) -> bool,
{
    if let Some(decision) = decisions.get(&node_id) {
        *decision
    } else {
        let decision = access_filter(node_id);
        let _ = decisions.insert(node_id, decision);
        decision
    }
}

fn navigation_guarantees<A>(
    node: &Node,
    states: &HashMap<NodeId, NodeState>,
    decisions: &mut HashMap<NodeId, bool>,
    access_filter: &mut A,
) -> Vec<NodeId>
where
    A: FnMut(NodeId) -> bool,
{
    let mut targets = Vec::new();
    for connection in &node.connections {
        let target = connection.target;
        if states.contains_key(&target) {
            continue;
        }
        if !allowed(target, decisions, access_filter) {
            continue;
        }
        if connection.tier == ConnectionTier::Navigation && !targets.contains(&target) {
            targets.push(target);
        }
    }
    targets
}

fn opening_cost<A>(
    node: &Node,
    states: &HashMap<NodeId, NodeState>,
    decisions: &mut HashMap<NodeId, bool>,
    access_filter: &mut A,
) -> Result<u64, String>
where
    A: FnMut(NodeId) -> bool,
{
    let count = navigation_guarantees(node, states, decisions, access_filter).len() as u64;
    count
        .checked_mul(PREVIEW_TENTHS)
        .and_then(|cost| cost.checked_add(NARRATIVE_TENTHS))
        .ok_or_else(|| "Kmap attention cost overflow".to_owned())
}

fn affordable(spent: u64, cost: u64, budget: f64) -> Result<bool, String> {
    let total = spent
        .checked_add(cost)
        .ok_or_else(|| "Kmap attention cost overflow".to_owned())? as f64
        / 10.0;
    let tolerance = 8.0 * f64::EPSILON * total.abs().max(budget.abs()).max(1.0);
    Ok(total <= budget || total - budget <= tolerance)
}

fn choose<R>(
    candidates: &[(usize, f64, u64)],
    temperature: f64,
    random: &mut R,
) -> Result<usize, String>
where
    R: FnMut() -> Result<f64, String>,
{
    if temperature == 0.0 {
        let maximum = candidates
            .iter()
            .map(|candidate| candidate.1)
            .max_by(f64::total_cmp)
            .ok_or_else(|| "Kmap candidate set was empty".to_owned())?;
        let maxima: Vec<usize> = candidates
            .iter()
            .enumerate()
            .filter(|(_, candidate)| candidate.1.total_cmp(&maximum).is_eq())
            .map(|(index, _)| index)
            .collect();
        let index = (random_unit(random)? * maxima.len() as f64) as usize;
        return maxima
            .get(index.min(maxima.len() - 1))
            .copied()
            .ok_or_else(|| "Kmap maximum candidate set was empty".to_owned());
    }

    let maximum_log = candidates
        .iter()
        .map(|candidate| candidate.1.ln())
        .max_by(f64::total_cmp)
        .ok_or_else(|| "Kmap candidate set was empty".to_owned())?;
    let weights: Vec<f64> = candidates
        .iter()
        .map(|candidate| ((candidate.1.ln() - maximum_log) / temperature).exp())
        .collect();
    let total: f64 = weights.iter().sum();
    let threshold = random_unit(random)? * total;
    let mut cumulative = 0.0;
    for (index, weight) in weights.into_iter().enumerate() {
        cumulative += weight;
        if threshold < cumulative {
            return Ok(index);
        }
    }
    Ok(candidates.len() - 1)
}

fn random_unit<R>(random: &mut R) -> Result<f64, String>
where
    R: FnMut() -> Result<f64, String>,
{
    let value = random()?;
    if value.is_finite() && (0.0..1.0).contains(&value) {
        Ok(value)
    } else {
        Err("Kmap random value must be finite and in [0, 1)".to_owned())
    }
}

fn os_random_unit() -> Result<f64, String> {
    let mut bytes = [0_u8; 8];
    getrandom::fill(&mut bytes).map_err(|error| format!("Kmap randomness failed: {error}"))?;
    let value = u64::from_ne_bytes(bytes) >> 11;
    Ok(value as f64 / (1_u64 << 53) as f64)
}

#[cfg(test)]
mod tests {
    use super::*;
    use kcode_k1_kmap_format::ConnectionTier::{Automated as A, Navigation as N};
    use kcode_k1_kmap_format::{Connection, Weight};

    fn id(v: u8) -> NodeId {
        NodeId([v; 12])
    }

    fn node() -> Node {
        Node::new("n", "h", "x", vec![]).unwrap()
    }

    fn g(edges: &[(u8, u8, ConnectionTier, f64)]) -> HashMap<NodeId, Node> {
        let mut graph = HashMap::new();
        let _ = graph.insert(id(0), node());
        for (source, target, tier, weight) in edges {
            let source = graph.entry(id(*source)).or_insert_with(node);
            source.connections.push(Connection {
                target: id(*target),
                tier: *tier,
                weight: Weight::new(*weight, 1.0).unwrap(),
            });
            let _ = graph.entry(id(*target)).or_insert_with(node);
        }
        graph
    }

    fn open(g: &HashMap<NodeId, Node>, b: f64, t: f64, r: &[f64]) -> Result<OpenResult, String> {
        let mut values = r.iter().copied();
        let load = |key| Ok(g.get(&key).cloned());
        let access = |_| true;
        let random = || Ok(values.next().unwrap_or(0.0));
        open_node_with_random(id(0), b, t, load, access, random)
    }

    #[test]
    fn acceptance() {
        let graph = g(&[
            (0, 1, N, 1.0),
            (0, 1, A, 1.0),
            (0, 2, A, 1.0),
            (0, 2, A, 1.0),
        ]);
        let calls = std::cell::Cell::new(0);
        let result = open_node_with_random(
            id(0),
            0.0,
            0.0,
            |key| Ok((key != id(2)).then(|| graph[&key].clone())),
            |key| {
                calls.set(calls.get() + 1);
                key != id(2)
            },
            || Ok(0.0),
        )
        .unwrap();
        assert_eq!(calls.get(), 2);
        let opened: Vec<_> = result
            .nodes
            .iter()
            .map(|node| node.narrative.is_some())
            .collect();
        assert_eq!(
            (opened, result.automatic_attention_spent),
            (vec![true, false], 0.0)
        );

        let graph = g(&[(0, 1, A, 1.0)]);
        for (b, yes, cost) in [(0.3, false, 0.3), (1.3, true, 1.3)] {
            let r = open(&graph, b, 0.0, &[0.0, 0.0]).unwrap();
            assert_eq!(r.nodes[1].narrative.is_some(), yes);
            assert_eq!(r.automatic_attention_spent, cost);
        }

        let graph = g(&[(0, 1, A, 1.0), (1, 2, N, 1.0)]);
        let s = open(&graph, 1.3, 0.0, &[0.0, 0.0]).unwrap();
        let f = open(&graph, 1.6, 0.0, &[0.0, 0.0]).unwrap();
        assert!(s.nodes[1].narrative.is_none());
        assert_eq!(s.automatic_attention_spent, 0.3);
        assert_eq!(f.automatic_attention_spent, 1.6);
        assert!(f.nodes[1].narrative.is_some() && f.nodes[2].narrative.is_none());

        let graph = g(&[
            (0, 1, N, 1.0),
            (0, 2, N, 1.0),
            (1, 3, A, 1.0),
            (2, 3, A, 1.0),
            (2, 4, A, 1.0),
        ]);
        let r = open(&graph, 2.3, 0.0, &[0.0, 0.0, 0.6]).unwrap();
        assert!(r.nodes.iter().any(|node| node.node_id == id(3)));
        assert!(!r.nodes.iter().any(|node| node.node_id == id(4)));

        let graph = g(&[(0, 1, A, 1.0), (0, 2, A, 0.5)]);
        for (t, random, selected) in [(1.0, 0.7, 2), (0.5, 0.79, 1), (0.25, 0.95, 2)] {
            let r = open(&graph, 0.3, t, &[random]).unwrap();
            assert_eq!(r.nodes[1].node_id, id(selected));
        }

        let graph = g(&[(0, 1, N, 1.0), (0, 2, N, 0.8), (1, 3, A, 1.0)]);
        let r = open(&graph, 2.0, 0.0, &[0.0, 0.0]).unwrap();
        assert!(r.nodes[2].narrative.is_some());

        let graph = g(&[(0, 1, A, 0.5), (0, 2, A, 1.0), (0, 3, A, 1.0)]);
        let r = open(&graph, 0.3, 0.0, &[0.999_999_999_999_999_9]).unwrap();
        assert_eq!(r.nodes[1].node_id, id(3));

        for b in [f64::NAN, f64::INFINITY, -0.1] {
            assert!(open(&graph, b, 0.0, &[]).unwrap_err().contains("budget"));
        }
        for t in [f64::NAN, f64::INFINITY, -0.1] {
            assert!(
                open(&graph, 0.3, t, &[])
                    .unwrap_err()
                    .contains("temperature")
            );
        }
        for random in [f64::NAN, f64::INFINITY, -0.1, 1.0] {
            assert!(
                open(&graph, 0.3, 0.0, &[random])
                    .unwrap_err()
                    .contains("[0, 1)")
            );
        }
        let mut missing = g(&[(0, 9, A, 1.0)]);
        let _ = missing.remove(&id(9));
        assert!(
            open(&missing, 0.3, 0.0, &[0.0])
                .unwrap_err()
                .contains("missing")
        );
    }
}