open-standards/wai/wai-pack/src/live.rs

//! The live packer ([packaging] §9.2, §4.3): frames pushed one at a time, each piece
//! coded as the segmenter closes it, by the coder the on-demand packer uses
//! ([`crate::pack::Coder`]), so a piece is the same objects either way, and the pieces
//! gathered into windows, each a stage index (§6.1) sealed by a version-2 receipt that
//! names the previous window's (parent kind `live-window`, §9.2).
//!
//! `wai.pack.scenecut/1` decides a frame's cut once `lookahead_frames` more have arrived
//! (§4.3), so a live and an on-demand packer under one policy cut the same pieces and
//! write the same objects. A live packer writes no floor: a supplied floor is made from
//! the whole source after the fact (§7.3), so it is the on-demand package's alone, and a
//! window lists the primary objects, its bases and refinements.
//!
//! A piece's objects are published when the frame that closes it has been pushed: its
//! base then, and layer `k` `layer_lag_ms[k]` after it (§8.6), which [`Published`] states
//! for the delivery simulator (`crate::delivery`).
//!
//! What it holds. Each object's bytes go to the caller in its [`Published`] and are not
//! kept: an open piece keeps, per object, its entry's fields (hash, prefix digest,
//! capability, lengths, layer, tier, deadline, SSE), from which its window's index is
//! built. Besides those it holds the frames from the first of the piece being cut (at
//! most the policy's longest piece plus `lookahead_frames`), each piece's first frame
//! for the plan file (eight bytes a piece, for the session), and the windows sealed and
//! not yet taken ([`LivePacker::take_sealed`]). With windows of `n` pieces the open
//! records are at most `n` pieces'; with one window over the session (`window_pieces`
//! 0) they are every piece's, sealed by [`LivePacker::finish`].
//!
//! [packaging]: ../../../extensions/packaging.md

use crate::pack::{CodedPiece, Coder, PackError};
use crate::policy::{Media, Policy};
use crate::scenecut::{SegParams, Segmenter};
use ed25519_dalek::SigningKey;
use wai::package::{PackageReceiptV2, Parent};
use wai::quantum_ops::GrantRef;
use wai::stage_index::Role;

/// What a live packer seals windows with.
pub struct LiveOpts {
    pub signer: SigningKey,
    pub signer_id: String,
    /// Pieces per window; 0 for one window over the whole session.
    pub window_pieces: u64,
    pub grant: GrantRef,
}

impl LiveOpts {
    /// One window, unbounded grant `wai.video.package`.
    pub fn new(signer: SigningKey, signer_id: impl Into<String>) -> LiveOpts {
        LiveOpts { signer, signer_id: signer_id.into(), window_pieces: 0, grant: GrantRef::unbounded("wai.video.package") }
    }
}

/// One object published.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Published {
    /// The absolute piece number.
    pub piece: u64,
    /// The piece's first frame and its frames.
    pub start: u64,
    pub duration: u64,
    /// The frame whose push closed the piece: the object is published once it has
    /// arrived (and, for layer `k`, `layer_lag_ms[k]` after that).
    pub at_frame: u64,
    pub layer: u8,
    pub role: Role,
    pub capability: String,
    /// A refinement's deadline (§8.6); 0 for a base.
    pub deadline_ms: u32,
    pub hash: [u8; 32],
    pub envelope: Vec<u8>,
}

/// What [`LivePacker::finish`] gives: the objects it published last, every window, and
/// the plan (each piece's first frame).
pub type Finished = (Vec<Published>, Vec<LiveWindow>, Vec<u64>);

/// A window: its stage index and the receipt that seals it, and the pieces it holds.
#[derive(Debug, Clone)]
pub struct LiveWindow {
    pub index: Vec<u8>,
    pub receipt: PackageReceiptV2,
    /// Its window number from 0, its first piece and its pieces.
    pub number: u64,
    pub first_piece: u64,
    pub pieces: u64,
}

/// The live packer of a one-component video policy.
pub struct LivePacker {
    policy: Policy,
    w: u32,
    h: u32,
    opts: LiveOpts,
    seg: Option<Segmenter>,
    /// The frames from the first of the piece being cut, the first at `held_from`.
    held: Vec<Vec<u8>>,
    held_from: u64,
    pushed: u64,
    starts: Vec<u64>,
    open: Vec<CodedPiece>,
    first_piece: u64,
    first_entry: u64,
    prev_index: [u8; 32],
    prev_receipt: Option<[u8; 32]>,
    /// Windows sealed and not yet taken, and how many were sealed in all.
    windows: Vec<LiveWindow>,
    sealed: u64,
}

