open-standards/wai/wai-web/src/int_world.rs

//! In-browser **byte-exact interactive world model** — zero JS ML runtime.
//!
//! The same pure-i64 Rust that renders `wai.world.model`'s integer reference world
//! (`wai::int_world`) compiles to wasm32 and produces the IDENTICAL frames here: an
//! action-conditioned next-frame model whose transition + synthesis are integer
//! convs + integer IGDN. So the page can drive the world with the arrow keys, render
//! every frame in pure i64, and re-check each frame's hash against the native
//! reference. And every keystroke is a signed action op whose trace replays to one
//! portable **session hash** ([`world_session_hash`]) — replay-equivalence, a
//! generative interactive world you can put in a receipt.

use wasm_bindgen::prelude::*;

use wai::int_hyper_synth::{decode as synth_decode, IntModel};

/// A stateful integer world model: hold the transition `T` + synthesis `g_s`, carry
/// the latent state `z`, and advance it one action at a time — the live-play path.
#[wasm_bindgen]
pub struct WorldModel {
    t: IntModel,
    gs: IntModel,
    c: usize,
    lat: usize,
    z: Vec<i64>,
    frame_px: usize,
}

#[wasm_bindgen]
impl WorldModel {
    /// Load the two WHM1 models (`world_T.bin`, `world_gs.bin`); `c`/`lat` from
    /// `world.json`.
    #[wasm_bindgen(constructor)]
    pub fn new(t_bin: &[u8], gs_bin: &[u8], c: u32, lat: u32) -> Result<WorldModel, JsError> {
        Ok(WorldModel {
            t: IntModel::from_bytes(t_bin).map_err(JsError::new)?,
            gs: IntModel::from_bytes(gs_bin).map_err(JsError::new)?,
            c: c as usize,
            lat: lat as usize,
            z: Vec::new(),
            frame_px: 0,
        })
    }

    /// Reset the world to an initial latent `z0` (`c·lat·lat` ints at `2^afrac`);
    /// returns the rendered frame 0 (HWC RGB `u8`).
    pub fn reset(&mut self, z0: &[i32]) -> Box<[u8]> {
        self.z = z0.iter().map(|&v| v as i64).collect();
        let (f, h, _) = synth_decode(&self.gs, &self.z, self.c, self.lat, self.lat);
        self.frame_px = h;
        f.into_boxed_slice()
    }

    /// Advance the integer transition by one `action` (index into the reference
    /// action set: 0 stay, 1 +x, 2 −x, 3 +y, 4 −y) and render — pure i64, byte-exact.
    pub fn step(&mut self, action: u32) -> Box<[u8]> {
        self.z = wai::int_world::step(&self.t, &self.z, action as usize, self.c, self.lat);
        let (f, _, _) = synth_decode(&self.gs, &self.z, self.c, self.lat, self.lat);
        f.into_boxed_slice()
    }

    /// Advance the transition by raw `(dx, dy)` deltas and render — the shared-world
    /// path (a merged multi-actor log decodes to a delta sequence via
    /// [`replay_world_ops`]).
    #[wasm_bindgen(js_name = stepDelta)]
    pub fn step_delta(&mut self, dx: i32, dy: i32) -> Box<[u8]> {
        self.z = wai::int_world::step_delta(&self.t, &self.z, dx as i64, dy as i64, self.c, self.lat);
        let (f, _, _) = synth_decode(&self.gs, &self.z, self.c, self.lat, self.lat);
        f.into_boxed_slice()
    }

    /// Advance the transition by an explicit set of action **planes** (length = the
    /// model's `A`; values at `2^afrac`) and render — the multi-agent path, where a
    /// step carries every agent's action (e.g. `[dxA,dyA,dxB,dyB]` for two agents).
    #[wasm_bindgen(js_name = stepPlanes)]
    pub fn step_planes(&mut self, planes: &[i32]) -> Box<[u8]> {
        let q: Vec<i64> = planes.iter().map(|&v| v as i64).collect();
        self.z = wai::int_world::step_planes(&self.t, &self.z, &q, self.c, self.lat);
        let (f, _, _) = synth_decode(&self.gs, &self.z, self.c, self.lat, self.lat);
        f.into_boxed_slice()
    }

    #[wasm_bindgen(getter, js_name = framePx)]
    pub fn frame_px(&self) -> u32 {
        self.frame_px as u32
    }
}

