cubelinux-2/cubecode/src/cell.rs

//! Code/data mapping: how functions, modules, kernels, layers and
//! checkpoints are represented as CUBELinux records (PDF Package 4).
//!
//! The PDF's intent: "Treat each function, module, or model component as a
//! record: header: title, type ('fn', 'kernel', 'layer', 'checkpoint'),
//! created/last updated, owner, etc. Use association flags and CZYX links for
//! call graphs, dependency edges, variant implementations, version history."
//!
//! We make that concrete: a [`CodeCell`] is the typed bundle a record's
//! header *should* carry (we store it as `Option<String>` doc_type + a typed
//! `Kind` enum in the title-space so the existing `CubeHeader` wire codec is
//! reused without a new serialize format). The VM reads the body as bytecode
//! and resolves `CALL_LINK n` to `linked_records[n]`.

use cubecoords::{CubeHeader, Czyx};

/// The kind of code artifact a record holds.
///
/// These are the PDF's example `type`s. They are stored as the record's
/// `doc_type` string so the binary header wire format (cubestore) carries
/// them for free, and so cubefs can surface them as a file's "type".
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Kind {
    /// A plain function: a self-contained bytecode routine.
    Fn,
    /// A kernel: a routine the host treats as a primitive building block
    /// (e.g. a vector op offloaded to `syscall`).
    Kernel,
    /// A model layer: weights/params for an AI component, addressable as a
    /// record so other layers link to it.
    Layer,
    /// A model checkpoint: a snapshot of weights at a training step.
    Checkpoint,
    /// A variant: an alternate implementation of another routine.
    Variant,
    /// An unknown/unspecified artifact kind.
    Other,
}

impl Kind {
    /// The canonical `doc_type` string persisted in the record header.
    pub fn as_str(&self) -> &'static str {
        match self {
            Kind::Fn => "fn",
            Kind::Kernel => "kernel",
            Kind::Layer => "layer",
            Kind::Checkpoint => "checkpoint",
            Kind::Variant => "variant",
            Kind::Other => "other",
        }
    }

    /// Parse a `doc_type` string back into a [`Kind`].
    pub fn from_doc_type(s: &str) -> Kind {
        match s {
            "fn" => Kind::Fn,
            "kernel" => Kind::Kernel,
            "layer" => Kind::Layer,
            "checkpoint" => Kind::Checkpoint,
            "variant" => Kind::Variant,
            _ => Kind::Other,
        }
    }
}

/// A typed, named code cell mapped onto a cube record.
///
/// Decision: rather than invent a new serialization, a `CodeCell` is built
/// from (and flattened into) the existing [`CubeHeader`] + body bytes: the
/// `Kind` rides in `doc_type`, the human name in `title`, and the bytecode
/// in the body. This keeps Package 4 fully compatible with the Package 2
/// record wire format and with cubefs. `code_kind`/`set_kind` bridge the
/// enum to the string.
#[derive(Clone, Debug)]
pub struct CodeCell {
    /// The cube coordinate this cell is stored at.
    pub label: Czyx,
    /// Typed header.
    pub header: CubeHeader,
    /// Decoded bytecode (re-derived from the body on construction).
    pub code: Vec<crate::opcode::Op>,
}

impl CodeCell {
    /// Build a code cell from a header + raw body, deriving the decoded
    /// bytecode. Returns `None` if the body is not valid bytecode.
    pub fn from_record(label: Czyx, header: &CubeHeader, body: &[u8]) -> Option<CodeCell> {
        let code = crate::opcode::decode(body).ok()?;
        // Carry a clone of the header for callers that want the metadata
        // (owner, kind, timestamps) without re-fetching.
        let mut header = header.clone();
        header.size_bytes = Some(body.len() as u64);
        header.refresh_flags();
        Some(CodeCell {
            label,
            header,
            code,
        })
    }

    /// The artifact kind, read from `doc_type`.
    pub fn kind(&self) -> Kind {
        Kind::from_doc_type(self.header.doc_type.as_deref().unwrap_or("other"))
    }

    /// The human title, if any.
    pub fn name(&self) -> Option<&str> {
        self.header.title.as_deref()
    }

    /// Indices of records this cell links to (its call-graph / dependency
    /// edges). `CALL_LINK n` addresses `linked_records[n]`.
    pub fn links(&self) -> &[Czyx] {
        &self.header.linked_records
    }

    /// Serialize for storage: returns the body bytes (re-encoded from the
    /// decoded `code` so a `CodeCell` constructed in memory round-trips
    /// exactly through the codec, not through a stale original buffer).
    pub fn body(&self) -> Vec<u8> {
        crate::opcode::encode(&self.code)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cubecoords::Czyx;

    fn cell(kind: Kind, links: usize) -> CodeCell {
        let mut h = CubeHeader::new();
        h.title = Some("adder".into());
        h.doc_type = Some(kind.as_str().into());
        let link_coords: Vec<Czyx> = (0..links)
            .map(|i| Czyx::new(7, 0, 0, i as u8 + 1))
            .collect();
        h.linked_records = link_coords;
        h.refresh_flags();
        let code = vec![crate::opcode::Op::Const(1), crate::opcode::Op::Halt];
        CodeCell {
            label: Czyx::new(1, 1, 1, 1),
            header: h,
            code,
        }
    }

    #[test]
    fn kind_roundtrips_through_doc_type() {
        for k in [
            Kind::Fn,
            Kind::Kernel,
            Kind::Layer,
            Kind::Checkpoint,
            Kind::Variant,
        ] {
            assert_eq!(Kind::from_doc_type(k.as_str()), k);
        }
    }

    #[test]
    fn from_record_roundtrips() {
        let c = cell(Kind::Fn, 2);
        let body = c.body();
        let back = CodeCell::from_record(c.label, &c.header, &body).unwrap();
        assert_eq!(back.kind(), Kind::Fn);
        assert_eq!(back.name(), Some("adder"));
        assert_eq!(back.links().len(), 2);
        assert_eq!(back.code, c.code);
    }

    #[test]
    fn invalid_body_rejected() {
        let mut h = CubeHeader::new();
        h.doc_type = Some("fn".into());
        h.refresh_flags();
        assert!(CodeCell::from_record(Czyx::new(1, 2, 3, 4), &h, &[0xFF]).is_none());
    }
}