superfluid_engine/mock/
strategy.rs

1//! Strategy registration.
2
3use superfluid_abi::{
4    cert_domain, cert_mode, exactness, sampling, strategy_cap, CapabilityReq,
5    ExactnessCert, LaneAdmit, SamplingParams, StrategyRegistration,
6    TapSpec, RecordArena, Status,
7};
8
9#[derive(Debug, Clone)]
10pub struct Cert {
11    pub cert_id: u32,
12    pub exactness: u8,
13    pub sampling_modes: u8,
14    pub grammar_allowed: bool,
15    pub logit_bias_allowed: bool,
16    pub param_domain: u32,
17    pub verification_shape_class: u32,
18    pub admissible_host_identities: Vec<[u8; 32]>,
19}
20
21impl Cert {
22    pub fn to_abi(&self, arena: &mut RecordArena) -> ExactnessCert {
23        ExactnessCert {
24            cert_id: self.cert_id,
25            exactness: self.exactness,
26            sampling_modes: self.sampling_modes,
27            grammar_allowed: self.grammar_allowed as u8,
28            logit_bias_allowed: self.logit_bias_allowed as u8,
29            param_domain: self.param_domain,
30            verification_shape_class: self.verification_shape_class,
31            admissible_host_identities: arena.push_records(&self.admissible_host_identities),
32        }
33    }
34}
35
36#[derive(Debug, Clone)]
37#[allow(dead_code)]
38pub struct RegisteredStrategy {
39    pub slot: u32,
40    pub strategy_id: String,
41    pub impl_hash: [u8; 32],
42    pub config_hash: [u8; 32],
43    pub host_compatible: bool,
44    pub certs: Vec<Cert>,
45    pub reserved_bytes: u64,
46}
47
48#[allow(dead_code)]
49pub struct RegistrationCheck<'a> {
50    pub reg: &'a StrategyRegistration,
51    pub capabilities: Vec<CapabilityReq>,
52    pub taps: Vec<TapSpec>,
53    pub target_archs: Vec<String>,
54}
55
56impl RegistrationCheck<'_> {
57    pub fn validate(&self, target_layers: u32, target_arch: &str) -> Result<bool, Status> {
58        let mut host_compatible = false;
59        for cap in &self.capabilities {
60            match cap.kind_id {
61                strategy_cap::PROPOSAL_LINEAR
62                | strategy_cap::PROPOSAL_BLOCK
63                | strategy_cap::PROPOSAL_TREE
64                | strategy_cap::MASK_ANCESTOR => {}
65                strategy_cap::HOST_COMPATIBLE => host_compatible = true,
66                _ => return Err(Status::RegistrationRefused),
67            }
68        }
69        for tap in &self.taps {
70            if tap.layer >= target_layers {
71                return Err(Status::RegistrationRefused);
72            }
73        }
74        if !self.target_archs.is_empty() && !self.target_archs.iter().any(|a| a == target_arch) {
75            return Err(Status::RegistrationRefused);
76        }
77        Ok(host_compatible)
78    }
79}
80
81pub fn resolve_certificate(
82    strategy: &RegisteredStrategy,
83    admit: &LaneAdmit,
84) -> Option<(u32, u8)> {
85    let mode_bit = match admit.sampling {
86        sampling::GPU_GREEDY => cert_mode::GREEDY,
87        sampling::GPU_GUMBEL => cert_mode::GUMBEL,
88        sampling::HOST => cert_mode::HOST,
89        _ => return None,
90    };
91    let zero_identity = admit.host_sampler_identity == [0u8; 32];
92    strategy
93        .certs
94        .iter()
95        .filter(|c| {
96            if admit.required_cert_id != 0 && c.cert_id != admit.required_cert_id {
97                return false;
98            }
99            if c.sampling_modes & mode_bit == 0 {
100                return false;
101            }
102            if admit.sampling == sampling::HOST
103                && (zero_identity || !c.admissible_host_identities.contains(&admit.host_sampler_identity))
104            {
105                return false;
106            }
107            if !param_domain_matches(c.param_domain, &admit.params) {
108                return false;
109            }
110            if admit.grammar_handle != 0 && !c.grammar_allowed {
111                return false;
112            }
113            c.exactness >= admit.minimum_exactness
114        })
115        .max_by_key(|c| c.exactness)
116        .map(|c| (c.cert_id, c.exactness))
117}
118
119fn param_domain_matches(domain: u32, p: &SamplingParams) -> bool {
120    let penalty_free = p.freq_penalty == 0.0
121        && p.presence_penalty == 0.0
122        && (p.repeat_penalty == 0.0 || p.repeat_penalty == 1.0);
123    match domain {
124        cert_domain::GREEDY_ONLY => p.temperature == 0.0 && penalty_free,
125        cert_domain::PENALTY_FREE => penalty_free,
126        cert_domain::ANY => true,
127        _ => false,
128    }
129}
130
131pub fn default_certs(strategy_id: &str, next_cert_id: &mut u32) -> Vec<Cert> {
132    match strategy_id {
133        "prompt-lookup" => {
134            let id = *next_cert_id;
135            *next_cert_id += 1;
136            vec![Cert {
137                cert_id: id,
138                exactness: exactness::SEED_PATH_INVARIANT,
139                sampling_modes: cert_mode::GREEDY,
140                grammar_allowed: false,
141                logit_bias_allowed: false,
142                param_domain: cert_domain::GREEDY_ONLY,
143                verification_shape_class: 1,
144                admissible_host_identities: Vec::new(),
145            }]
146        }
147        _ => Vec::new(),
148    }
149}
superfluid_engine/mock/strategy.rs

superfluid_engine/mock/
strategy.rs
