Skip to main content

superfluid_engine/
plan.rs

1//! Plan decoding + structural validation.
2
3use std::collections::HashSet;
4
5use superfluid_abi::{
6    array::{read_array, read_versioned_struct},
7    sampling, LaneAdmit, LaneCommit, LaneDecode, LanePrefill,
8    LaneRetire, TickPlan, RecordArena, Status,
9};
10
11#[derive(Debug)]
12pub struct DecodedPlan {
13    pub plan_seq: u64,
14    pub flags: u32,
15    pub prefill_token_budget: u32,
16    pub max_decode_lanes: u32,
17    pub admits: Vec<LaneAdmit>,
18    pub commits: Vec<LaneCommit>,
19    pub prefills: Vec<LanePrefill>,
20    pub decodes: Vec<LaneDecode>,
21    pub retires: Vec<LaneRetire>,
22    pub victim_lanes: Vec<u64>,
23    pub evictable_cache_classes: u64,
24    pub protected_quota_bytes: u64,
25    pub max_evict_bytes: u64,
26    pub on_partial_emit: Option<
27        unsafe extern "C" fn(user: *mut core::ffi::c_void, lane_tag: u64, tokens: *const u32, n_tokens: u32),
28    >,
29    pub partial_emit_user: *mut core::ffi::c_void,
30}
31
32fn decode_records<T: superfluid_abi::AbiRecord>(
33    arr: &superfluid_abi::Array,
34    arena: &RecordArena,
35) -> Result<Vec<T>, Status> {
36    let bounds = arena.bounds_of(arr).ok_or(Status::RejectBounds)?;
37    // SAFETY: bounds_of proved the array lies inside one of the arena's
38    // live segments, which outlive this call; read_array enforces the
39    // remaining normative rules.
40    let iter = unsafe { read_array::<T>(arr, &bounds) }.map_err(|e| e.status())?;
41    Ok(iter.collect())
42}
43
44fn unique_tags(tags: impl Iterator<Item = u64>) -> Result<HashSet<u64>, Status> {
45    let mut set = HashSet::new();
46    for t in tags {
47        if !set.insert(t) {
48            return Err(Status::RejectDuplicateLane);
49        }
50    }
51    Ok(set)
52}
53
54impl DecodedPlan {
55    pub fn decode(plan: &TickPlan, arena: &RecordArena) -> Result<DecodedPlan, Status> {
56        // SAFETY: `plan` is a live reference; struct_size is its writer's
57        // declared size, bounded below by the min-prefix check inside.
58        let plan = unsafe { read_versioned_struct(plan, plan.struct_size.max(1)) }
59            .map_err(|e| e.status())?;
60
61        let admits: Vec<LaneAdmit> = decode_records(&plan.admits, arena)?;
62        let commits: Vec<LaneCommit> = decode_records(&plan.commits, arena)?;
63        let prefills: Vec<LanePrefill> = decode_records(&plan.prefills, arena)?;
64        let decodes: Vec<LaneDecode> = decode_records(&plan.decodes, arena)?;
65        let retires: Vec<LaneRetire> = decode_records(&plan.retires, arena)?;
66        let victim_lanes: Vec<u64> = decode_records(&plan.shed_policy.victim_lanes, arena)?;
67
68        let admit_tags = unique_tags(admits.iter().map(|r| r.lane_tag))?;
69        let commit_tags = unique_tags(commits.iter().map(|r| r.lane_tag))?;
70        let prefill_tags = unique_tags(prefills.iter().map(|r| r.lane_tag))?;
71        let decode_tags = unique_tags(decodes.iter().map(|r| r.lane_tag))?;
72        let retire_tags = unique_tags(retires.iter().map(|r| r.lane_tag))?;
73
74        if admit_tags.intersection(&retire_tags).next().is_some() {
75            return Err(Status::RejectIllegalCombination);
76        }
77        if admit_tags.intersection(&commit_tags).next().is_some() {
78            return Err(Status::RejectIllegalCombination);
79        }
80        if commit_tags.intersection(&retire_tags).next().is_some() {
81            return Err(Status::RejectIllegalCombination);
82        }
83        if decode_tags.intersection(&retire_tags).next().is_some() {
84            return Err(Status::RejectIllegalCombination);
85        }
86        if prefill_tags.intersection(&retire_tags).next().is_some() {
87            return Err(Status::RejectIllegalCombination);
88        }
89
90        for a in &admits {
91            if a.sampling == sampling::HOST {
92                if let Some(d) = decodes.iter().find(|d| d.lane_tag == a.lane_tag) {
93                    if d.max_new_tokens != 1 {
94                        return Err(Status::RejectHostRules);
95                    }
96                }
97            }
98        }
99
100        if decodes.len() as u32 > plan.max_decode_lanes {
101            return Err(Status::RejectBudget);
102        }
103
104        Ok(DecodedPlan {
105            plan_seq: plan.plan_seq,
106            flags: plan.flags,
107            prefill_token_budget: plan.prefill_token_budget,
108            max_decode_lanes: plan.max_decode_lanes,
109            admits,
110            commits,
111            prefills,
112            decodes,
113            retires,
114            victim_lanes,
115            evictable_cache_classes: plan.shed_policy.evictable_cache_classes,
116            protected_quota_bytes: plan.shed_policy.protected_quota_bytes,
117            max_evict_bytes: plan.shed_policy.max_evict_bytes,
118            on_partial_emit: plan.on_partial_emit,
119            partial_emit_user: plan.partial_emit_user,
120        })
121    }
122}