/// Replay a signed action trace → the portable **session hash** (hex), in the
/// browser. `actions_json` is `[{lamport,tick,dx,dy,buttons}]`; each is signed with
/// the actor key derived from `key_seed_hex` (32-byte seed), then replayed under
/// canonical total order — the same `wai.world.model` replay-equivalence the native
/// reference computes. Lets the page show that the keystrokes the user just made
/// hash to one portable, signed session identity.
#[wasm_bindgen(js_name = worldSessionHash)]
pub fn world_session_hash(
    model_id_hex: &str, seed: u64, width: u32, height: u32, fps: u32, key_seed_hex: &str, actions_json: &[u8],
) -> Result<String, JsError> {
    use ed25519_dalek::SigningKey;
    use wai::world_model::{hex32, replay, sign_action, ModelHeader};

    let model_id = hex32_to_bytes(model_id_hex).ok_or_else(|| JsError::new("bad model_id hex"))?;
    let seed_bytes = hex32_to_bytes(key_seed_hex).ok_or_else(|| JsError::new("bad key seed hex"))?;
    let key = SigningKey::from_bytes(&seed_bytes);
    let header = ModelHeader { model_id, seed, width, height, fps };

    let acts: serde_json::Value =
        serde_json::from_slice(actions_json).map_err(|e| JsError::new(&e.to_string()))?;
    let mut ops = Vec::new();
    for a in acts.as_array().ok_or_else(|| JsError::new("actions must be an array"))? {
        ops.push(sign_action(
            &key,
            a["lamport"].as_u64().unwrap_or(0),
            a["tick"].as_u64().unwrap_or(0),
            a["dx"].as_i64().unwrap_or(0) as i32,
            a["dy"].as_i64().unwrap_or(0) as i32,
            a["buttons"].as_u64().unwrap_or(0) as u32,
            b"",
        ));
    }
    let session = replay(&header, &ops, true).map_err(JsError::new)?;
    Ok(hex32(&session.session_hash))
}

/// Sign one action op for a SHARED world and return its serialized bytes (a
/// `CausalOp`). Each peer signs with its own actor key (`actor_seed_hex`) and
/// broadcasts these bytes; peers concatenate everyone's ops into one log and
/// [`replay_world_ops`] it. The signature binds the action to the actor — an op can
/// only claim the identity that signed it.
#[wasm_bindgen(js_name = signWorldAction)]
pub fn sign_world_action(actor_seed_hex: &str, lamport: u64, tick: u64, dx: i32, dy: i32, buttons: u32) -> Result<Box<[u8]>, JsError> {
    use ed25519_dalek::SigningKey;
    let seed = hex32_to_bytes(actor_seed_hex).ok_or_else(|| JsError::new("bad actor seed hex"))?;
    let key = SigningKey::from_bytes(&seed);
    let op = wai::world_model::sign_action(&key, lamport, tick, dx, dy, buttons, b"");
    Ok(op.to_bytes().into_boxed_slice())
}

/// The result of replaying a merged multi-actor action log: the convergent session
/// hash + the canonical `(dx, dy)` delta sequence to render.
#[wasm_bindgen]
pub struct WorldReplay {
    session_hash: String,
    deltas: Vec<i32>,
    slots: Vec<u32>,
}

#[wasm_bindgen]
impl WorldReplay {
    #[wasm_bindgen(getter, js_name = sessionHash)]
    pub fn session_hash(&self) -> String {
        self.session_hash.clone()
    }
    /// Flattened `[dx0, dy0, dx1, dy1, …]` in canonical total order.
    #[wasm_bindgen(getter)]
    pub fn deltas(&self) -> Box<[i32]> {
        self.deltas.clone().into_boxed_slice()
    }
    /// Per-op agent SLOT (the op's actor's index in canonical actor order) — so a
    /// peer can place each op's `(dx,dy)` in its agent's plane lane for multi-agent
    /// co-control. Same on every peer (the actor set + ordering is canonical).
    #[wasm_bindgen(getter)]
    pub fn slots(&self) -> Box<[u32]> {
        self.slots.clone().into_boxed_slice()
    }
}

/// Replay a MERGED multi-actor action log (`log` = everyone's `CausalOp`s
/// concatenated, any order) → the convergent session hash + the canonical delta
/// sequence. Every peer that holds the same op set computes the IDENTICAL result —
/// replay-equivalence — so rendering the deltas yields the same byte-exact world on
/// every tab. Signatures are verified.
#[wasm_bindgen(js_name = replayWorldOps)]
pub fn replay_world_ops(
    model_id_hex: &str, seed: u64, width: u32, height: u32, fps: u32, log: &[u8],
) -> Result<WorldReplay, JsError> {
    use wai::world::CausalOp;
    use wai::world_model::{hex32, replay, ModelHeader};

    let model_id = hex32_to_bytes(model_id_hex).ok_or_else(|| JsError::new("bad model_id hex"))?;
    let ops = if log.is_empty() {
        Vec::new()
    } else {
        CausalOp::parse_log(log).map_err(|_| JsError::new("malformed action log"))?
    };
    let header = ModelHeader { model_id, seed, width, height, fps };
    let session = replay(&header, &ops, true).map_err(JsError::new)?;
    let deltas = session.actions.iter().flat_map(|a| [a.dx, a.dy]).collect();
    // Assign each distinct actor a stable slot by canonical actor order, so every
    // peer maps the same op to the same agent lane.
    let mut actors: Vec<[u8; 32]> = session.actions.iter().map(|a| a.actor).collect();
    actors.sort_unstable();
    actors.dedup();
    let slots = session.actions.iter().map(|a| actors.iter().position(|x| x == &a.actor).unwrap() as u32).collect();
    Ok(WorldReplay { session_hash: hex32(&session.session_hash), deltas, slots })
}

fn hex32_to_bytes(s: &str) -> Option<[u8; 32]> {
    let s = s.as_bytes();
    if s.len() != 64 {
        return None;
    }
    let mut o = [0u8; 32];
    for (i, b) in o.iter_mut().enumerate() {
        let hi = (s[i * 2] as char).to_digit(16)?;
        let lo = (s[i * 2 + 1] as char).to_digit(16)?;
        *b = (hi * 16 + lo) as u8;
    }
    Some(o)
}