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

rustc_codegen_spirv/
custom_decorations.rs
