use alloc::collections::BTreeMap;
use alloc::string::String;
use alloc::vec::Vec;
use serde::{Deserialize, Serialize};
use crate::config::{ConfigSchema, ConfigValue};
macro_rules! id_type {
($name:ident) => {
#[derive(
Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize,
)]
#[serde(transparent)]
pub struct $name(pub u32);
};
}
id_type!(NodeId);
id_type!(KernelId);
id_type!(BufferId);
id_type!(StepId);
id_type!(FeedbackId);
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct GraphId(pub [u8; 32]);
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct PlanId(pub [u8; 32]);
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct SubgraphId(pub [u8; 32]);
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct ImplementationId(pub [u8; 32]);
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct NodeTypeRef {
pub type_name: String,
pub version: u32,
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Capability(pub String);
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Target(u32);
#[allow(non_upper_case_globals)]
impl Target {
pub const McuAot: Self = Self(0);
pub const Host: Self = Self(1);
pub const BlutDurable: Self = Self(2);
pub const KNOWN: [Self; 3] = [Self::McuAot, Self::Host, Self::BlutDurable];
pub const fn token(self) -> u32 {
self.0
}
pub const fn from_token(token: u32) -> Self {
Self(token)
}
pub const fn is_known(self) -> bool {
self.0 <= 2
}
const fn name(self) -> Option<&'static str> {
match self.0 {
0 => Some("McuAot"),
1 => Some("Host"),
2 => Some("BlutDurable"),
_ => None,
}
}
}
impl core::fmt::Debug for Target {
fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self.name() {
Some(name) => formatter.write_str(name),
None => write!(formatter, "Target({})", self.0),
}
}
}
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct ExecutionRealm(u32);
#[allow(non_upper_case_globals)]
impl ExecutionRealm {
pub const McuAot: Self = Self(0);
pub const HostStream: Self = Self(1);
pub const BlutDurable: Self = Self(2);
pub const KNOWN: [Self; 3] = [Self::McuAot, Self::HostStream, Self::BlutDurable];
pub const fn token(self) -> u32 {
self.0
}
pub const fn from_token(token: u32) -> Self {
Self(token)
}
pub const fn is_known(self) -> bool {
self.0 <= 2
}
pub const fn target(self) -> Target {
match self.0 {
0 => Target::McuAot,
1 => Target::Host,
2 => Target::BlutDurable,
other => Target::from_token(other),
}
}
const fn name(self) -> Option<&'static str> {
match self.0 {
0 => Some("McuAot"),
1 => Some("HostStream"),
2 => Some("BlutDurable"),
_ => None,
}
}
}
impl core::fmt::Debug for ExecutionRealm {
fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self.name() {
Some(name) => formatter.write_str(name),
None => write!(formatter, "ExecutionRealm({})", self.0),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum Determinism {
BitExact,
NumericallyEquivalent,
Seeded,
Nondeterministic,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum Effect {
Pure,
Idempotent,
Transactional,
AtMostOnce,
AtLeastOnce,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum Partiality {
Atomic,
ExplicitGaps,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct FailureContract {
pub domains: Vec<String>,
}
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Layout(u32);
#[allow(non_upper_case_globals)]
impl Layout {
pub const Canonical: Self = Self(0);
pub const ChannelMajor: Self = Self(1);
pub const TimeMajor: Self = Self(2);
pub const Packed: Self = Self(3);
pub const Opaque: Self = Self(4);
pub const KNOWN: [Self; 5] = [
Self::Canonical,
Self::ChannelMajor,
Self::TimeMajor,
Self::Packed,
Self::Opaque,
];
pub const fn token(self) -> u32 {
self.0
}
pub const fn from_token(token: u32) -> Self {
Self(token)
}
pub const fn is_known(self) -> bool {
self.0 <= 4
}
const fn name(self) -> Option<&'static str> {
match self.0 {
0 => Some("Canonical"),
1 => Some("ChannelMajor"),
2 => Some("TimeMajor"),
3 => Some("Packed"),
4 => Some("Opaque"),
_ => None,
}
}
}
impl core::fmt::Debug for Layout {
fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self.name() {
Some(name) => formatter.write_str(name),
None => write!(formatter, "Layout({})", self.0),
}
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(transparent)]
pub struct DomainToken(String);
impl DomainToken {
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
pub fn as_str(&self) -> &str {
&self.0
}
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
}
impl From<&str> for DomainToken {
fn from(value: &str) -> Self {
Self(value.into())
}
}
impl From<String> for DomainToken {
fn from(value: String) -> Self {
Self(value)
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct DomainType {
pub root: DomainToken,
pub view: DomainToken,
}
impl DomainType {
pub fn new(root: impl Into<DomainToken>, view: impl Into<DomainToken>) -> Self {
Self {
root: root.into(),
view: view.into(),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ExtentContract {
pub rank: u8,
pub maximum_shape: Vec<u64>,
pub max_elements: u64,
pub ragged: bool,
pub sparse: bool,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum LeaseAccess {
ReadOnly,
ExclusiveWrite,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum LeaseLifetime {
Step,
Invocation,
Session,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LeaseContract {
pub access: LeaseAccess,
pub lifetime: LeaseLifetime,
pub zero_copy_permitted: bool,
pub contiguous_required: bool,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResourceEnvelope {
pub peak_bytes: u64,
pub scratch_bytes: u64,
pub threads: u16,
pub device: Option<String>,
}
impl ResourceEnvelope {
pub const fn bounded(peak_bytes: u64, scratch_bytes: u64, threads: u16) -> Self {
Self {
peak_bytes,
scratch_bytes,
threads,
device: None,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct PortDescriptor {
pub name: String,
pub semantic_type: String,
pub optional: bool,
pub layouts: Vec<Layout>,
pub max_bytes: u64,
pub domain: DomainType,
pub proof: ProofContract,
pub policy: PolicyContract,
pub fidelity: FidelityContract,
pub extent: ExtentContract,
pub lease: LeaseContract,
}
impl PortDescriptor {
pub fn opaque(
name: impl Into<String>,
semantic_type: impl Into<String>,
max_bytes: u64,
) -> Self {
Self {
name: name.into(),
semantic_type: semantic_type.into(),
optional: false,
layouts: alloc::vec![Layout::Canonical],
max_bytes,
domain: DomainType {
root: DomainToken::new("blob-ref"),
view: DomainToken::new("atom"),
},
proof: ProofContract {
requires: Vec::new(),
provides: Vec::new(),
invalidates: Vec::new(),
},
policy: PolicyContract {
requires: Vec::new(),
adds: Vec::new(),
},
fidelity: FidelityContract {
minimum_input: 0,
maximum_loss: 0,
},
extent: ExtentContract {
rank: 0,
maximum_shape: Vec::new(),
max_elements: 1,
ragged: false,
sparse: false,
},
lease: LeaseContract {
access: LeaseAccess::ReadOnly,
lifetime: LeaseLifetime::Invocation,
zero_copy_permitted: false,
contiguous_required: false,
},
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProofContract {
pub requires: Vec<String>,
pub provides: Vec<String>,
pub invalidates: Vec<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct PolicyContract {
pub requires: Vec<String>,
pub adds: Vec<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct FidelityContract {
pub minimum_input: u16,
pub maximum_loss: u16,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum StateScope {
Stateless,
Invocation,
Session,
Durable,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum CheckpointMode {
Disabled,
Optional,
Required,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CheckpointContract {
pub mode: CheckpointMode,
pub max_snapshot_bytes: u64,
pub max_interval_invocations: u32,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct StateContract {
pub scope: StateScope,
pub max_bytes: u64,
pub checkpoint: CheckpointContract,
}
impl StateContract {
pub const fn stateless() -> Self {
Self {
scope: StateScope::Stateless,
max_bytes: 0,
checkpoint: CheckpointContract {
mode: CheckpointMode::Disabled,
max_snapshot_bytes: 0,
max_interval_invocations: 0,
},
}
}
pub const fn checkpointable(&self) -> bool {
!matches!(self.checkpoint.mode, CheckpointMode::Disabled)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SessionContract {
pub namespace: String,
pub max_concurrent_sessions: u32,
pub max_idle_millis: u64,
pub reset_on_plan_change: bool,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct DelayContract {
pub invocations: u32,
pub initial: DelayInitial,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum DelayInitial {
Absent,
Zeroed,
ContentId([u8; 32]),
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct FeedbackEdge {
pub from: PortRef,
pub to: PortRef,
pub delay: DelayContract,
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct PortMap {
pub outer: String,
pub inner: String,
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct SubgraphConfigMap {
pub outer: String,
pub node: NodeId,
pub inner: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SubgraphLowering {
pub subgraph: SubgraphId,
pub input_map: Vec<PortMap>,
pub output_map: Vec<PortMap>,
pub config_map: Vec<SubgraphConfigMap>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SubgraphNode {
pub id: NodeId,
pub node_type: NodeTypeRef,
pub config: BTreeMap<String, ConfigValue>,
pub child: Option<SubgraphId>,
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct SubgraphInterfacePort {
pub name: String,
pub inner: PortRef,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SubgraphSchema {
pub id: SubgraphId,
pub version: u32,
pub nodes: Vec<SubgraphNode>,
pub edges: Vec<Edge>,
pub inputs: Vec<SubgraphInterfacePort>,
pub outputs: Vec<SubgraphInterfacePort>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MaterializedSubgraph {
pub graph: Graph,
pub inputs: Vec<SubgraphInterfacePort>,
pub outputs: Vec<SubgraphInterfacePort>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct NodeDescriptor {
pub type_name: String,
pub version: u32,
pub inputs: Vec<PortDescriptor>,
pub outputs: Vec<PortDescriptor>,
pub capabilities: Vec<Capability>,
pub targets: Vec<Target>,
pub resources: ResourceEnvelope,
pub determinism: Determinism,
pub config: ConfigSchema,
pub state: StateContract,
pub subgraph: Option<SubgraphLowering>,
pub proof: ProofContract,
pub policy: PolicyContract,
pub fidelity: FidelityContract,
pub partiality: Partiality,
pub failure: FailureContract,
pub effect: Effect,
pub retry_limit: u16,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct NodeInstance {
pub id: NodeId,
pub descriptor: String,
pub descriptor_version: u32,
pub config: BTreeMap<String, ConfigValue>,
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct PortRef {
pub node: NodeId,
pub port: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Edge {
pub from: PortRef,
pub to: PortRef,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Graph {
pub version: u32,
pub nodes: Vec<NodeInstance>,
pub edges: Vec<Edge>,
#[serde(default)]
pub feedback: Vec<FeedbackEdge>,
#[serde(default)]
pub invocation_inputs: Vec<PortRef>,
pub required_capabilities: Vec<Capability>,
pub required_proofs: Vec<String>,
pub policy: Vec<String>,
pub minimum_fidelity: u16,
pub session: Option<SessionContract>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct KernelDescriptor {
pub id: KernelId,
pub implements: Vec<NodeTypeRef>,
pub implementation_id: ImplementationId,
pub conversion: Option<LayoutConversion>,
pub target: Target,
pub input_layouts: Vec<Layout>,
pub output_layouts: Vec<Layout>,
pub resources: ResourceEnvelope,
pub determinism: Determinism,
pub lowering: String,
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct LayoutConversion {
pub semantic_type: String,
pub from: Layout,
pub to: Layout,
pub max_input_bytes: u64,
pub max_output_bytes: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum OutputBinding {
Buffer(BufferId),
Terminal,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum InputBinding {
Buffer(BufferId),
Invocation(u32),
Feedback(FeedbackId),
Absent,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CompiledPortContract {
pub name: String,
pub semantic_type: String,
pub optional: bool,
pub layout: Layout,
pub max_bytes: u64,
pub domain: DomainType,
pub proof: ProofContract,
pub policy: PolicyContract,
pub fidelity: FidelityContract,
pub extent: ExtentContract,
pub lease: LeaseContract,
}
impl CompiledPortContract {
pub fn opaque(
name: impl Into<String>,
semantic_type: impl Into<String>,
layout: Layout,
max_bytes: u64,
) -> Self {
let port = PortDescriptor::opaque(name, semantic_type, max_bytes);
Self {
name: port.name,
semantic_type: port.semantic_type,
optional: port.optional,
layout,
max_bytes: port.max_bytes,
domain: port.domain,
proof: port.proof,
policy: port.policy,
fidelity: port.fidelity,
extent: port.extent,
lease: port.lease,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CompiledNode {
pub id: StepId,
pub semantic_nodes: Vec<NodeId>,
pub semantic_types: Vec<NodeTypeRef>,
pub semantic_configs: Vec<BTreeMap<String, ConfigValue>>,
pub kernel: KernelId,
pub implementation_id: ImplementationId,
pub resources: ResourceEnvelope,
pub determinism: Determinism,
pub lowering: String,
pub conversion: Option<LayoutConversion>,
pub input_ports: Vec<String>,
pub output_ports: Vec<String>,
pub input_contracts: Vec<CompiledPortContract>,
pub output_contracts: Vec<CompiledPortContract>,
pub input_bindings: Vec<InputBinding>,
pub output_bindings: Vec<OutputBinding>,
pub partiality: Partiality,
pub failure: FailureContract,
pub effect: Effect,
pub retry_limit: u16,
pub state: StateContract,
pub subgraph_path: Vec<SubgraphId>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct FeedbackPlan {
pub id: FeedbackId,
pub from_step: StepId,
pub from_port: u32,
pub to_step: StepId,
pub to_port: u32,
pub delay: DelayContract,
pub state_bytes: u64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BufferPlan {
pub id: BufferId,
pub layout: Layout,
pub capacity_bytes: u64,
pub producer: StepId,
pub consumers: Vec<StepId>,
pub last_consumer: StepId,
pub aliases: Option<BufferId>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CompiledPlan {
pub schema_version: u32,
pub graph_id: GraphId,
pub plan_id: PlanId,
pub realm: ExecutionRealm,
pub order: Vec<NodeId>,
pub nodes: Vec<CompiledNode>,
pub buffers: Vec<BufferPlan>,
pub feedback: Vec<FeedbackPlan>,
pub invocation_ports: Vec<PortRef>,
pub propagated_proofs: Vec<String>,
pub propagated_policy: Vec<String>,
pub resulting_fidelity: u16,
pub peak_bytes: u64,
pub persistent_state_bytes: u64,
pub session: Option<SessionContract>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AuthorizedPlan {
plan: CompiledPlan,
}
impl AuthorizedPlan {
pub(crate) const fn new(plan: CompiledPlan) -> Self {
Self { plan }
}
pub const fn as_plan(&self) -> &CompiledPlan {
&self.plan
}
pub fn into_plan(self) -> CompiledPlan {
self.plan
}
}
impl core::ops::Deref for AuthorizedPlan {
type Target = CompiledPlan;
fn deref(&self) -> &Self::Target {
&self.plan
}
}