kcode-k1-kmap-loader 0.2.0

Stateless access-filtered attention loader for K1 Kmaps
Documentation
use std::collections::HashMap;

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

mod selection;

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,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum OpenMode {
    Full,
    NavigationOnly,
}

pub fn open_node<L, A>(
    node_id: NodeId,
    budget: f64,
    temperature: f64,
    mode: OpenMode,
    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,
        mode,
        load_node,
        access_filter,
        selection::os_random_unit,
    )
}

fn open_node_with_random<L, A, R>(
    node_id: NodeId,
    budget: f64,
    temperature: f64,
    mode: OpenMode,
    load_node: L,
    access_filter: A,
    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 mut engine = Engine {
        budget,
        temperature,
        load_node,
        access_filter,
        random,
        decisions: HashMap::new(),
    };
    let root = engine.required(node_id)?;
    match mode {
        OpenMode::Full => engine.full(node_id, root),
        OpenMode::NavigationOnly => engine.navigation_only(node_id, root),
    }
}

struct Engine<L, A, R> {
    budget: f64,
    temperature: f64,
    load_node: L,
    access_filter: A,
    random: R,
    decisions: HashMap<NodeId, bool>,
}

impl<L, A, R> Engine<L, A, R>
where
    L: FnMut(NodeId) -> Result<Option<Node>, String>,
    A: FnMut(NodeId) -> bool,
    R: FnMut() -> Result<f64, String>,
{
    fn full(&mut self, node_id: NodeId, root: Node) -> Result<OpenResult, String> {
        let mut outputs = vec![loaded(node_id, &root, true)];
        let mut states = HashMap::from([(node_id, NodeState::Opened)]);
        let mut frontier = Vec::new();
        let mut spent = 0_u64;

        for connection in &root.connections {
            let target = connection.target;
            if matches!(states.get(&target), Some(NodeState::Opened)) || !self.allowed(target) {
                continue;
            }
            if connection.tier == ConnectionTier::Navigation && !states.contains_key(&target) {
                let node = self.required(target)?;
                insert_preview(target, node, &mut outputs, &mut states);
            }
            frontier.push(Occurrence::new(target, connection.weight.value, 1));
        }

        loop {
            let mut candidates = Vec::new();
            let mut opening_costs = HashMap::new();
            for (index, occurrence) in frontier.iter().enumerate() {
                if matches!(states.get(&occurrence.target), Some(NodeState::Opened)) {
                    continue;
                }
                let effective = selection::score(occurrence);
                if !effective.is_finite() || effective <= 0.0 {
                    continue;
                }
                let cost = match states.get(&occurrence.target) {
                    None => PREVIEW_TENTHS,
                    Some(NodeState::Previewed(node)) => {
                        if let Some(cost) = opening_costs.get(&occurrence.target) {
                            *cost
                        } else {
                            let previews = self.navigation_targets(node, &states, None).len();
                            let cost = opening_cost(previews)?;
                            let _ = opening_costs.insert(occurrence.target, cost);
                            cost
                        }
                    }
                    Some(NodeState::Opened) => continue,
                };
                if affordable(spent, cost, self.budget)? {
                    candidates.push((index, effective, cost));
                }
            }
            if candidates.is_empty() {
                break;
            }

            let choice = selection::choose(&candidates, self.temperature, &mut self.random)?;
            let (selected, _, cost) = candidates[choice];
            let occurrence = frontier[selected].clone();
            match states.get(&occurrence.target).cloned() {
                None => {
                    let node = self.required(occurrence.target)?;
                    insert_preview(occurrence.target, node, &mut outputs, &mut states);
                }
                Some(NodeState::Previewed(node)) => {
                    let guarantees = self.navigation_targets(&node, &states, None);
                    let _ = states.insert(occurrence.target, NodeState::Opened);
                    let output = outputs
                        .iter_mut()
                        .find(|node| node.node_id == occurrence.target)
                        .ok_or_else(|| "previewed Kmap node had no output".to_owned())?;
                    output.narrative = Some(node.narrative.clone());
                    frontier.retain(|entry| entry.target != occurrence.target);
                    for target in guarantees {
                        let guaranteed = self.required(target)?;
                        insert_preview(target, guaranteed, &mut outputs, &mut states);
                    }
                    let depth = increment_depth(occurrence.depth)?;
                    frontier.extend(self.outgoing(&node, depth, &states));
                }
                Some(NodeState::Opened) => {
                    return Err("opened Kmap node remained selectable".to_owned());
                }
            }
            spent = attention(spent.checked_add(cost))?;
        }
        Ok(result(outputs, spent))
    }

    fn navigation_only(&mut self, node_id: NodeId, root: Node) -> Result<OpenResult, String> {
        let mut outputs = vec![loaded(node_id, &root, false)];
        let mut states = HashMap::from([(node_id, NodeState::Opened)]);
        let root_targets = self.navigation_targets(&root, &states, None);
        let root_cost = preview_cost(root_targets.len())?;
        if !affordable(0, root_cost, self.budget)? {
            return Ok(result(outputs, 0));
        }
        let mut root_nodes = Vec::with_capacity(root_targets.len());
        for target in root_targets {
            root_nodes.push((target, self.required(target)?));
        }
        for (target, node) in &root_nodes {
            outputs.push(loaded(*target, node, false));
            let _ = states.insert(*target, NodeState::Opened);
        }

        let mut spent = root_cost;
        let mut frontier = self.outgoing(&root, 1, &states);
        for (_, node) in &root_nodes {
            frontier.extend(self.outgoing(node, 2, &states));
        }
        loop {
            let mut candidates = Vec::new();
            for (index, occurrence) in frontier.iter().enumerate() {
                if states.contains_key(&occurrence.target) {
                    continue;
                }
                let effective = selection::score(occurrence);
                if effective.is_finite() && effective > 0.0 {
                    candidates.push((index, effective, 0));
                }
            }
            if candidates.is_empty() {
                break;
            }

            let choice = selection::choose(&candidates, self.temperature, &mut self.random)?;
            let (selected, _, _) = candidates[choice];
            let occurrence = frontier[selected].clone();
            let node = self.required(occurrence.target)?;
            let children = self.navigation_targets(&node, &states, Some(occurrence.target));
            let count = attention(children.len().checked_add(1))?;
            let cost = preview_cost(count)?;
            if !affordable(spent, cost, self.budget)? {
                break;
            }
            let next_depth = increment_depth(occurrence.depth)?;
            let new_spent = attention(spent.checked_add(cost))?;
            let mut child_nodes = Vec::with_capacity(children.len());
            for target in children {
                child_nodes.push((target, self.required(target)?));
            }

            let _ = states.insert(occurrence.target, NodeState::Opened);
            for (target, _) in &child_nodes {
                let _ = states.insert(*target, NodeState::Opened);
            }
            let mut additions = self.outgoing(&node, next_depth, &states);
            for (_, child) in &child_nodes {
                let child_depth = increment_depth(next_depth)?;
                additions.extend(self.outgoing(child, child_depth, &states));
            }
            outputs.push(loaded(occurrence.target, &node, false));
            for (target, child) in &child_nodes {
                outputs.push(loaded(*target, child, false));
            }
            frontier.retain(|entry| !states.contains_key(&entry.target));
            frontier.extend(additions);
            spent = new_spent;
        }
        Ok(result(outputs, spent))
    }

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

    fn allowed(&mut self, node_id: NodeId) -> bool {
        let access_filter = &mut self.access_filter;
        *self
            .decisions
            .entry(node_id)
            .or_insert_with(|| access_filter(node_id))
    }

    fn navigation_targets(
        &mut self,
        node: &Node,
        states: &HashMap<NodeId, NodeState>,
        planned: Option<NodeId>,
    ) -> Vec<NodeId> {
        let mut targets = Vec::new();
        for connection in &node.connections {
            let target = connection.target;
            if states.contains_key(&target) || planned == Some(target) || !self.allowed(target) {
                continue;
            }
            if connection.tier == ConnectionTier::Navigation && !targets.contains(&target) {
                targets.push(target);
            }
        }
        targets
    }

    fn outgoing(
        &mut self,
        node: &Node,
        depth: usize,
        states: &HashMap<NodeId, NodeState>,
    ) -> Vec<Occurrence> {
        let mut entries = Vec::new();
        for connection in &node.connections {
            let target = connection.target;
            if matches!(states.get(&target), Some(NodeState::Opened)) || !self.allowed(target) {
                continue;
            }
            entries.push(Occurrence::new(target, connection.weight.value, depth));
        }
        entries
    }
}

