spirt/spv/
spec.rs

1//! SPIR-V specification parsing/indexing.
2
3use crate::FxIndexSet;
4use arrayvec::ArrayVec;
5use lazy_static::lazy_static;
6use rustc_hash::FxHashMap;
7use std::borrow::Cow;
8use std::collections::BTreeMap;
9use std::{fmt, iter};
10
11use self::indexed::FlatIdx as _;
12
13pub const HEADER_LEN: usize = 5;
14
15pub struct Spec {
16    pub magic: u32,
17
18    /// Pre-cached IDs for "well-known" names.
19    pub well_known: WellKnown,
20
21    pub instructions: indexed::NamedIdxMap<Opcode, InstructionDef, indexed::KhrSegmented>,
22
23    pub operand_kinds: indexed::NamedIdxMap<OperandKind, OperandKindDef, indexed::Flat>,
24
25    // HACK(eddyb) ad-hoc interning, to reduce the cost of tracking operand names
26    // down to a single extra byte per operand (see `PackedOperandNameAndKind`).
27    operand_names: FxIndexSet<&'static str>,
28
29    // HACK(eddyb) the `OnceLock`s allow lazy parsing, avoiding overhead.
30    // FIXME(eddyb) fix type complexity once `LazyLock` stabilizes.
31    #[allow(clippy::type_complexity)]
32    ext_inst_sets: BTreeMap<
33        &'static str,
34        (std::sync::OnceLock<ExtInstSetDesc>, Box<dyn Fn() -> ExtInstSetDesc + Send + Sync>),
35    >,
36}
37
38/// Simplified information for pretty-printing "extended instruction" sets.
39pub struct ExtInstSetDesc {
40    /// Shorter name to use during pretty-printing.
41    pub short_alias: Option<Cow<'static, str>>,
42
43    pub instructions: BTreeMap<u32, ExtInstSetInstructionDesc>,
44}
45
46/// Simplified [`InstructionDef`] for pretty-printing "extended instruction" sets.
47pub struct ExtInstSetInstructionDesc {
48    pub name: Cow<'static, str>,
49    pub operand_names: Vec<Cow<'static, str>>,
50
51    /// Whether this instruction is non-semantic debuginfo and should therefore
52    /// be pretty-printed on a single line, as a comment.
53    //
54    // FIXME(eddyb) allow customizing the formatting, but this works for now.
55    pub is_debuginfo: bool,
56}
57
58macro_rules! def_well_known {
59    ($($group:ident: $ty:ty = [$($entry:ident),+ $(,)?]),+ $(,)?) => {
60        // FIXME(eddyb) decide whether to split this type into one per-group.
61        #[allow(non_snake_case)]
62        pub struct WellKnown {
63            $($(pub $entry: $ty,)+)+
64        }
65
66        #[allow(non_camel_case_types)]
67        struct PerWellKnownGroup<$($group),+> {
68            $($group: $group),+
69        }
70
71        impl WellKnown {
72            fn lookup_with(lookup_fns: PerWellKnownGroup<$(impl Fn(&'static str) -> $ty),+>) -> Self {
73                Self {
74                    $($($entry: (lookup_fns.$group)(stringify!($entry)),)+)+
75                }
76            }
77        }
78    };
79}
80
81// FIXME(eddyb) maybe sort some of these groups alphabetically.
82def_well_known! {
83    opcode: Opcode = [
84        OpNop,
85
86        OpCapability,
87        OpExtension,
88        OpExtInstImport,
89        OpExtInst,
90
91        OpMemoryModel,
92
93        OpEntryPoint,
94        OpExecutionMode,
95        OpExecutionModeId,
96
97        OpString,
98        OpSource,
99        OpSourceContinued,
100        OpSourceExtension,
101        OpName,
102        OpMemberName,
103        OpModuleProcessed,
104
105        OpDecorate,
106        OpMemberDecorate,
107        OpDecorateId,
108        OpDecorateString,
109        OpMemberDecorateString,
110
111        // Deprecated in favor of `OpDecorate`/`OpMemberDecorate`.
112        OpDecorationGroup,
113        OpGroupDecorate,
114        OpGroupMemberDecorate,
115
116        OpLine,
117        OpNoLine,
118
119        OpTypeVoid,
120        OpTypeBool,
121        OpTypeInt,
122        OpTypeFloat,
123        OpTypeVector,
124        OpTypeMatrix,
125        OpTypeArray,
126        OpTypeRuntimeArray,
127        OpTypeStruct,
128        OpTypeForwardPointer,
129        OpTypePointer,
130        OpTypeFunction,
131        OpTypeImage,
132        OpTypeSampler,
133        OpTypeSampledImage,
134        OpTypeAccelerationStructureKHR,
135
136        OpConstantFalse,
137        OpConstantTrue,
138        OpConstant,
139        OpUndef,
140
141        OpVariable,
142
143        OpFunction,
144        OpFunctionParameter,
145        OpFunctionEnd,
146
147        OpLabel,
148        OpPhi,
149        OpSelectionMerge,
150        OpLoopMerge,
151
152        OpUnreachable,
153        OpReturn,
154        OpReturnValue,
155        OpBranch,
156        OpBranchConditional,
157        OpSwitch,
158
159        OpFunctionCall,
160
161        OpLoad,
162        OpStore,
163        OpArrayLength,
164        OpAccessChain,
165        OpInBoundsAccessChain,
166        OpPtrAccessChain,
167        OpInBoundsPtrAccessChain,
168        OpBitcast,
169    ],
170    operand_kind: OperandKind = [
171        Capability,
172        AddressingModel,
173        MemoryModel,
174        SourceLanguage,
175        StorageClass,
176        FunctionControl,
177        Decoration,
178        LinkageType,
179        SelectionControl,
180        LoopControl,
181
182        LiteralInteger,
183        LiteralExtInstInteger,
184        LiteralString,
185        LiteralContextDependentNumber,
186    ],
187    // FIXME(eddyb) find a way to namespace these to avoid conflicts.
188    addressing_model: u32 = [
189        Logical,
190    ],
191    storage_class: u32 = [
192        Function,
193
194        UniformConstant,
195        Input,
196        Output,
197
198        IncomingRayPayloadKHR,
199        IncomingCallableDataKHR,
200        HitAttributeKHR,
201        RayPayloadKHR,
202        CallableDataKHR,
203    ],
204    decoration: u32 = [
205        LinkageAttributes,
206
207        ArrayStride,
208
209        Block,
210        RowMajor,
211        Offset,
212    ],
213    linkage_type: u32 = [
214        Import,
215        Export,
216    ],
217}
218
219#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
220pub struct Opcode(u16);
221
222impl indexed::FlatIdx for Opcode {
223    fn to_usize(self) -> usize {
224        self.0.into()
225    }
226}
227
228impl Opcode {
229    /// Lookup the name & definition for `opcode` in the lazily-loaded [`Spec`],
230    /// returning `None` if it's not a known opcode.
231    pub fn try_from_u16_with_name_and_def(
232        opcode: u16,
233    ) -> Option<(Self, &'static str, &'static InstructionDef)> {
234        let opcode = Self(opcode);
235        let (name, def) = Spec::get().instructions.get_named(opcode)?;
236        Some((opcode, name, def))
237    }
238
239    pub fn as_u16(self) -> u16 {
240        self.0
241    }
242
243    /// Lookup the name & definition for this opcode in the lazily-loaded [`Spec`].
244    #[inline]
245    pub fn name_and_def(self) -> (&'static str, &'static InstructionDef) {
246        Spec::get().instructions.get_named(self).unwrap()
247    }
248
249    /// Lookup the name for this opcode in the lazily-loaded [`Spec`].
250    #[inline]
251    pub fn name(self) -> &'static str {
252        self.name_and_def().0
253    }
254
255    /// Lookup the definition for this opcode in the lazily-loaded [`Spec`].
256    #[inline]
257    pub fn def(self) -> &'static InstructionDef {
258        self.name_and_def().1
259    }
260}
261
262#[derive(PartialEq, Eq)]
263pub struct InstructionDef {
264    pub category: InstructionCategory,
265
266    // FIXME(eddyb) consider nesting "Result Type ID" in "Result ID".
267    pub has_result_type_id: bool,
268    pub has_result_id: bool,
269
270    pub req_operands: ArrayVec<PackedOperandNameAndKind, 14>,
271    pub opt_operands: ArrayVec<PackedOperandNameAndKind, 2>,
272    pub rest_operands: Option<RestOperandsUnit>,
273}
274
275#[derive(Copy, Clone, Debug, PartialEq, Eq)]
276pub enum InstructionCategory {
277    Type,
278    Const,
279    ControlFlow,
280    Other,
281}
282
283/// Whether the trailing `*` "operand" (i.e. repeated arbitrarily many times),
284/// consists of just one operand, or two per repeat (used by e.g. `OpPhi`).
285#[derive(PartialEq, Eq)]
286pub enum RestOperandsUnit {
287    One(OperandKind),
288    Two([OperandKind; 2]),
289}
290
291#[derive(Copy, Clone, PartialEq, Eq)]
292pub enum OperandMode {
293    Required,
294    Optional,
295}
296
297impl InstructionDef {
298    /// Return a (potentially infinite) iterator of [`OperandKind`]s, along with
299    /// the [`OperandMode`] indicating whether an operand is expected (`Required`),
300    /// or that an operand's absence signals the end of operands (`Optional`),
301    /// which is also the exit signal for the "rest operands" infinite iterators.
302    pub fn all_operands(&self) -> impl Iterator<Item = (OperandMode, OperandKind)> + '_ {
303        self.all_operands_with_names().map(|(mode, name_and_kind)| (mode, name_and_kind.kind()))
304    }
305
306    /// Like `all_operands`, but providing access to the operand names as well.
307    pub fn all_operands_with_names(
308        &self,
309    ) -> impl Iterator<Item = (OperandMode, PackedOperandNameAndKind)> + '_ {
310        self.req_operands
311            .iter()
312            .copied()
313            .map(|name_and_kind| (OperandMode::Required, name_and_kind))
314            .chain(
315                self.opt_operands
316                    .iter()
317                    .copied()
318                    .map(|name_and_kind| (OperandMode::Optional, name_and_kind)),
319            )
320            .chain(self.rest_operands.iter().flat_map(|rest_unit| {
321                // If the rest operands come in pairs, only the first operand in
322                // the pair is optional, the second one must be present when the
323                // first one is (i.e. only the pair as a whole is optional).
324                let (opt_a, req_b) = match *rest_unit {
325                    RestOperandsUnit::One(kind) => (kind, None),
326                    RestOperandsUnit::Two([a_kind, b_kind]) => (a_kind, Some(b_kind)),
327                };
328                iter::repeat(
329                    iter::once((OperandMode::Optional, PackedOperandNameAndKind::unnamed(opt_a)))
330                        .chain(req_b.map(|kind| {
331                            (OperandMode::Required, PackedOperandNameAndKind::unnamed(kind))
332                        })),
333                )
334                .flatten()
335            }))
336    }
337}
338
339#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
340pub struct OperandKind(u8);
341
342impl indexed::FlatIdx for OperandKind {
343    fn to_usize(self) -> usize {
344        self.0.into()
345    }
346}
347
348impl fmt::Debug for OperandKind {
349    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
350        write!(f, "OperandKind({} => {:?})", self.0, self.name())
351    }
352}
353
354impl OperandKind {
355    /// Lookup the name & definition for this operand kind in the lazily-loaded [`Spec`].
356    #[inline]
357    pub fn name_and_def(self) -> (&'static str, &'static OperandKindDef) {
358        Spec::get().operand_kinds.get_named(self).unwrap()
359    }
360
361    /// Lookup the name for this operand kind in the lazily-loaded [`Spec`].
362    #[inline]
363    pub fn name(self) -> &'static str {
364        self.name_and_def().0
365    }
366
367    /// Lookup the definition for this operand kind in the lazily-loaded [`Spec`].
368    #[inline]
369    pub fn def(self) -> &'static OperandKindDef {
370        self.name_and_def().1
371    }
372}
373
374// HACK(eddyb) only needed because there are more than 256 unique operand names,
375// but less than 64 `OperandKind`s, so we can split 16 kind:name bits as 6:10.
376#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
377pub struct PackedOperandNameAndKind(u16);
378
379impl fmt::Debug for PackedOperandNameAndKind {
380    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
381        self.name_and_kind().fmt(f)
382    }
383}
384
385impl PackedOperandNameAndKind {
386    /// `operand_names[EMPTY_NAME_IDX]` must be reserved to contain `""`.
387    const EMPTY_NAME_IDX: usize = 0;
388
389    #[inline]
390    fn unnamed(kind: OperandKind) -> Self {
391        Self::pack(Self::EMPTY_NAME_IDX, kind)
392    }
393
394    #[inline]
395    fn pack(name_idx: usize, kind: OperandKind) -> Self {
396        let packed = Self(((name_idx as u16) << 6) | (kind.0 as u16));
397        assert_eq!(packed.unpack(), (name_idx, kind));
398        packed
399    }
400
401    #[inline]
402    fn unpack(self) -> (usize, OperandKind) {
403        ((self.0 >> 6) as usize, OperandKind((self.0 & ((1 << 6) - 1)) as u8))
404    }
405
406    /// Unpack this `PackedOperandNameAndKind` into just its `OperandKind`.
407    #[inline]
408    pub fn kind(self) -> OperandKind {
409        self.unpack().1
410    }
411
412    /// Unpack this `PackedOperandNameAndKind` into a name and `OperandKind`.
413    #[inline]
414    pub fn name_and_kind(self) -> (&'static str, OperandKind) {
415        let (name_idx, kind) = self.unpack();
416        (Spec::get().operand_names.get_index(name_idx).unwrap(), kind)
417    }
418}
419
420pub enum OperandKindDef {
421    BitEnum {
422        empty_name: &'static str,
423        bits: indexed::NamedIdxMap<BitIdx, Enumerant, indexed::FlatWithHoles>,
424    },
425
426    ValueEnum {
427        variants: indexed::NamedIdxMap<u16, Enumerant, indexed::KhrSegmented>,
428    },
429
430    Id,
431    Literal {
432        size: LiteralSize,
433    },
434}
435
436#[derive(Copy, Clone, PartialEq, Eq)]
437pub struct BitIdx(pub u8);
438
439impl BitIdx {
440    /// Returns `Some(BitIdx(i))` if and only if `x == (1 << i)`.
441    pub fn of_single_set_bit(x: u32) -> Option<Self> {
442        if x.is_power_of_two() { Some(Self(x.trailing_zeros() as u8)) } else { None }
443    }
444
445    /// Returns an iterator of [`BitIdx`]s, from which `x` can be reconstructed
446    /// by OR-ing together `1 << i` for every `BitIdx(i)`.
447    ///
448    /// The iterator is ordered: lower bit indices appear before higher ones.
449    pub fn of_all_set_bits(mut x: u32) -> impl Iterator<Item = Self> {
450        let mut consumed_bits = 0;
451        iter::from_fn(move || {
452            if x == 0 {
453                None
454            } else {
455                let tz = x.trailing_zeros() as u8;
456                let idx = Self(consumed_bits + tz);
457
458                // Consume a sequence of bits `100...00`, where `tz` is just the
459                // count of zeros, so `tz + 1` is the whole sequence's length.
460                x >>= tz + 1;
461                consumed_bits += tz + 1;
462
463                Some(idx)
464            }
465        })
466    }
467}
468
469impl indexed::FlatIdx for BitIdx {
470    fn to_usize(self) -> usize {
471        self.0.into()
472    }
473}
474
475#[derive(PartialEq, Eq)]
476pub struct Enumerant {
477    pub req_params: ArrayVec<PackedOperandNameAndKind, 3>,
478    pub rest_params: Option<OperandKind>,
479}
480
481impl Enumerant {
482    /// Return a (potentially infinite) iterator of [`OperandKind`]s, along with
483    /// the [`OperandMode`] indicating whether an operand is expected (`Required`),
484    /// or that an operand's absence signals the end of operands (`Optional`),
485    /// which is also the exit signal for the "rest operands" infinite iterators.
486    pub fn all_params(&self) -> impl Iterator<Item = (OperandMode, OperandKind)> + '_ {
487        self.all_params_with_names().map(|(mode, name_and_kind)| (mode, name_and_kind.kind()))
488    }
489
490    /// Like `all_params`, but providing access to the operand names as well.
491    pub fn all_params_with_names(
492        &self,
493    ) -> impl Iterator<Item = (OperandMode, PackedOperandNameAndKind)> + '_ {
494        self.req_params.iter().copied().map(|kind| (OperandMode::Required, kind)).chain(
495            self.rest_params.into_iter().flat_map(|kind| {
496                iter::repeat((OperandMode::Optional, PackedOperandNameAndKind::unnamed(kind)))
497            }),
498        )
499    }
500}
501
502pub enum LiteralSize {
503    /// The literal is always one word (but may occupy only part of it).
504    Word,
505
506    /// The literal is a word-encoded byte array, that ends with a `0` byte.
507    NulTerminated,
508
509    /// The literal uses as many words as required by its type, which is known
510    /// contextually (`OpConstant`'s result type or `OpSwitch`'s selector type).
511    FromContextualType,
512}
513
514fn sanitize_operand_name<'a>(name: &Option<raw::CowStr<'a>>) -> &'a str {
515    name.as_ref()
516        .and_then(|name| match name {
517            &raw::CowStr::Borrowed(s) => {
518                s.strip_prefix('\'')?.strip_suffix('\'').filter(|s| {
519                    // HACK(eddyb) it's pretty bad that SPIR-V uses spaces
520                    // in operand names, but by constraining the rest of
521                    // the character set (to be identifier-like), we get
522                    // to remove spaces (to get `FooBar`), or even replace
523                    // them with `_` (to get `Foo_Bar` or even `foo_bar`).
524                    s.starts_with(|c: char| c.is_ascii_alphabetic())
525                        && s.chars().all(|c| c.is_ascii_alphanumeric() || c == ' ')
526                })
527            }
528            raw::CowStr::Owned(s) => {
529                assert!(s.contains("', +\n'"), "unexpected non-zero-copy {s:?}");
530                None
531            }
532        })
533        .unwrap_or("")
534}
535
536// HACK(eddyb) make sure parsing JSON doesn't start failing randomly.
537#[test]
538fn get_spec_and_all_ext_inst_sets() {
539    let spec = Spec::get();
540    for name in spec.ext_inst_sets.keys() {
541        spec.get_ext_inst_set_by_lowercase_name(name);
542    }
543}
544
545impl Spec {
546    /// Return a lazily-loaded [`Spec`] (only does significant work for the first call).
547    #[inline(always)]
548    #[must_use]
549    pub fn get() -> &'static Spec {
550        lazy_static! {
551            static ref SPEC: Spec = {
552                mod khr_spv_grammar_jsons {
553                    include!(concat!(env!("OUT_DIR"), "/khr_spv_grammar_jsons.rs"));
554                }
555
556                let raw_core_grammar: raw::CoreGrammar<'static> =
557                    serde_json::from_str(khr_spv_grammar_jsons::SPIRV_CORE_GRAMMAR).unwrap();
558
559                let mut spec = Spec::from_raw(raw_core_grammar);
560
561                // FIXME(eddyb) this should be moved somewhere better.
562                for (name, json) in khr_spv_grammar_jsons::EXTINST_NAMES_AND_GRAMMARS {
563                    let lazy_init = move || {
564                        let is_debuginfo_ext_inst_set = name.contains(".debuginfo.");
565                        let extinst_grammar: raw::ExtInstGrammar<'static> =
566                            serde_json::from_str(json).unwrap();
567                        let instructions = extinst_grammar
568                            .instructions
569                            .iter()
570                            .map(|inst| {
571                                (
572                                    inst.opcode.into(),
573                                    ExtInstSetInstructionDesc {
574                                        name: inst.opname.into(),
575                                        operand_names: inst
576                                            .operands
577                                            .iter()
578                                            .map(|operand| {
579                                                sanitize_operand_name(&operand.name)
580                                            })
581                                            .take_while(|name| !name.is_empty())
582                                            .map(|name| name.into())
583                                            .collect(),
584
585                                        is_debuginfo: is_debuginfo_ext_inst_set
586                                            && inst.opname.strip_prefix("Debug")
587                                                .is_some_and(|next| next.starts_with(|c: char| c.is_ascii_uppercase()))
588                                    },
589                                )
590                            })
591                            .collect::<BTreeMap<_, _>>();
592                        ExtInstSetDesc { short_alias: None, instructions }
593                    };
594                    spec.ext_inst_sets.insert(
595                        name,
596                        (
597                            Default::default(),
598                            Box::new(lazy_init),
599                        ),
600                    );
601                }
602
603                spec
604            };
605        }
606        &SPEC
607    }
608
609    /// Return a lazily-parsed [`ExtInstSetDesc`], if a known one exists for this
610    /// `OpExtInstImport` name (required to be lowercase, due to Khronos' choice
611    /// of case insensitivity, but **not checked by this function**).
612    pub fn get_ext_inst_set_by_lowercase_name(
613        &self,
614        lowercase_ext_inst_set_name: &str,
615    ) -> Option<&ExtInstSetDesc> {
616        self.ext_inst_sets
617            .get(lowercase_ext_inst_set_name)
618            .map(|(once_cell, init)| once_cell.get_or_init(init))
619    }
620
621    /// Implementation detail of [`Spec::get`], indexes the raw data to produce a [`Spec`].
622    fn from_raw(raw_core_grammar: raw::CoreGrammar<'static>) -> Self {
623        /// Helper for picking a name when the same index has multiple names.
624        fn preferred_name_between_dups<'a>(a: &'a str, b: &'a str) -> &'a str {
625            // Prefer standard / Khronos extensions over vendor extensions.
626            let is_khr_and_vnd = |khr: &str, vnd: &str| {
627                let base = khr.trim_end_matches("KHR");
628                vnd.starts_with(base) && vnd.len() > base.len()
629            };
630            if is_khr_and_vnd(a, b) {
631                a
632            } else if is_khr_and_vnd(b, a) {
633                b
634            } else {
635                // Worst case, use the first in alphabetical order.
636                a.min(b)
637            }
638        }
639
640        // HACK(eddyb) ad-hoc interning, to reduce the cost of tracking operand names
641        // down to a single extra byte per operand (see `PackedOperandNameAndKind`).
642        let mut operand_names = FxIndexSet::default();
643        assert_eq!(operand_names.insert_full("").0, PackedOperandNameAndKind::EMPTY_NAME_IDX);
644        let mut pack_operand_name_and_kind = |name: &Option<raw::CowStr<'static>>, kind| {
645            let (name_idx, _) = operand_names.insert_full(sanitize_operand_name(name));
646            PackedOperandNameAndKind::pack(name_idx, kind)
647        };
648
649        // Constructing the full `OperandKindDef` may require looking up other
650        // `OperandKind`s by name, so build the lookup table for that up-front.
651        let operand_kind_by_name: FxHashMap<_, _> = raw_core_grammar
652            .operand_kinds
653            .iter()
654            .filter(|o| !matches!(o.category, raw::OperandKindCategory::Composite))
655            .enumerate()
656            .map(|(i, o)| (o.kind, OperandKind(i.try_into().unwrap())))
657            .collect();
658
659        let operand_kinds: Vec<_> = raw_core_grammar
660            .operand_kinds
661            .iter()
662            .filter_map(|o| {
663                let mut enumerant_from_raw = |e: &raw::OperandKindEnumerant<'static>| {
664                    let mut all_params = e
665                        .parameters
666                        .iter()
667                        .map(|p| (&p.name, &p.quantifier, operand_kind_by_name[p.kind]));
668
669                    let rest_params = match all_params.clone().next_back() {
670                        Some((_, Some(raw::Quantifier::Rest), kind)) => {
671                            all_params.next_back();
672                            Some(kind)
673                        }
674                        _ => None,
675                    };
676
677                    let mut req_params = ArrayVec::new();
678                    for (name, quantifier, kind) in all_params {
679                        assert!(quantifier.is_none());
680                        req_params
681                            .try_push(pack_operand_name_and_kind(name, kind))
682                            .map_err(|err| format!("{}/{name:?}: {err}", o.kind))
683                            .unwrap();
684                    }
685
686                    Enumerant {
687                        req_params,
688                        rest_params,
689                    }
690                };
691
692                let def = match o.category {
693                    raw::OperandKindCategory::BitEnum => {
694                        assert!(o.bases.is_none());
695
696                        let enumerants = o.enumerants.as_ref().unwrap();
697                        let mut empty_name = None;
698                        let mut bits = vec![];
699                        for e in enumerants {
700                            let new_name = e.enumerant;
701
702                            // `BitEnum` enumerants with `"value" : "0x0000"`
703                            // are only really provided to give a canonical name
704                            // to the state with no bits set (usually `"None"`).
705                            if e.value == 0 {
706                                assert!(e.parameters.is_empty());
707
708                                empty_name = Some(match empty_name {
709                                    None => new_name,
710                                    Some(prev_name) => {
711                                        preferred_name_between_dups(prev_name, new_name)
712                                    }
713                                });
714
715                                continue;
716                            }
717
718                            let new_enumerant = enumerant_from_raw(e);
719
720                            let bit_idx = BitIdx::of_single_set_bit(e.value).unwrap();
721
722                            // Make room for our new value, if necessary.
723                            let i = bit_idx.to_usize();
724                            if i >= bits.len() {
725                                bits.resize_with(i + 1, || None);
726                            }
727                            let slot = &mut bits[i];
728
729                            *slot = Some(match slot.take() {
730                                None => (new_name, new_enumerant),
731                                Some((prev_name, prev_enumerant)) => {
732                                    // Only allow aliases that do not meaningfully differ.
733                                    assert!(
734                                        prev_enumerant == new_enumerant,
735                                        "{} bits {} and {} share a bit index but differ in definition",
736                                        o.kind,
737                                        prev_name,
738                                        new_name,
739                                    );
740
741                                    (
742                                        preferred_name_between_dups(prev_name, new_name),
743                                        new_enumerant,
744                                    )
745                                }
746                            });
747                        }
748
749                        // FIXME(eddyb) automate this in `indexed::NamedIdxMap`.
750                        let bits = indexed::NamedIdxMap {
751                            idx_by_name: enumerants
752                                .iter()
753                                .filter_map(|e| {
754                                    Some((e.enumerant, BitIdx::of_single_set_bit(e.value)?))
755                                })
756                                .collect(),
757                            storage: bits,
758                        };
759
760                        OperandKindDef::BitEnum {
761                            empty_name: empty_name.unwrap_or("None"),
762                            bits,
763                        }
764                    }
765                    raw::OperandKindCategory::ValueEnum => {
766                        assert!(o.bases.is_none());
767
768                        let enumerants = o.enumerants.as_ref().unwrap();
769                        let variants = indexed::KhrSegmentedVec::from_in_order_iter(
770                            enumerants.iter().map(|e| {
771                                (
772                                    e.value.try_into().unwrap(),
773                                    (e.enumerant, enumerant_from_raw(e)),
774                                )
775                            }),
776                            // `merge_duplicates` closure:
777                            |(prev_name, prev_enumerant), (new_name, new_enumerant)| {
778                                // Only allow aliases that do not meaningfully differ.
779                                assert!(
780                                    prev_enumerant == new_enumerant,
781                                    "{} variants {} and {} share a value but differ in definition",
782                                    o.kind,
783                                    prev_name,
784                                    new_name,
785                                );
786
787                                (
788                                    preferred_name_between_dups(prev_name, new_name),
789                                    new_enumerant,
790                                )
791                            },
792                        );
793
794                        // FIXME(eddyb) automate this in `indexed::NamedIdxMap`.
795                        let variants = indexed::NamedIdxMap {
796                            idx_by_name: enumerants
797                                .iter()
798                                .map(|e| (e.enumerant, e.value.try_into().unwrap()))
799                                .collect(),
800                            storage: variants,
801                        };
802
803                        OperandKindDef::ValueEnum { variants }
804                    }
805                    raw::OperandKindCategory::Id => {
806                        assert!(o.enumerants.is_none() && o.bases.is_none());
807                        OperandKindDef::Id
808                    }
809                    raw::OperandKindCategory::Literal => {
810                        assert!(o.enumerants.is_none() && o.bases.is_none());
811                        let size = match o.kind {
812                            "LiteralInteger"
813                            | "LiteralExtInstInteger"
814                            | "LiteralSpecConstantOpInteger"
815                            | "LiteralFloat" => LiteralSize::Word,
816                            "LiteralString" => LiteralSize::NulTerminated,
817                            "LiteralContextDependentNumber" => LiteralSize::FromContextualType,
818                            _ => unreachable!(),
819                        };
820                        OperandKindDef::Literal { size }
821                    }
822                    raw::OperandKindCategory::Composite => {
823                        return None;
824                    }
825                };
826                Some((o.kind, def))
827            })
828            .collect();
829
830        // FIXME(eddyb) automate this in `indexed::NamedIdxMap`.
831        assert_eq!(operand_kind_by_name.len(), operand_kinds.len());
832        let operand_kinds =
833            indexed::NamedIdxMap { idx_by_name: operand_kind_by_name, storage: operand_kinds };
834
835        let operand_kind_pairs_by_name: FxHashMap<_, _> = raw_core_grammar
836            .operand_kinds
837            .iter()
838            .filter(|o| matches!(o.category, raw::OperandKindCategory::Composite))
839            .map(|o| {
840                assert!(o.enumerants.is_none());
841                let mut bases: [_; 2] = o.bases.as_ref().unwrap()[..].try_into().unwrap();
842
843                // HACK(eddyb) work around https://github.com/KhronosGroup/SPIRV-Headers/issues/38.
844                if o.kind == "PairLiteralIntegerIdRef" {
845                    assert_eq!(bases, ["LiteralInteger", "IdRef"]);
846                    bases[0] = "LiteralContextDependentNumber";
847                }
848
849                (o.kind, [
850                    operand_kinds.lookup(bases[0]).unwrap(),
851                    operand_kinds.lookup(bases[1]).unwrap(),
852                ])
853            })
854            .collect();
855
856        let id_result_type = operand_kinds.lookup("IdResultType").unwrap();
857        let id_result = operand_kinds.lookup("IdResult").unwrap();
858
859        let instructions = indexed::KhrSegmentedVec::from_in_order_iter(
860            raw_core_grammar.instructions.iter().map(|inst| {
861                // Helper for checking if `inst.opname` starts with `prefix`
862                // followed by an uppercase letter indicating the start of
863                // the first "word" for the intra-category instruction name.
864                let has_categorical_prefix = |prefix| {
865                    inst.opname
866                        .strip_prefix(prefix)
867                        .is_some_and(|next| next.starts_with(|c: char| c.is_ascii_uppercase()))
868                };
869
870                let category_from_prefix = if has_categorical_prefix("OpType") {
871                    Some(InstructionCategory::Type)
872                } else if matches!(inst.opname, "OpConstant" | "OpSpecConstant")
873                    || has_categorical_prefix("OpConstant")
874                    || has_categorical_prefix("OpSpecConstant")
875                {
876                    Some(InstructionCategory::Const)
877                } else if has_categorical_prefix("OpIgnore")
878                    || has_categorical_prefix("OpTerminate")
879                    || inst.opname == "OpEmitMeshTasksEXT"
880                {
881                    // HACK(eddyb) not category prefixes, but they help with
882                    // working around `Reserved` extensions with control-flow
883                    // instructions. False positives will be caught by the
884                    // assert further down, if `category_from_class` differs.
885                    Some(InstructionCategory::ControlFlow)
886                } else {
887                    None
888                };
889                let category_from_class = match inst.class {
890                    "Type-Declaration" => Some(InstructionCategory::Type),
891                    "Constant-Creation" => Some(InstructionCategory::Const),
892                    "Control-Flow" => Some(InstructionCategory::ControlFlow),
893
894                    // HACK(eddyb) work around all pipe instructions being in
895                    // the `Pipe` class, even when e.g. `Constant-Creation`
896                    // would be more appropriate (for `OpConstantPipeStorage`).
897                    "Pipe" => category_from_prefix.filter(|&category| {
898                        assert_eq!(
899                            (inst.opname, category),
900                            ("OpConstantPipeStorage", InstructionCategory::Const)
901                        );
902                        true
903                    }),
904
905                    // HACK(eddyb) work around extensions getting initially
906                    // added to catch-all classes like `Reserved` or `@exclude`.
907                    "Reserved" | "@exclude" => category_from_prefix,
908
909                    _ => None,
910                };
911                match (category_from_prefix, category_from_class) {
912                    // Control-flow instructions don't (all) have prefixes.
913                    (None, Some(InstructionCategory::ControlFlow)) => {}
914
915                    _ => assert!(
916                        category_from_prefix == category_from_class,
917                        "instruction name `{}` implies category `{:?}`, \
918                         but class `{}` implies category `{:?}`",
919                        inst.opname,
920                        category_from_prefix,
921                        inst.class,
922                        category_from_class,
923                    ),
924                }
925
926                let mut def = InstructionDef {
927                    // FIXME(eddyb) should `Other` be replaced with `Option`?
928                    category: category_from_class.unwrap_or(InstructionCategory::Other),
929
930                    has_result_type_id: false,
931                    has_result_id: false,
932
933                    req_operands: ArrayVec::new(),
934                    opt_operands: ArrayVec::new(),
935                    rest_operands: None,
936                };
937
938                #[derive(Copy, Clone, Debug, PartialEq, PartialOrd)]
939                enum Seq {
940                    IdResultType,
941                    IdResult,
942                    Required,
943                    Optional,
944                    Rest,
945                }
946                let mut seq = None;
947
948                for o in &inst.operands {
949                    let single = operand_kinds.lookup(o.kind);
950
951                    let next_seq = match o.quantifier {
952                        _ if single == Some(id_result_type) => {
953                            assert!(o.quantifier.is_none());
954                            assert!(!def.has_result_type_id);
955                            def.has_result_type_id = true;
956                            Seq::IdResultType
957                        }
958                        _ if single == Some(id_result) => {
959                            assert!(o.quantifier.is_none());
960                            assert!(!def.has_result_id);
961                            def.has_result_id = true;
962                            Seq::IdResult
963                        }
964                        None => {
965                            def.req_operands
966                                .try_push(pack_operand_name_and_kind(&o.name, single.unwrap()))
967                                .map_err(|err| format!("{}/{:?}: {err}", inst.opname, o.name))
968                                .unwrap();
969                            Seq::Required
970                        }
971                        Some(raw::Quantifier::Optional) => {
972                            def.opt_operands
973                                .try_push(pack_operand_name_and_kind(&o.name, single.unwrap()))
974                                .map_err(|err| format!("{}/{:?}: {err}", inst.opname, o.name))
975                                .unwrap();
976                            Seq::Optional
977                        }
978                        Some(raw::Quantifier::Rest) => {
979                            def.rest_operands = Some(match single {
980                                Some(kind) => RestOperandsUnit::One(kind),
981                                None => RestOperandsUnit::Two(operand_kind_pairs_by_name[o.kind]),
982                            });
983                            Seq::Rest
984                        }
985                    };
986                    assert!(seq <= Some(next_seq), "{next_seq:?} -> {seq:?}");
987                    seq = Some(next_seq);
988                }
989
990                // `IdResultType` without `IdResult` is impossible.
991                if def.has_result_type_id {
992                    assert!(def.has_result_id);
993                }
994
995                (inst.opcode, (inst.opname, def))
996            }),
997            // `merge_duplicates` closure:
998            |(prev_name, prev_def), (new_name, new_def)| {
999                // Only allow aliases that do not meaningfully differ.
1000                assert!(
1001                    prev_def == new_def,
1002                    "instructions {prev_name} and {new_name} share an opcode but differ in definition",
1003                );
1004
1005                (preferred_name_between_dups(prev_name, new_name), new_def)
1006            },
1007        );
1008
1009        // FIXME(eddyb) automate this in `indexed::NamedIdxMap`.
1010        let instructions = indexed::NamedIdxMap {
1011            idx_by_name: raw_core_grammar
1012                .instructions
1013                .iter()
1014                .map(|inst| (inst.opname, Opcode(inst.opcode)))
1015                .collect(),
1016            storage: instructions,
1017        };
1018
1019        let addressing_models =
1020            match &operand_kinds[operand_kinds.lookup("AddressingModel").unwrap()] {
1021                OperandKindDef::ValueEnum { variants } => variants,
1022                _ => unreachable!(),
1023            };
1024        let storage_classes = match &operand_kinds[operand_kinds.lookup("StorageClass").unwrap()] {
1025            OperandKindDef::ValueEnum { variants } => variants,
1026            _ => unreachable!(),
1027        };
1028        let decorations = match &operand_kinds[operand_kinds.lookup("Decoration").unwrap()] {
1029            OperandKindDef::ValueEnum { variants } => variants,
1030            _ => unreachable!(),
1031        };
1032        let linkage_types = match &operand_kinds[operand_kinds.lookup("LinkageType").unwrap()] {
1033            OperandKindDef::ValueEnum { variants } => variants,
1034            _ => unreachable!(),
1035        };
1036
1037        // FIXME(eddyb) if this is computed earlier, `IdResultType` and `IdResult`
1038        // wouldn't be looked up twice - but for now, this is mildly cleaner.
1039        let well_known = WellKnown::lookup_with(PerWellKnownGroup {
1040            opcode: |name| instructions.lookup(name).unwrap(),
1041            operand_kind: |name| operand_kinds.lookup(name).unwrap(),
1042            addressing_model: |name| addressing_models.lookup(name).unwrap().into(),
1043            storage_class: |name| storage_classes.lookup(name).unwrap().into(),
1044            decoration: |name| decorations.lookup(name).unwrap().into(),
1045            linkage_type: |name| linkage_types.lookup(name).unwrap().into(),
1046        });
1047
1048        Self {
1049            magic: raw_core_grammar.magic_number,
1050
1051            instructions,
1052            well_known,
1053            operand_kinds,
1054
1055            operand_names,
1056
1057            ext_inst_sets: BTreeMap::new(),
1058        }
1059    }
1060}
1061
1062/// Deserialization for the `.grammar.json` files, without any post-processing.
1063pub mod raw {
1064    use serde::Deserialize;
1065    use smallvec::SmallVec;
1066
1067    #[derive(Deserialize)]
1068    #[serde(deny_unknown_fields)]
1069    pub struct CoreGrammar<'a> {
1070        #[serde(borrow)]
1071        pub copyright: Vec<CowStr<'a>>,
1072
1073        #[serde(deserialize_with = "dew_u32_maybe_hex")]
1074        pub magic_number: u32,
1075
1076        pub major_version: u8,
1077        pub minor_version: u8,
1078        pub revision: u8,
1079
1080        pub instruction_printing_class: Vec<InstructionPrintingClass<'a>>,
1081        pub instructions: Vec<Instruction<'a>>,
1082        pub operand_kinds: Vec<OperandKind<'a>>,
1083    }
1084
1085    #[derive(Deserialize)]
1086    #[serde(deny_unknown_fields)]
1087    pub struct ExtInstGrammar<'a> {
1088        #[serde(borrow)]
1089        pub copyright: Option<Vec<CowStr<'a>>>,
1090
1091        pub version: Option<u8>,
1092        pub revision: u8,
1093
1094        pub instructions: Vec<Instruction<'a>>,
1095        #[serde(default)]
1096        pub operand_kinds: Vec<OperandKind<'a>>,
1097    }
1098
1099    #[derive(Deserialize)]
1100    #[serde(deny_unknown_fields)]
1101    pub struct InstructionPrintingClass<'a> {
1102        pub tag: &'a str,
1103        pub heading: Option<&'a str>,
1104    }
1105
1106    #[derive(Deserialize)]
1107    #[serde(deny_unknown_fields)]
1108    pub struct Instruction<'a> {
1109        pub opname: &'a str,
1110        #[serde(default)]
1111        pub class: &'a str,
1112        pub opcode: u16,
1113        #[serde(default)]
1114        pub operands: Vec<Operand<'a>>,
1115
1116        #[serde(default)]
1117        pub extensions: SmallVec<[&'a str; 1]>,
1118        #[serde(default)]
1119        pub capabilities: SmallVec<[&'a str; 1]>,
1120        // HACK(eddyb) some `extinst.*.json` use this form.
1121        pub capability: Option<&'a str>,
1122
1123        pub version: Option<&'a str>,
1124        #[serde(rename = "lastVersion")]
1125        pub last_version: Option<&'a str>,
1126    }
1127
1128    #[derive(Deserialize)]
1129    #[serde(deny_unknown_fields)]
1130    pub struct Operand<'a> {
1131        pub kind: &'a str,
1132        pub quantifier: Option<Quantifier>,
1133        #[serde(borrow)]
1134        pub name: Option<CowStr<'a>>,
1135    }
1136
1137    #[derive(Deserialize)]
1138    pub enum Quantifier {
1139        #[serde(rename = "?")]
1140        Optional,
1141
1142        #[serde(rename = "*")]
1143        Rest,
1144    }
1145
1146    #[derive(Deserialize)]
1147    #[serde(deny_unknown_fields)]
1148    pub struct OperandKind<'a> {
1149        pub category: OperandKindCategory,
1150        pub kind: &'a str,
1151        pub doc: Option<&'a str>,
1152
1153        pub enumerants: Option<Vec<OperandKindEnumerant<'a>>>,
1154
1155        pub bases: Option<Vec<&'a str>>,
1156    }
1157
1158    #[derive(Deserialize)]
1159    pub enum OperandKindCategory {
1160        BitEnum,
1161        ValueEnum,
1162
1163        Id,
1164        Literal,
1165        Composite,
1166    }
1167
1168    #[derive(Deserialize)]
1169    #[serde(deny_unknown_fields)]
1170    pub struct OperandKindEnumerant<'a> {
1171        pub enumerant: &'a str,
1172
1173        #[serde(deserialize_with = "dew_u32_maybe_hex")]
1174        pub value: u32,
1175
1176        #[serde(default)]
1177        pub parameters: Vec<Operand<'a>>,
1178
1179        #[serde(default)]
1180        pub extensions: SmallVec<[&'a str; 1]>,
1181        #[serde(default)]
1182        pub capabilities: SmallVec<[&'a str; 1]>,
1183
1184        pub version: Option<&'a str>,
1185        #[serde(rename = "lastVersion")]
1186        pub last_version: Option<&'a str>,
1187    }
1188
1189    // HACK(eddyb) `Cow<'a, str>` that works w/ zero-copy deserialization, even
1190    // when nested (`serde` only special-cases `Cow` used directly as a field type).
1191    #[derive(Deserialize, Debug)]
1192    #[serde(untagged)]
1193    pub enum CowStr<'a> {
1194        Borrowed(&'a str),
1195        Owned(String),
1196    }
1197
1198    /// Helper to generate functions usable with `deserialize_with` (hence "dew"),
1199    /// that deserialize to an intermediary type, which is then passed through the
1200    /// supplied closure, which is allowed to error. This is similar to the serde
1201    /// attribute `#[serde(try_from = "...")]`, but that only works for whole types.
1202    macro_rules! dew_and_then {
1203        ($($name:ident: |$x:ident: $in_ty:ty| -> $out_ty:ty $body:block),* $(,)?) => {
1204            $(fn $name<'de, D>(deserializer: D) -> Result<$out_ty, D::Error>
1205            where
1206                D: serde::Deserializer<'de>,
1207            {
1208                let x = Deserialize::deserialize(deserializer)?;
1209
1210                // HACK(eddyb) this is a `try {...}`-like use of a closure.
1211                #[allow(clippy::redundant_closure_call)]
1212                (|$x: $in_ty| -> Result<$out_ty, _> { $body })(x)
1213                    .map_err(serde::de::Error::custom)
1214            })*
1215        };
1216    }
1217
1218    dew_and_then! {
1219        dew_u32_maybe_hex: |x: DecOrHex<'_, u32>| -> u32 { x.try_into() },
1220    }
1221
1222    #[derive(Deserialize)]
1223    #[serde(untagged)]
1224    pub enum DecOrHex<'a, T> {
1225        Dec(T),
1226        MaybeHex(&'a str),
1227    }
1228
1229    impl TryInto<u32> for DecOrHex<'_, u32> {
1230        type Error = String;
1231        fn try_into(self) -> Result<u32, Self::Error> {
1232            match self {
1233                DecOrHex::Dec(x) => Ok(x),
1234                DecOrHex::MaybeHex(s) => {
1235                    // HACK(eddyb) some decimal numbers are kept as strings.
1236                    if let Ok(x) = s.parse() {
1237                        return Ok(x);
1238                    }
1239
1240                    s.strip_prefix("0x")
1241                        .ok_or_else(|| {
1242                            format!("DecOrHex string form doesn't start with 0x: {s:?}")
1243                        })?
1244                        .chars()
1245                        .try_fold(0u32, |r, c| {
1246                            // HACK(eddyb) this uses `checked_mul` because `checked_shl`
1247                            // doesn't handle information loss (bits being shifted off).
1248                            Ok(r.checked_mul(16).ok_or("DecOrHex hex overflow into u32")?
1249                                + c.to_digit(16)
1250                                    .ok_or("DecOrHex hex has non-hex-nibble character")?)
1251                        })
1252                }
1253            }
1254        }
1255    }
1256}
1257
1258/// Utilities for indexing data in a variety of ways (names, compact indices, etc.).
1259// FIXME(eddyb) move this out of here?
1260pub mod indexed {
1261    use rustc_hash::FxHashMap;
1262    use smallvec::SmallVec;
1263
1264    pub trait StorageShape<I, T> {
1265        type Storage;
1266        fn get_by_idx(storage: &Self::Storage, idx: I) -> Option<&T>;
1267    }
1268
1269    pub trait FlatIdx: Copy {
1270        fn to_usize(self) -> usize;
1271    }
1272
1273    impl FlatIdx for u16 {
1274        fn to_usize(self) -> usize {
1275            self.into()
1276        }
1277    }
1278
1279    /// Flat array ([`Vec`]) storage, likely used with compact indices.
1280    pub enum Flat {}
1281
1282    impl<I: FlatIdx, T> StorageShape<I, T> for Flat {
1283        type Storage = Vec<T>;
1284        fn get_by_idx(storage: &Self::Storage, idx: I) -> Option<&T> {
1285            storage.get(idx.to_usize())
1286        }
1287    }
1288
1289    /// Like [`Flat`], but the [`Vec`] elements are wrapped in [`Option`].
1290    pub enum FlatWithHoles {}
1291
1292    impl<I: FlatIdx, T> StorageShape<I, T> for FlatWithHoles {
1293        type Storage = Vec<Option<T>>;
1294        fn get_by_idx(storage: &Self::Storage, idx: I) -> Option<&T> {
1295            storage.get(idx.to_usize())?.as_ref()
1296        }
1297    }
1298
1299    /// Segmented sparse storage, taking advantage of Khronos' predictable
1300    /// reservation policy for SPIR-V instruction opcodes and `ValueEnum`s:
1301    /// * indices in `0..=4096` are reserved for the standard, and mostly
1302    ///   allocated without gaps (~84% density at the time of writing)
1303    /// * indices in `4096..` are allocated in blocks of `64`; while sparser
1304    ///   than the standard range, the blockiness allows some optimizations
1305    pub enum KhrSegmented {}
1306
1307    /// Khronos-oriented segmented sparse array (see [`KhrSegmented`]).
1308    pub struct KhrSegmentedVec<T> {
1309        /// Concatenation of values for indices lower than `4096`, with values
1310        /// for indices in a `64`-sized/aligned block starting at/above `4096`.
1311        ///
1312        /// Gaps are present (as `None`), but only if there are more values at
1313        /// some point after the gap, in the `0..=4096` index range, or in the
1314        /// same `64`-sized/aligned block (i.e. tailing gaps are elided).
1315        flattened: Vec<Option<T>>,
1316
1317        /// Starting indices in `flattened` for every `64`-sized/aligned block.
1318        ///
1319        /// For example, if an index `i >= 4096` is present, its value can be
1320        /// found at `flattened[block_starts[(i - 4096) / 64] + (i % 64)]`.
1321        block_starts: SmallVec<[u16; 8]>,
1322    }
1323
1324    impl<T> KhrSegmentedVec<T> {
1325        /// If `idx` is not in an out-of-range block, returns the pair of a
1326        /// "segment range" and an "intra-segment index".
1327        ///
1328        /// For example, if an index `i` is present, then `idx_to_segmented(i)`
1329        /// will return `Some((seg_range, intra_seg_idx))`, and the value can be
1330        /// found at `flattened[seg_range][intra_seg_idx]`.
1331        fn idx_to_segmented(&self, idx: u16) -> Option<(std::ops::Range<usize>, usize)> {
1332            let (block, intra_seg_idx) = if let Some(in_blocks_idx) = idx.checked_sub(4096) {
1333                (Some(usize::from(in_blocks_idx / 64)), idx % 64)
1334            } else {
1335                (None, idx)
1336            };
1337            let next_block = block.map_or(0, |b| b + 1);
1338
1339            let seg_start =
1340                block.map_or(Some(0), |b| self.block_starts.get(b).copied().map(usize::from))?;
1341            let seg_end = self
1342                .block_starts
1343                .get(next_block)
1344                .copied()
1345                .map_or(self.flattened.len(), usize::from);
1346
1347            Some((seg_start..seg_end, usize::from(intra_seg_idx)))
1348        }
1349
1350        /// Add a new value, with an index greater than all previous indices.
1351        ///
1352        /// An exception is made for duplicates, which have to be handled by the
1353        /// `merge_duplicates` closure, instead of being outright disallowed.
1354        fn insert_in_order(&mut self, idx: u16, value: T, merge_duplicates: impl Fn(T, T) -> T) {
1355            let last_idx_plus_one = self.block_starts.len().checked_sub(1).map_or(
1356                self.flattened.len(),
1357                |last_block_idx| {
1358                    4096 + 64 * last_block_idx
1359                        + (self.flattened.len() - usize::from(self.block_starts[last_block_idx]))
1360                },
1361            );
1362            if let Some(last_idx) = last_idx_plus_one.checked_sub(1) {
1363                // HACK(eddyb) the condition being `<` instead of `<=` allows
1364                // for special handling of duplicates (via `merge_duplicates`).
1365                if usize::from(idx) < last_idx {
1366                    panic!(
1367                        "KhrSegmentedVec::insert_in_order: out of order indices ({idx} after {last_idx})",
1368                    );
1369                }
1370            }
1371
1372            // Reserve new blocks if needed (so `idx_to_segmented` can't fail).
1373            if let Some(block) = idx.checked_sub(4096).map(|i| i / 64) {
1374                let needed_blocks = usize::from(block).checked_add(1).unwrap();
1375                if needed_blocks > self.block_starts.len() {
1376                    self.block_starts
1377                        .resize(needed_blocks, self.flattened.len().try_into().unwrap());
1378                }
1379            }
1380            let (seg_range, intra_seg_idx) = self.idx_to_segmented(idx).unwrap();
1381
1382            // The check at the start ensures we're never trying to insert in
1383            // an "already completed" segment.
1384            assert_eq!(seg_range.end, self.flattened.len());
1385
1386            let slot_idx = seg_range.start + intra_seg_idx;
1387            let needed_slots = slot_idx.checked_add(1).unwrap();
1388            if needed_slots > self.flattened.len() {
1389                self.flattened.resize_with(needed_slots, || None);
1390            }
1391            let slot = &mut self.flattened[slot_idx];
1392            if let Some(prev) = slot.take() {
1393                *slot = Some(merge_duplicates(prev, value));
1394            } else {
1395                *slot = Some(value);
1396            }
1397        }
1398
1399        /// Construct a [`KhrSegmentedVec`] out of an iterator with ordered indices.
1400        ///
1401        /// An exception is made for duplicates, which have to be handled by the
1402        /// `merge_duplicates` closure, instead of being outright disallowed.
1403        pub fn from_in_order_iter(
1404            it: impl IntoIterator<Item = (u16, T)>,
1405            merge_duplicates: impl Fn(T, T) -> T,
1406        ) -> Self {
1407            let iter = it.into_iter();
1408
1409            let mut this = Self {
1410                flattened: Vec::with_capacity(
1411                    iter.size_hint().0.checked_next_power_of_two().unwrap_or(0),
1412                ),
1413                block_starts: SmallVec::new(),
1414            };
1415
1416            for (idx, value) in iter {
1417                // FIXME(eddyb) the check at the start of `insert_in_order` may
1418                // be less efficient than if we checked the ordering here instead.
1419                this.insert_in_order(idx, value, &merge_duplicates);
1420            }
1421
1422            this
1423        }
1424    }
1425
1426    impl<I: FlatIdx, T> StorageShape<I, T> for KhrSegmented {
1427        type Storage = KhrSegmentedVec<T>;
1428        fn get_by_idx(storage: &Self::Storage, idx: I) -> Option<&T> {
1429            let (seg_range, intra_seg_idx) =
1430                storage.idx_to_segmented(idx.to_usize().try_into().ok()?)?;
1431
1432            storage.flattened.get(seg_range)?.get(intra_seg_idx)?.as_ref()
1433        }
1434    }
1435
1436    pub struct NamedIdxMap<I, T, S: StorageShape<I, (&'static str, T)>> {
1437        pub(super) idx_by_name: FxHashMap<&'static str, I>,
1438        pub(super) storage: S::Storage,
1439    }
1440
1441    impl<I, T, S: StorageShape<I, (&'static str, T)>> NamedIdxMap<I, T, S> {
1442        /// Get an index from a name.
1443        pub fn lookup(&self, name: &str) -> Option<I>
1444        where
1445            I: Copy,
1446        {
1447            self.idx_by_name.get(name).copied()
1448        }
1449
1450        pub fn get_named(&self, idx: I) -> Option<(&'static str, &T)> {
1451            let (name, value) = S::get_by_idx(&self.storage, idx)?;
1452            Some((name, value))
1453        }
1454
1455        pub fn get(&self, idx: I) -> Option<&T> {
1456            let (_name, value) = self.get_named(idx)?;
1457            Some(value)
1458        }
1459    }
1460
1461    impl<I, T, S: StorageShape<I, (&'static str, T)>> std::ops::Index<I> for NamedIdxMap<I, T, S> {
1462        type Output = T;
1463        fn index(&self, idx: I) -> &T {
1464            self.get(idx).unwrap()
1465        }
1466    }
1467}