superfluid_engine/
plan.rs1use 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 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 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}