#[derive(Clone)]
enum NodeState {
    Previewed(Node),
    Opened,
}

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

impl Occurrence {
    fn new(target: NodeId, value: f64, depth: usize) -> Self {
        Self {
            target,
            value,
            depth,
        }
    }
}

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>,
    states: &mut HashMap<NodeId, NodeState>,
) {
    if states.contains_key(&node_id) {
        return;
    }
    outputs.push(loaded(node_id, &node, false));
    let _ = states.insert(node_id, NodeState::Previewed(node));
}

fn result(nodes: Vec<LoadedNode>, spent: u64) -> OpenResult {
    OpenResult {
        nodes,
        automatic_attention_spent: spent as f64 / 10.0,
    }
}

fn opening_cost(previews: usize) -> Result<u64, String> {
    attention(preview_cost(previews)?.checked_add(NARRATIVE_TENTHS))
}

fn increment_depth(depth: usize) -> Result<usize, String> {
    depth
        .checked_add(1)
        .ok_or_else(|| "Kmap traversal depth overflow".to_owned())
}

fn attention<T>(value: Option<T>) -> Result<T, String> {
    value.ok_or_else(|| "Kmap attention cost overflow".to_owned())
}

fn preview_cost(count: usize) -> Result<u64, String> {
    let count = attention(u64::try_from(count).ok())?;
    attention(count.checked_mul(PREVIEW_TENTHS))
}

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