impl LivePacker {
    /// A live packer of frames of `w × h` under `policy_bytes`, refusing as the on-demand
    /// packer does before any frame: the policy's own refusals, another medium, a rung it
    /// does not code.
    pub fn new(policy_bytes: &[u8], w: u32, h: u32, opts: LiveOpts) -> Result<LivePacker, PackError> {
        let policy = Policy::parse_bytes(policy_bytes).map_err(|e| PackError::new(e.code(), e.to_string()))?;
        if let Some(r) = crate::verify::policy_refusal(&policy) {
            return Err(PackError::new(r.code, r.detail));
        }
        let [comp] = policy.components.as_slice() else {
            return Err(PackError::new("pack-media-unsupported", "this revision packs one component"));
        };
        if comp.media != Media::Video {
            return Err(PackError::new("pack-media-unsupported", format!("component {} is {}", comp.id, comp.media.label())));
        }
        Coder::new(&policy, comp, w, h)?;
        let seg = Segmenter::new(SegParams::from_policy(&policy.segmentation, comp.timebase), w, h);
        Ok(LivePacker {
            policy,
            w,
            h,
            opts,
            seg: Some(seg),
            held: Vec::new(),
            held_from: 0,
            pushed: 0,
            starts: Vec::new(),
            open: Vec::new(),
            first_piece: 0,
            first_entry: 0,
            prev_index: [0u8; 32],
            prev_receipt: None,
            windows: Vec::new(),
            sealed: 0,
        })
    }

    /// The policy, every default written in.
    pub fn policy(&self) -> &Policy {
        &self.policy
    }

    fn code(&mut self, pieces: Vec<crate::scenecut::Piece>) -> Result<Vec<Published>, PackError> {
        let mut out = Vec::new();
        let at_frame = self.pushed.saturating_sub(1);
        for piece in pieces {
            let comp = &self.policy.components[0];
            let coder = Coder { policy: &self.policy, comp, w: self.w, h: self.h };
            let from = (piece.start - self.held_from) as usize;
            let mut cp = coder.code(&piece, &self.held[from..from + piece.len as usize])?;
            let number = self.first_piece + self.open.len() as u64;
            // Each object's bytes go to the caller; the piece keeps its entry's fields.
            for s in &mut cp.stages {
                out.push(Published {
                    piece: number,
                    start: piece.start,
                    duration: piece.len,
                    at_frame,
                    layer: s.layer,
                    role: s.role,
                    capability: s.capability.clone(),
                    deadline_ms: s.deadline_ms,
                    hash: s.hash,
                    envelope: std::mem::take(&mut s.envelope),
                });
            }
            // The frames before the next piece are no longer needed.
            let end = (piece.start + piece.len - self.held_from) as usize;
            self.held.drain(..end);
            self.held_from = piece.start + piece.len;
            self.starts.push(piece.start);
            self.open.push(cp);
            if self.opts.window_pieces > 0 && self.open.len() as u64 == self.opts.window_pieces {
                self.close_window()?;
            }
        }
        Ok(out)
    }

    /// Push the next frame (RGB8, `3·w·h` bytes); the objects of every piece it closes.
    pub fn push(&mut self, rgb: &[u8]) -> Result<Vec<Published>, PackError> {
        if rgb.len() as u64 != 3 * self.w as u64 * self.h as u64 {
            return Err(PackError::new("y4m-size-change", "a frame of another size"));
        }
        self.held.push(rgb.to_vec());
        self.pushed += 1;
        let pieces = self.seg.as_mut().expect("open until finish").push(rgb);
        self.code(pieces)
    }

