rustc_codegen_spirv/
custom_decorations.rs

1//! SPIR-V decorations specific to `rustc_codegen_spirv`, produced during
2//! the original codegen of a crate, and consumed by the `linker`.
3
4use crate::builder_spirv::BuilderSpirv;
5use crate::custom_insts::{self, CustomInst};
6use either::Either;
7use rspirv::dr::{Instruction, Module, Operand};
8use rspirv::spirv::{Decoration, Op, Word};
9use rustc_data_structures::fx::FxIndexMap;
10use rustc_span::{FileName, SourceFile};
11use rustc_span::{Span, source_map::SourceMap};
12use smallvec::SmallVec;
13use std::borrow::Cow;
14use std::marker::PhantomData;
15use std::sync::Arc;
16use std::{fmt, iter, slice, str};
17
18/// Decorations not native to SPIR-V require some form of encoding into existing
19/// SPIR-V constructs, for which we use `OpDecorateString` with decoration type
20/// `UserTypeGOOGLE` and some encoded Rust value as the decoration string.
21///
22/// Each decoration type has to implement this trait, and use a different
23/// `ENCODING_PREFIX` from any other decoration type, to disambiguate them.
24///
25/// Also, all decorations have to be stripped by the linker at some point,
26/// ideally as soon as they're no longer needed, because no other tools
27/// processing the SPIR-V would understand them correctly.
28///
29/// TODO: uses `non_semantic` instead of piggybacking off of `UserTypeGOOGLE`
30/// <https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/KHR/SPV_KHR_non_semantic_info.html>
31pub trait CustomDecoration<'a>: Sized {
32    const ENCODING_PREFIX: &'static str;
33
34    fn encode(self, w: &mut impl fmt::Write) -> fmt::Result;
35    fn decode(s: &'a str) -> Self;
36
37    fn encode_to_inst(self, id: Word) -> Instruction {
38        let mut encoded = Self::ENCODING_PREFIX.to_string();
39        self.encode(&mut encoded).unwrap();
40
41        Instruction::new(
42            Op::DecorateString,
43            None,
44            None,
45            vec![
46                Operand::IdRef(id),
47                Operand::Decoration(Decoration::UserTypeGOOGLE),
48                Operand::LiteralString(encoded),
49            ],
50        )
51    }
52
53    fn try_decode_from_inst(inst: &Instruction) -> Option<(Word, LazilyDecoded<'_, Self>)> {
54        if inst.class.opcode == Op::DecorateString
55            && inst.operands[1].unwrap_decoration() == Decoration::UserTypeGOOGLE
56        {
57            let id = inst.operands[0].unwrap_id_ref();
58            let prefixed_encoded = inst.operands[2].unwrap_literal_string();
59            let encoded = prefixed_encoded.strip_prefix(Self::ENCODING_PREFIX)?;
60
61            Some((
62                id,
63                LazilyDecoded {
64                    encoded,
65                    _marker: PhantomData,
66                },
67            ))
68        } else {
69            None
70        }
71    }
72
73    fn decode_all(module: &Module) -> DecodeAllIter<'_, Self> {
74        module
75            .annotations
76            .iter()
77            .filter_map(Self::try_decode_from_inst as fn(_) -> _)
78    }
79
80    fn remove_all(module: &mut Module) {
81        module
82            .annotations
83            .retain(|inst| Self::try_decode_from_inst(inst).is_none());
84    }
85}
86
87// HACK(eddyb) return type of `CustomDecoration::decode_all`, in lieu of
88// `-> impl Iterator<Item = (Word, LazilyDecoded<'_, Self>)` in the trait.
89type DecodeAllIter<'a, D> = iter::FilterMap<
90    slice::Iter<'a, Instruction>,
91    fn(&'a Instruction) -> Option<(Word, LazilyDecoded<'a, D>)>,
92>;
93
94/// Helper allowing full decoding to be avoided where possible.
95//
96// FIXME(eddyb) is this even needed? (decoding impls are now much cheaper)
97pub struct LazilyDecoded<'a, D> {
98    encoded: &'a str,
99    _marker: PhantomData<D>,
100}
101
102impl<'a, D: CustomDecoration<'a>> LazilyDecoded<'a, D> {
103    pub fn decode(&self) -> D {
104        D::decode(self.encoded)
105    }
106}
107
108pub struct ZombieDecoration<'a> {
109    pub reason: Cow<'a, str>,
110}
111
112impl<'a> CustomDecoration<'a> for ZombieDecoration<'a> {
113    const ENCODING_PREFIX: &'static str = "Z";
114
115    fn encode(self, w: &mut impl fmt::Write) -> fmt::Result {
116        let Self { reason } = self;
117        w.write_str(&reason)
118    }
119    fn decode(s: &'a str) -> Self {
120        Self { reason: s.into() }
121    }
122}
123
124/// Equivalent of `CustomInst::SetDebugSrcLoc` (see `crate::custom_insts`),
125/// for global definitions (i.e. outside functions), where limitations of
126/// `rspirv`/`spirt` prevent us from using anything other than decorations.
127//
128// NOTE(eddyb) `CustomInst::SetDebugSrcLoc` is modelled after `DebugLine` from
129// `NonSemantic.Shader.DebugInfo`, might be good to invest in SPIR-T being able
130// to use `NonSemantic.Shader.DebugInfo` directly, in all situations.
131#[derive(Copy, Clone)]
132pub struct SrcLocDecoration<'a> {
133    pub file_name: &'a str,
134    pub line_start: u32,
135    pub line_end: u32,
136    pub col_start: u32,
137    pub col_end: u32,
138}
139
140impl<'a> CustomDecoration<'a> for SrcLocDecoration<'a> {
141    const ENCODING_PREFIX: &'static str = "L";
142
143    fn encode(self, w: &mut impl fmt::Write) -> fmt::Result {
144        let Self {
145            file_name,
146            line_start,
147            line_end,
148            col_start,
149            col_end,
150        } = self;
151        write!(
152            w,
153            "{file_name}:{line_start}:{col_start}-{line_end}:{col_end}"
154        )
155    }
156    fn decode(s: &'a str) -> Self {
157        #[derive(Copy, Clone, Debug)]
158        struct InvalidSrcLoc<'a>(
159            // HACK(eddyb) only exists for `fmt::Debug` in case of error.
160            #[allow(dead_code)] &'a str,
161        );
162        let err = InvalidSrcLoc(s);
163
164        let (s, col_end) = s.rsplit_once(':').ok_or(err).unwrap();
165        let (s, line_end) = s.rsplit_once('-').ok_or(err).unwrap();
166        let (s, col_start) = s.rsplit_once(':').ok_or(err).unwrap();
167        let (s, line_start) = s.rsplit_once(':').ok_or(err).unwrap();
168        let file_name = s;
169
170        Self {
171            file_name,
172            line_start: line_start.parse().unwrap(),
173            line_end: line_end.parse().unwrap(),
174            col_start: col_start.parse().unwrap(),
175            col_end: col_end.parse().unwrap(),
176        }
177    }
178}
179
180impl<'tcx> SrcLocDecoration<'tcx> {
181    pub fn from_rustc_span(span: Span, builder: &BuilderSpirv<'tcx>) -> Option<Self> {
182        // We may not always have valid spans.
183        // FIXME(eddyb) reduce the sources of this as much as possible.
184        if span.is_dummy() {
185            return None;
186        }
187
188        let (file, line_col_range) = builder.file_line_col_range_for_debuginfo(span);
189        let ((line_start, col_start), (line_end, col_end)) =
190            (line_col_range.start, line_col_range.end);
191
192        Some(Self {
193            file_name: file.file_name,
194            line_start,
195            line_end,
196            col_start,
197            col_end,
198        })
199    }
200}
201
202/// Helper type to delay most of the work necessary to turn a `SrcLocDecoration`
203/// back into an usable `Span`, until it's actually needed (i.e. for an error).
204pub struct SpanRegenerator<'a> {
205    source_map: &'a SourceMap,
206    module: Either<&'a Module, &'a spirt::Module>,
207
208    src_loc_decorations: Option<FxIndexMap<Word, LazilyDecoded<'a, SrcLocDecoration<'a>>>>,
209
210    // HACK(eddyb) this has no really good reason to belong here, but it's easier
211    // to handle it together with `SrcLocDecoration`, than separately.
212    zombie_decorations: Option<FxIndexMap<Word, LazilyDecoded<'a, ZombieDecoration<'a>>>>,
213
214    // HACK(eddyb) this is mostly replicating SPIR-T's module-level debuginfo.
215    spv_debug_info: Option<SpvDebugInfo<'a>>,
216}
217
218#[derive(Default)]
219struct SpvDebugInfo<'a> {
220    /// ID of `OpExtInstImport` for our custom "extended instruction set",
221    /// if present (see `crate::custom_insts` for more details).
222    custom_ext_inst_set_import: Option<Word>,
223
224    // HACK(eddyb) this is only needed because `OpExtInst`s can't have immediates,
225    // and must resort to referencing `OpConstant`s instead.
226    id_to_op_constant_operand: FxIndexMap<Word, &'a Operand>,
227
228    id_to_op_string: FxIndexMap<Word, &'a str>,
229    files: FxIndexMap<&'a str, SpvDebugFile<'a>>,
230}
231
232impl<'a> SpvDebugInfo<'a> {
233    fn collect(module: Either<&'a Module, &'a spirt::Module>) -> Self {
234        let mut this = Self::default();
235
236        let module = match module {
237            Either::Left(module) => module,
238
239            // HACK(eddyb) the SPIR-T codepath is simpler, and kind of silly,
240            // but we need the `SpvDebugFile`'s `regenerated_rustc_source_file`
241            // caching, so for now it reuses `SpvDebugInfo` overall.
242            Either::Right(module) => {
243                let cx = module.cx_ref();
244                match &module.debug_info {
245                    spirt::ModuleDebugInfo::Spv(debug_info) => {
246                        for sources in debug_info.source_languages.values() {
247                            for (&file_name, src) in &sources.file_contents {
248                                // FIXME(eddyb) what if the file is already present,
249                                // should it be considered ambiguous overall?
250                                this.files
251                                    .entry(&cx[file_name])
252                                    .or_default()
253                                    .op_source_parts = [&src[..]].into_iter().collect();
254                            }
255                        }
256                    }
257                }
258                return this;
259            }
260        };
261
262        // FIXME(eddyb) avoid repeating this across different passes/helpers.
263        this.custom_ext_inst_set_import = module
264            .ext_inst_imports
265            .iter()
266            .find(|inst| {
267                assert_eq!(inst.class.opcode, Op::ExtInstImport);
268                inst.operands[0].unwrap_literal_string() == &custom_insts::CUSTOM_EXT_INST_SET[..]
269            })
270            .map(|inst| inst.result_id.unwrap());
271
272        this.id_to_op_constant_operand.extend(
273            module
274                .types_global_values
275                .iter()
276                .filter(|inst| inst.class.opcode == Op::Constant)
277                .map(|inst| (inst.result_id.unwrap(), &inst.operands[0])),
278        );
279
280        let mut insts = module.debug_string_source.iter().peekable();
281        while let Some(inst) = insts.next() {
282            match inst.class.opcode {
283                Op::String => {
284                    this.id_to_op_string.insert(
285                        inst.result_id.unwrap(),
286                        inst.operands[0].unwrap_literal_string(),
287                    );
288                }
289                Op::Source if inst.operands.len() == 4 => {
290                    let file_name_id = inst.operands[2].unwrap_id_ref();
291                    if let Some(&file_name) = this.id_to_op_string.get(&file_name_id) {
292                        let mut file = SpvDebugFile::default();
293                        file.op_source_parts
294                            .push(inst.operands[3].unwrap_literal_string());
295                        while let Some(&next_inst) = insts.peek() {
296                            if next_inst.class.opcode != Op::SourceContinued {
297                                break;
298                            }
299                            insts.next();
300
301                            file.op_source_parts
302                                .push(next_inst.operands[0].unwrap_literal_string());
303                        }
304
305                        // FIXME(eddyb) what if the file is already present,
306                        // should it be considered ambiguous overall?
307                        this.files.insert(file_name, file);
308                    }
309                }
310                _ => {}
311            }
312        }
313        this
314    }
315}
316
317// HACK(eddyb) this is mostly replicating SPIR-T's module-level debuginfo.
318#[derive(Default)]
319struct SpvDebugFile<'a> {
320    /// Source strings from one `OpSource`, and any number of `OpSourceContinued`.
321    op_source_parts: SmallVec<[&'a str; 1]>,
322
323    regenerated_rustc_source_file: Option<Arc<SourceFile>>,
324}
325
326impl<'a> SpanRegenerator<'a> {
327    pub fn new(source_map: &'a SourceMap, module: &'a Module) -> Self {
328        Self {
329            source_map,
330            module: Either::Left(module),
331
332            src_loc_decorations: None,
333            zombie_decorations: None,
334
335            spv_debug_info: None,
336        }
337    }
338
339    pub fn new_spirt(source_map: &'a SourceMap, module: &'a spirt::Module) -> Self {
340        Self {
341            source_map,
342            module: Either::Right(module),
343
344            src_loc_decorations: None,
345            zombie_decorations: None,
346
347            spv_debug_info: None,
348        }
349    }
350
351    pub fn src_loc_for_id(&mut self, id: Word) -> Option<SrcLocDecoration<'a>> {
352        self.src_loc_decorations
353            .get_or_insert_with(|| {
354                SrcLocDecoration::decode_all(self.module.left().unwrap()).collect()
355            })
356            .get(&id)
357            .map(|src_loc| src_loc.decode())
358    }
359
360    // HACK(eddyb) this has no really good reason to belong here, but it's easier
361    // to handle it together with `SrcLocDecoration`, than separately.
362    pub(crate) fn zombie_for_id(&mut self, id: Word) -> Option<ZombieDecoration<'a>> {
363        self.zombie_decorations
364            .get_or_insert_with(|| {
365                ZombieDecoration::decode_all(self.module.left().unwrap()).collect()
366            })
367            .get(&id)
368            .map(|zombie| zombie.decode())
369    }
370
371    /// Extract the equivalent `SrcLocDecoration` from a debug instruction that
372    /// specifies some source location (both the standard `OpLine`, and our own
373    /// custom instruction, i.e. `CustomInst::SetDebugSrcLoc`, are supported).
374    pub fn src_loc_from_debug_inst(&mut self, inst: &Instruction) -> Option<SrcLocDecoration<'a>> {
375        let spv_debug_info = self
376            .spv_debug_info
377            .get_or_insert_with(|| SpvDebugInfo::collect(self.module));
378
379        let (file_id, line_start, line_end, col_start, col_end) = match inst.class.opcode {
380            Op::Line => {
381                let file = inst.operands[0].unwrap_id_ref();
382                let line = inst.operands[1].unwrap_literal_bit32();
383                let col = inst.operands[2].unwrap_literal_bit32();
384                (file, line, line, col, col)
385            }
386            Op::ExtInst
387                if Some(inst.operands[0].unwrap_id_ref())
388                    == spv_debug_info.custom_ext_inst_set_import =>
389            {
390                match CustomInst::decode(inst) {
391                    CustomInst::SetDebugSrcLoc {
392                        file,
393                        line_start,
394                        line_end,
395                        col_start,
396                        col_end,
397                    } => {
398                        let const_u32 = |operand: Operand| {
399                            spv_debug_info.id_to_op_constant_operand[&operand.unwrap_id_ref()]
400                                .unwrap_literal_bit32()
401                        };
402                        (
403                            file.unwrap_id_ref(),
404                            const_u32(line_start),
405                            const_u32(line_end),
406                            const_u32(col_start),
407                            const_u32(col_end),
408                        )
409                    }
410                    custom_inst => {
411                        unreachable!("src_loc_from_debug_inst({inst:?} => {custom_inst:?})")
412                    }
413                }
414            }
415            _ => unreachable!("src_loc_from_debug_inst({inst:?})"),
416        };
417
418        spv_debug_info
419            .id_to_op_string
420            .get(&file_id)
421            .map(|&file_name| SrcLocDecoration {
422                file_name,
423                line_start,
424                line_end,
425                col_start,
426                col_end,
427            })
428    }
429
430    fn regenerate_rustc_source_file(&mut self, file_name: &str) -> Option<&SourceFile> {
431        let spv_debug_file = self
432            .spv_debug_info
433            .get_or_insert_with(|| SpvDebugInfo::collect(self.module))
434            .files
435            .get_mut(file_name)?;
436
437        let file = &mut spv_debug_file.regenerated_rustc_source_file;
438        if file.is_none() {
439            // FIXME(eddyb) reduce allocations here by checking if the file is
440            // already loaded, and not allocating just to compare the source,
441            // but at least it's cheap when `OpSourceContinued` isn't used.
442            let src = match &spv_debug_file.op_source_parts[..] {
443                &[part] => Cow::Borrowed(part),
444                parts => parts.concat().into(),
445            };
446
447            // HACK(eddyb) in case the file has changed, and because `SourceMap`
448            // is strictly monotonic, we need to come up with some other name.
449            let mut sm_file_name_candidates = [FileName::Custom(file_name.to_string())]
450                .into_iter()
451                .chain((0..).map(|i| FileName::Custom(format!("outdated({i}) {file_name}"))));
452
453            *file = sm_file_name_candidates.find_map(|sm_file_name_candidate| {
454                let sf = self
455                    .source_map
456                    .new_source_file(sm_file_name_candidate, src.clone().into_owned());
457
458                // Only use this `FileName` candidate if we either:
459                // 1. reused a `SourceFile` with the right `src`/`external_src`
460                // 2. allocated a new `SourceFile` with our choice of `src`
461                self.source_map.ensure_source_file_source_present(&sf);
462                let sf_src_matches = sf
463                    .src
464                    .as_ref()
465                    .map(|sf_src| sf_src[..] == src[..])
466                    .or_else(|| {
467                        sf.external_src
468                            .borrow()
469                            .get_source()
470                            .map(|sf_src| sf_src[..] == src[..])
471                    })
472                    .unwrap_or(false);
473
474                if sf_src_matches { Some(sf) } else { None }
475            });
476        }
477        file.as_deref()
478    }
479
480    pub fn src_loc_to_rustc(&mut self, src_loc: SrcLocDecoration<'_>) -> Option<Span> {
481        let SrcLocDecoration {
482            file_name,
483            line_start,
484            line_end,
485            col_start,
486            col_end,
487        } = src_loc;
488
489        let file = self.regenerate_rustc_source_file(file_name)?;
490
491        // FIXME(eddyb) avoid some of the duplicated work when this closure is
492        // called with `line`/`col` values that are near eachother - thankfully,
493        // this code should only be hit on the error reporting path anyway.
494        let line_col_to_bpos = |line: u32, col: u32| {
495            let line_idx_in_file = line.checked_sub(1)? as usize;
496            let line_bpos_range = file.line_bounds(line_idx_in_file);
497            let line_contents = file.get_line(line_idx_in_file)?;
498
499            // Increment the `BytePos` until we reach the right `col_display`.
500            let (mut cur_bpos, mut cur_col_display) = (line_bpos_range.start, 0);
501            let mut line_chars = line_contents.chars();
502            while cur_bpos < line_bpos_range.end && cur_col_display < col {
503                // Add each char's `BytePos` and `col_display` contributions.
504                let ch = line_chars.next()?;
505                cur_bpos.0 += ch.len_utf8() as u32;
506                cur_col_display += rustc_span::char_width(ch) as u32;
507            }
508            Some(cur_bpos)
509        };
510
511        Some(Span::with_root_ctxt(
512            line_col_to_bpos(line_start, col_start)?,
513            line_col_to_bpos(line_end, col_end)?,
514        ))
515    }
516}
rustc_codegen_spirv/custom_decorations.rs

rustc_codegen_spirv/
custom_decorations.rs
