453 lines
16 KiB
Rust
453 lines
16 KiB
Rust
//! Checkpoint envelope: a JSON manifest plus named binary chunks, checksummed.
|
|
//!
|
|
//! Ports `agent/envelope.ts`. The layout is frozen (files written by earlier versions of this
|
|
//! format must keep loading):
|
|
//!
|
|
//! ```text
|
|
//! magic ASCII, caller-chosen, identifies the file kind
|
|
//! manifestLength u32 little-endian
|
|
//! manifest UTF-8 JSON, includes schemaVersion and the chunk name list
|
|
//! [ length, bytes ] * n u32 little-endian length then the chunk, in manifest order
|
|
//! checksum u32 little-endian CRC32 over every preceding byte
|
|
//! ```
|
|
|
|
use crate::agent::AgentState;
|
|
use crate::decoder::DecoderState;
|
|
use crate::error::{bail, Error, Result};
|
|
use crate::json::JsonValue;
|
|
use crate::lif::LifState;
|
|
use crate::ordered::NumberMap;
|
|
use crate::plasticity::PlasticityState;
|
|
|
|
/// Envelope schema; bumped only for a layout change, not for manifest fields.
|
|
pub const ENVELOPE_SCHEMA_VERSION: f64 = 2.0;
|
|
|
|
/// CRC32 (IEEE 802.3, reflected polynomial 0xedb88320) over `bytes`.
|
|
///
|
|
/// Detects accidental corruption of the manifest and every payload byte. Not a signature.
|
|
pub fn checksum(bytes: &[u8]) -> u32 {
|
|
let mut crc: u32 = 0xffff_ffff;
|
|
for byte in bytes {
|
|
crc ^= u32::from(*byte);
|
|
for _ in 0..8 {
|
|
crc = (crc >> 1) ^ (0xedb8_8320 & (!(crc & 1)).wrapping_add(1));
|
|
}
|
|
}
|
|
crc ^ 0xffff_ffff
|
|
}
|
|
|
|
/// A decoded envelope.
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct EnvelopeParts {
|
|
pub manifest: JsonValue,
|
|
/// Chunks in manifest order.
|
|
pub chunks: Vec<(String, Vec<u8>)>,
|
|
}
|
|
|
|
impl EnvelopeParts {
|
|
pub fn chunk(&self, name: &str) -> Option<&[u8]> {
|
|
self.chunks
|
|
.iter()
|
|
.find(|(key, _)| key == name)
|
|
.map(|(_, bytes)| bytes.as_slice())
|
|
}
|
|
}
|
|
|
|
/// Chunk names are restricted to letters, so a manifest can never name a prototype key.
|
|
fn is_chunk_name(name: &str) -> bool {
|
|
!name.is_empty() && name.chars().all(|c| c.is_ascii_alphabetic())
|
|
}
|
|
|
|
/// Encode `manifest` and `chunks` into one checksummed buffer.
|
|
///
|
|
/// `magic` must be ASCII; the chunk name list and `schemaVersion` are written into the manifest,
|
|
/// overriding any same-named fields.
|
|
pub fn encode_envelope(
|
|
magic: &str,
|
|
manifest: &JsonValue,
|
|
chunks: &[(String, Vec<u8>)],
|
|
) -> Result<Vec<u8>> {
|
|
if magic.is_empty() || !magic.is_ascii() {
|
|
bail!("Envelope magic must be a non-empty ASCII string");
|
|
}
|
|
let names: Vec<String> = chunks.iter().map(|(name, _)| name.clone()).collect();
|
|
if names.iter().any(|name| !is_chunk_name(name)) {
|
|
bail!("Invalid envelope chunk name");
|
|
}
|
|
let mut manifest = manifest.clone();
|
|
if !matches!(manifest, JsonValue::Object(_)) {
|
|
manifest = JsonValue::object();
|
|
}
|
|
manifest.set("schemaVersion", ENVELOPE_SCHEMA_VERSION.into());
|
|
manifest.set(
|
|
"chunks",
|
|
JsonValue::Array(names.iter().map(|name| name.as_str().into()).collect()),
|
|
);
|
|
let manifest_bytes = manifest.stringify().into_bytes();
|
|
|
|
let mut out = Vec::new();
|
|
out.extend_from_slice(magic.as_bytes());
|
|
out.extend_from_slice(&(manifest_bytes.len() as u32).to_le_bytes());
|
|
out.extend_from_slice(&manifest_bytes);
|
|
for (_, bytes) in chunks {
|
|
out.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
|
|
out.extend_from_slice(bytes);
|
|
}
|
|
let crc = checksum(&out);
|
|
out.extend_from_slice(&crc.to_le_bytes());
|
|
Ok(out)
|
|
}
|
|
|
|
/// Decode a buffer written by [`encode_envelope`], rejecting anything that is not structurally
|
|
/// this format with `magic`.
|
|
pub fn decode_envelope(bytes: &[u8], magic: &str) -> Result<EnvelopeParts> {
|
|
let magic_bytes = magic.as_bytes();
|
|
if bytes.len() < magic_bytes.len() + 4 || !bytes.starts_with(magic_bytes) {
|
|
bail!("Not a {magic} envelope");
|
|
}
|
|
let mut offset = magic_bytes.len();
|
|
let manifest_length = read_u32(bytes, offset) as usize;
|
|
offset += 4;
|
|
if offset + manifest_length > bytes.len() {
|
|
bail!("Envelope manifest is truncated");
|
|
}
|
|
let manifest_text = std::str::from_utf8(&bytes[offset..offset + manifest_length])
|
|
.map_err(|_| Error::new("Envelope manifest is not valid UTF-8"))?;
|
|
let manifest = JsonValue::parse(manifest_text)?;
|
|
offset += manifest_length;
|
|
let schema = manifest.get("schemaVersion").and_then(JsonValue::as_f64);
|
|
if schema != Some(ENVELOPE_SCHEMA_VERSION) {
|
|
bail!(
|
|
"Unsupported envelope schema: {}",
|
|
schema
|
|
.map(crate::jsmath::number_to_string)
|
|
.unwrap_or_else(|| "undefined".to_string())
|
|
);
|
|
}
|
|
if bytes.len() < 4 || checksum(&bytes[..bytes.len() - 4]) != read_u32(bytes, bytes.len() - 4) {
|
|
bail!("Envelope checksum mismatch");
|
|
}
|
|
let names = manifest
|
|
.get("chunks")
|
|
.and_then(JsonValue::as_array)
|
|
.map(<[JsonValue]>::to_vec);
|
|
let Some(names) = names else {
|
|
bail!("Invalid envelope chunks");
|
|
};
|
|
let mut chunk_names = Vec::with_capacity(names.len());
|
|
for name in &names {
|
|
match name.as_str() {
|
|
Some(name) if is_chunk_name(name) => chunk_names.push(name.to_string()),
|
|
_ => bail!("Invalid envelope chunks"),
|
|
}
|
|
}
|
|
let unique: std::collections::HashSet<&String> = chunk_names.iter().collect();
|
|
if unique.len() != chunk_names.len() {
|
|
bail!("Invalid envelope chunks");
|
|
}
|
|
|
|
let mut chunks = Vec::with_capacity(chunk_names.len());
|
|
for name in chunk_names {
|
|
if offset + 4 > bytes.len() {
|
|
bail!("Envelope chunk header is truncated");
|
|
}
|
|
let length = read_u32(bytes, offset) as usize;
|
|
offset += 4;
|
|
if offset + length > bytes.len() {
|
|
bail!("Envelope chunk {name} is truncated");
|
|
}
|
|
chunks.push((name, bytes[offset..offset + length].to_vec()));
|
|
offset += length;
|
|
}
|
|
if offset != bytes.len() - 4 {
|
|
bail!("Envelope contains trailing data");
|
|
}
|
|
Ok(EnvelopeParts { manifest, chunks })
|
|
}
|
|
|
|
fn read_u32(bytes: &[u8], offset: usize) -> u32 {
|
|
u32::from_le_bytes([
|
|
bytes[offset],
|
|
bytes[offset + 1],
|
|
bytes[offset + 2],
|
|
bytes[offset + 3],
|
|
])
|
|
}
|
|
|
|
/// Chunk names of the typed arrays in an agent checkpoint, in write order.
|
|
pub const AGENT_CHUNK_NAMES: [&str; 7] = [
|
|
"membrane",
|
|
"refractory",
|
|
"lastSpikeMs",
|
|
"visualDrive",
|
|
"plasticGains",
|
|
"plasticTraces",
|
|
"plasticTouched",
|
|
];
|
|
|
|
/// The scalar half of an [`AgentState`] plus its chunks, ready for [`encode_envelope`].
|
|
pub struct AgentChunks {
|
|
pub manifest: JsonValue,
|
|
pub chunks: Vec<(String, Vec<u8>)>,
|
|
}
|
|
|
|
/// Split an agent state into the manifest and chunks [`encode_envelope`] takes.
|
|
///
|
|
/// Chunk names and manifest key order are the original checkpoint's, so a host that already
|
|
/// writes those names keeps its file format.
|
|
pub fn agent_to_chunks(state: &AgentState) -> AgentChunks {
|
|
let network = &state.network;
|
|
let mut manifest = JsonValue::object();
|
|
manifest.set("agentVersion", 1.0.into());
|
|
manifest.set("remainder", state.remainder.into());
|
|
manifest.set("warmedUp", state.warmed_up.into());
|
|
|
|
let mut network_json = JsonValue::object();
|
|
network_json.set("rng", network.rng.into());
|
|
network_json.set("rewardRemaining", network.reward_remaining.into());
|
|
network_json.set("ms", network.ms.into());
|
|
network_json.set("populationRate", network.population_rate.into());
|
|
network_json.set("rates", network.rates.to_json());
|
|
manifest.set("network", network_json);
|
|
|
|
manifest.set("decoder", decoder_state_to_json(&state.decoder));
|
|
|
|
let mut plasticity = JsonValue::object();
|
|
plasticity.set("version", network.plasticity.version.as_str().into());
|
|
plasticity.set("topology", f64::from(network.plasticity.topology).into());
|
|
plasticity.set("enabled", network.plasticity.enabled.into());
|
|
plasticity.set("updates", network.plasticity.updates.into());
|
|
plasticity.set("signal", network.plasticity.signal.into());
|
|
manifest.set("plasticity", plasticity);
|
|
|
|
let chunks = vec![
|
|
("membrane".to_string(), f32_bytes(&network.membrane)),
|
|
("refractory".to_string(), network.refractory.clone()),
|
|
("lastSpikeMs".to_string(), f64_bytes(&network.last_spike_ms)),
|
|
("visualDrive".to_string(), f32_bytes(&network.visual_drive)),
|
|
(
|
|
"plasticGains".to_string(),
|
|
f32_bytes(&network.plasticity.gains),
|
|
),
|
|
(
|
|
"plasticTraces".to_string(),
|
|
f32_bytes(&network.plasticity.traces),
|
|
),
|
|
(
|
|
"plasticTouched".to_string(),
|
|
f64_bytes(&network.plasticity.touched),
|
|
),
|
|
];
|
|
AgentChunks { manifest, chunks }
|
|
}
|
|
|
|
/// Rebuild an agent state from a decoded envelope.
|
|
///
|
|
/// Only presence, alignment and the manifest shape are checked here; the values themselves are
|
|
/// validated by [`crate::agent::NeuralAgent::import_state`], which is also what makes a rejected
|
|
/// checkpoint a no-op.
|
|
pub fn agent_from_chunks(manifest: &JsonValue, parts: &EnvelopeParts) -> Result<AgentState> {
|
|
if manifest.get("agentVersion").and_then(JsonValue::as_f64) != Some(1.0) {
|
|
bail!("Unsupported agent checkpoint version");
|
|
}
|
|
let network = manifest.get("network");
|
|
let decoder = manifest.get("decoder");
|
|
let plasticity = manifest.get("plasticity");
|
|
let (Some(network), Some(decoder), Some(plasticity)) = (network, decoder, plasticity) else {
|
|
bail!("Agent checkpoint manifest is incomplete");
|
|
};
|
|
if network.is_null() || decoder.is_null() || plasticity.is_null() {
|
|
bail!("Agent checkpoint manifest is incomplete");
|
|
}
|
|
for name in AGENT_CHUNK_NAMES {
|
|
if parts.chunk(name).is_none() {
|
|
bail!("Checkpoint is missing {name}");
|
|
}
|
|
}
|
|
let number = |value: &JsonValue, field: &str| -> Result<f64> {
|
|
value
|
|
.get(field)
|
|
.and_then(JsonValue::as_f64)
|
|
.ok_or_else(|| Error::new(format!("Agent checkpoint manifest is incomplete: {field}")))
|
|
};
|
|
|
|
let plasticity_state = PlasticityState {
|
|
version: plasticity
|
|
.get("version")
|
|
.and_then(JsonValue::as_str)
|
|
.unwrap_or_default()
|
|
.to_string(),
|
|
topology: number(plasticity, "topology")? as u32,
|
|
enabled: plasticity
|
|
.get("enabled")
|
|
.and_then(JsonValue::as_bool)
|
|
.unwrap_or(false),
|
|
updates: number(plasticity, "updates")?,
|
|
signal: number(plasticity, "signal")?,
|
|
gains: read_f32(parts, "plasticGains")?,
|
|
traces: read_f32(parts, "plasticTraces")?,
|
|
touched: read_f64(parts, "plasticTouched")?,
|
|
};
|
|
|
|
Ok(AgentState {
|
|
version: 1,
|
|
remainder: number(manifest, "remainder")?,
|
|
warmed_up: manifest
|
|
.get("warmedUp")
|
|
.and_then(JsonValue::as_bool)
|
|
.unwrap_or(false),
|
|
network: LifState {
|
|
membrane: read_f32(parts, "membrane")?,
|
|
refractory: parts.chunk("refractory").unwrap_or_default().to_vec(),
|
|
last_spike_ms: read_f64(parts, "lastSpikeMs")?,
|
|
visual_drive: read_f32(parts, "visualDrive")?,
|
|
rng: crate::rng::Xorshift32::to_int32(number(network, "rng")?),
|
|
reward_remaining: number(network, "rewardRemaining")?,
|
|
ms: number(network, "ms")?,
|
|
population_rate: number(network, "populationRate")?,
|
|
rates: network
|
|
.get("rates")
|
|
.and_then(NumberMap::from_json)
|
|
.unwrap_or_default(),
|
|
plasticity: plasticity_state,
|
|
},
|
|
decoder: decoder_state_from_json(decoder)?,
|
|
})
|
|
}
|
|
|
|
/// A [`DecoderState`] as the manifest writes it.
|
|
pub fn decoder_state_to_json(state: &DecoderState) -> JsonValue {
|
|
let mut json = JsonValue::object();
|
|
json.set("version", f64::from(state.version).into());
|
|
json.set("calibrated", state.calibrated.into());
|
|
json.set("baseline", state.baseline.to_json());
|
|
json.set("heldUntil", state.held_until.to_json());
|
|
json.set("nextAllowed", state.next_allowed.to_json());
|
|
json.set("nextDecision", state.next_decision.into());
|
|
json.set(
|
|
"current",
|
|
match &state.current {
|
|
Some(channel) => channel.as_str().into(),
|
|
None => JsonValue::Null,
|
|
},
|
|
);
|
|
json.set("fatigue", state.fatigue.to_json());
|
|
// The macro group (`docs/design/macros.md` section 11). Always written, never required on the
|
|
// way back in: a manifest from before the group carries none of these three and the group
|
|
// then starts rested, which is what keeps the live checkpoint loadable across the change.
|
|
json.set("macroNextDecision", state.macro_next_decision.into());
|
|
json.set(
|
|
"macroCurrent",
|
|
match &state.macro_current {
|
|
Some(channel) => channel.as_str().into(),
|
|
None => JsonValue::Null,
|
|
},
|
|
);
|
|
json.set("macroFatigue", state.macro_fatigue.to_json());
|
|
json
|
|
}
|
|
|
|
/// Read a [`DecoderState`] back out of a manifest.
|
|
pub fn decoder_state_from_json(value: &JsonValue) -> Result<DecoderState> {
|
|
let map = |field: &str| -> Result<NumberMap> {
|
|
value
|
|
.get(field)
|
|
.and_then(NumberMap::from_json)
|
|
.ok_or_else(|| Error::new("Invalid decoder checkpoint"))
|
|
};
|
|
Ok(DecoderState {
|
|
version: value
|
|
.get("version")
|
|
.and_then(JsonValue::as_f64)
|
|
.unwrap_or(f64::NAN) as u32,
|
|
calibrated: value
|
|
.get("calibrated")
|
|
.and_then(JsonValue::as_bool)
|
|
.ok_or_else(|| Error::new("Invalid decoder checkpoint"))?,
|
|
baseline: map("baseline")?,
|
|
held_until: map("heldUntil")?,
|
|
next_allowed: map("nextAllowed")?,
|
|
next_decision: value
|
|
.get("nextDecision")
|
|
.and_then(JsonValue::as_f64)
|
|
.ok_or_else(|| Error::new("Invalid decoder checkpoint"))?,
|
|
current: match value.get("current") {
|
|
Some(JsonValue::String(channel)) => Some(channel.clone()),
|
|
_ => None,
|
|
},
|
|
fatigue: map("fatigue")?,
|
|
macro_next_decision: value
|
|
.get("macroNextDecision")
|
|
.and_then(JsonValue::as_f64)
|
|
.unwrap_or(0.0),
|
|
macro_current: match value.get("macroCurrent") {
|
|
Some(JsonValue::String(channel)) => Some(channel.clone()),
|
|
_ => None,
|
|
},
|
|
macro_fatigue: value
|
|
.get("macroFatigue")
|
|
.and_then(NumberMap::from_json)
|
|
.unwrap_or_default(),
|
|
})
|
|
}
|
|
|
|
/// `Zero-offset copy, so reading back an f32 array is aligned and exactly the right length.`
|
|
fn read_f32(parts: &EnvelopeParts, name: &str) -> Result<Vec<f32>> {
|
|
let bytes = parts
|
|
.chunk(name)
|
|
.ok_or_else(|| Error::new(format!("Checkpoint is missing {name}")))?;
|
|
if bytes.len() % 4 != 0 {
|
|
bail!("Checkpoint chunk {name} has a partial element");
|
|
}
|
|
Ok(bytes
|
|
.chunks_exact(4)
|
|
.map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
|
|
.collect())
|
|
}
|
|
|
|
fn read_f64(parts: &EnvelopeParts, name: &str) -> Result<Vec<f64>> {
|
|
let bytes = parts
|
|
.chunk(name)
|
|
.ok_or_else(|| Error::new(format!("Checkpoint is missing {name}")))?;
|
|
if bytes.len() % 8 != 0 {
|
|
bail!("Checkpoint chunk {name} has a partial element");
|
|
}
|
|
Ok(bytes
|
|
.chunks_exact(8)
|
|
.map(|chunk| {
|
|
f64::from_le_bytes([
|
|
chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7],
|
|
])
|
|
})
|
|
.collect())
|
|
}
|
|
|
|
/// Raw little-endian bytes of an f32 array, as `Float32Array` holds them.
|
|
pub fn f32_bytes(values: &[f32]) -> Vec<u8> {
|
|
let mut out = Vec::with_capacity(values.len() * 4);
|
|
for value in values {
|
|
out.extend_from_slice(&value.to_le_bytes());
|
|
}
|
|
out
|
|
}
|
|
|
|
/// Raw little-endian bytes of an f64 array, as `Float64Array` holds them.
|
|
pub fn f64_bytes(values: &[f64]) -> Vec<u8> {
|
|
let mut out = Vec::with_capacity(values.len() * 8);
|
|
for value in values {
|
|
out.extend_from_slice(&value.to_le_bytes());
|
|
}
|
|
out
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn crc32_matches_the_reference() {
|
|
assert_eq!(checksum(b""), 0);
|
|
assert_eq!(checksum(b"a"), 0xe8b7_be43);
|
|
assert_eq!(checksum(b"123456789"), 0xcbf4_3926);
|
|
}
|
|
}
|