    /// Seal the open pieces as a window (a no-op with none open).
    pub fn close_window(&mut self) -> Result<(), PackError> {
        if self.open.is_empty() {
            return Ok(());
        }
        let comp = &self.policy.components[0];
        let w = crate::pack::window_index(&self.policy, comp, &self.open, self.first_piece, self.first_entry, self.prev_index)?;
        let bytes = w.index.to_bytes();
        let parent = match self.prev_receipt {
            None => Parent::None,
            Some(h) => Parent::LiveWindow(h),
        };
        let receipt = PackageReceiptV2::seal(
            &self.opts.signer,
            &self.opts.signer_id,
            &self.policy.asset_id,
            self.policy.digest(),
            &[(comp.id.clone(), bytes.clone())],
            &crate::pack::packager_id(),
            wai::energy_meter::EnergyReading::unmetered(),
            self.opts.grant.clone(),
            parent,
        )
        .map_err(|r| PackError::new(r.code, r.detail))?;
        let window = LiveWindow { index: bytes, receipt, number: self.sealed, first_piece: self.first_piece, pieces: self.open.len() as u64 };
        self.first_piece += self.open.len() as u64;
        self.first_entry += w.index.entries.len() as u64;
        self.prev_index = w.index.digest();
        self.prev_receipt = Some(window.receipt.receipt_hash());
        self.open.clear();
        self.sealed += 1;
        self.windows.push(window);
        Ok(())
    }

    /// The windows sealed since the last call, in order; the packer keeps none of them.
    pub fn take_sealed(&mut self) -> Vec<LiveWindow> {
        std::mem::take(&mut self.windows)
    }

    /// The end of the source: the last pieces coded and published, the last window
    /// sealed. Returns the published objects, the windows not yet taken (every window,
    /// if none was), and the plan (each piece's first frame).
    pub fn finish(mut self) -> Result<Finished, PackError> {
        let pieces = self.seg.take().expect("open until finish").finish();
        let out = self.code(pieces)?;
        self.close_window()?;
        Ok((out, self.windows, self.starts))
    }
}

/// What a live and an on-demand pack of one source share (test 12): of `files`, the plan
/// files (`plan-<c>.txt`) and the objects the stage indexes (`index-<c>.wsix` at the top,
/// and `windows/<n>/index-<c>.wsix`) list as bases and refinements. A floor object, which
/// a live pack does not write, is left out.
pub fn primary_files(files: &std::collections::BTreeMap<String, Vec<u8>>) -> Result<std::collections::BTreeMap<String, Vec<u8>>, String> {
    let is_index = |p: &str| {
        let name = p.strip_prefix("windows/").and_then(|r| r.split_once('/')).map_or(p, |(_, n)| n);
        name.starts_with("index-") && name.ends_with(".wsix") && (p == name || p.starts_with("windows/"))
    };
    let mut out = std::collections::BTreeMap::new();
    for (p, b) in files {
        if p.starts_with("plan-") && p.ends_with(".txt") && !p.contains('/') {
            out.insert(p.clone(), b.clone());
        }
        if is_index(p) {
            let ix = wai::stage_index::StageIndex::parse(b).map_err(|e| format!("{p}: {e}"))?;
            for e in ix.entries.iter().filter(|e| matches!(e.role, Role::Base | Role::Refinement)) {
                let o = crate::objects::object_path(&e.content_hash);
                let bytes = files.get(&o).ok_or_else(|| format!("{o}: listed by {p}, not held"))?;
                out.insert(o, bytes.clone());
            }
        }
    }
    Ok(out)
}

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

    /// What a live packer holds (§9.2): with one window over the session, every open
    /// piece's records and none of its objects' bytes, which went out in its
    /// publications; and a sealed window, once taken, not at all.
    #[test]
    fn a_published_object_is_not_held() {
        let g = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../pack-conformance/packs/video-staged-01");
        let policy = std::fs::read(g.join("policy.toml")).expect("the golden's policy");
        let (hd, frames) = crate::source::y4m::read_all(&crate::fixture::generate("video-staged-01").expect("a fixture")).expect("its frames");
        let key = SigningKey::from_bytes(&[7; 32]);
        let mut lp = LivePacker::new(&policy, hd.w, hd.h, LiveOpts::new(key, "test")).expect("a live packer");
        let mut published = 0u64;
        for f in &frames {
            published += lp.push(&f.rgb).expect("a frame").iter().map(|p| p.envelope.len() as u64).sum::<u64>();
        }
        assert!(lp.open.len() >= 4, "every piece but the last is open: {}", lp.open.len());
        assert!(published > 0);
        assert!(lp.open.iter().flat_map(|cp| &cp.stages).all(|s| s.envelope.is_empty() && s.lens.byte_len > 0), "a published object's bytes are held");
        lp.close_window().expect("a window");
        assert_eq!(lp.take_sealed().len(), 1);
        assert!(lp.take_sealed().is_empty() && lp.open.is_empty(), "a taken window, or its pieces, are held");
    }
}