spirt/spv/
print.rs

1//! Pretty-printing SPIR-V operands.
2
3use crate::spv::{self, spec};
4use smallvec::SmallVec;
5use std::borrow::Cow;
6use std::fmt::Write;
7use std::{iter, mem, str};
8
9/// The smallest unit produced by printing a ("logical") SPIR-V operand.
10///
11/// All variants other than `Id` contain a fully formatted string, and the
12/// distinction between variants can be erased to obtain a plain-text version
13/// (also except `OperandKindNamespacePrefix` requiring an extra implicit `.`).
14//
15// FIXME(eddyb) should there be a `TokenKind` enum and then `Cow<'static, str>`
16// paired with a `TokenKind` in place of all of these individual variants?
17pub enum Token<ID> {
18    /// An inconsistency was detected in the operands to be printed.
19    /// For stylistic consistency, the error message is always found wrapped in
20    /// a block comment (i.e. the [`String`] is always of the form `"/* ... */"`).
21    Error(String),
22
23    // NOTE(eddyb) this implies a suffix `: ` not included in the string, and
24    // optionally some processing of the name (e.g. removing spaces).
25    OperandName(&'static str),
26
27    // FIXME(eddyb) perhaps encode the hierarchical structure of e.g. enumerand
28    // parameters, so that the SPIR-T printer can do layout for them.
29    Punctuation(&'static str),
30
31    // NOTE(eddyb) this implies a suffix `.` not included in the string.
32    OperandKindNamespacePrefix(&'static str),
33
34    EnumerandName(&'static str),
35
36    NumericLiteral(String),
37    StringLiteral(String),
38
39    /// Unprinted ID operand, of its original type (allowing post-processing).
40    Id(ID),
41}
42
43/// All the [`Token`]s outputted by printing one single ("logical") SPIR-V operand,
44/// which may be concatenated (after separately processing `ID`s) to obtain a
45/// complete plain-text version of the printed operand.
46pub struct TokensForOperand<ID> {
47    pub tokens: SmallVec<[Token<ID>; 3]>,
48}
49
50impl<ID> Default for TokensForOperand<ID> {
51    fn default() -> Self {
52        Self { tokens: SmallVec::new() }
53    }
54}
55
56impl TokensForOperand<String> {
57    pub fn concat_to_plain_text(self) -> String {
58        self.tokens
59            .into_iter()
60            .flat_map(|token| {
61                let (first, second): (Cow<'_, str>, _) = match token {
62                    Token::OperandName(s) => (s.into(), Some(": ".into())),
63                    Token::OperandKindNamespacePrefix(s) => (s.into(), Some(".".into())),
64                    Token::Punctuation(s) | Token::EnumerandName(s) => (s.into(), None),
65                    Token::Error(s)
66                    | Token::NumericLiteral(s)
67                    | Token::StringLiteral(s)
68                    | Token::Id(s) => (s.into(), None),
69                };
70                [first].into_iter().chain(second)
71            })
72            .reduce(|out, extra| (out.into_owned() + &extra).into())
73            .unwrap_or_default()
74            .into_owned()
75    }
76}
77
78// FIXME(eddyb) keep a `&'static spec::Spec` if that can even speed up anything.
79struct OperandPrinter<IMMS: Iterator<Item = spv::Imm>, ID, IDS: Iterator<Item = ID>> {
80    /// Input immediate operands to print from (may be grouped e.g. into literals).
81    imms: iter::Peekable<IMMS>,
82
83    /// Input ID operands to print from.
84    ids: iter::Peekable<IDS>,
85
86    /// Output for the current operand (drained by the `inst_operands` method).
87    out: TokensForOperand<ID>,
88}
89
90impl<IMMS: Iterator<Item = spv::Imm>, ID, IDS: Iterator<Item = ID>> OperandPrinter<IMMS, ID, IDS> {
91    fn is_exhausted(&mut self) -> bool {
92        self.imms.peek().is_none() && self.ids.peek().is_none()
93    }
94
95    fn enumerant_params(&mut self, enumerant: &spec::Enumerant) {
96        let mut first = true;
97        for (mode, name_and_kind) in enumerant.all_params_with_names() {
98            if mode == spec::OperandMode::Optional && self.is_exhausted() {
99                break;
100            }
101
102            self.out.tokens.push(Token::Punctuation(if first { "(" } else { ", " }));
103            first = false;
104
105            let (name, kind) = name_and_kind.name_and_kind();
106            self.operand(name, kind);
107        }
108        if !first {
109            self.out.tokens.push(Token::Punctuation(")"));
110        }
111    }
112
113    fn literal(&mut self, kind: spec::OperandKind, first_word: u32) {
114        // HACK(eddyb) easier to buffer these than to deal with iterators.
115        let mut words = SmallVec::<[u32; 16]>::new();
116        words.push(first_word);
117        while let Some(&spv::Imm::LongCont(cont_kind, word)) = self.imms.peek() {
118            self.imms.next();
119            assert_eq!(kind, cont_kind);
120            words.push(word);
121        }
122
123        let def = kind.def();
124        assert!(matches!(def, spec::OperandKindDef::Literal { .. }));
125
126        let literal_token = if kind == spec::Spec::get().well_known.LiteralString {
127            // FIXME(eddyb) deduplicate with `spv::extract_literal_string`.
128            let bytes: SmallVec<[u8; 64]> = words
129                .into_iter()
130                .flat_map(u32::to_le_bytes)
131                .take_while(|&byte| byte != 0)
132                .collect();
133            match str::from_utf8(&bytes) {
134                Ok(s) => Token::StringLiteral(format!("{s:?}")),
135                Err(e) => Token::Error(format!("/* {e} in {bytes:?} */")),
136            }
137        } else {
138            let mut words_msb_to_lsb =
139                words.into_iter().rev().skip_while(|&word| word == 0).peekable();
140            let most_significant_word = words_msb_to_lsb.next().unwrap_or(0);
141
142            // FIXME(eddyb) use a more advanced decision procedure for picking
143            // how to print integer(?) literals.
144            let mut s;
145            if words_msb_to_lsb.peek().is_none() && most_significant_word <= 0xffff {
146                s = format!("{most_significant_word}");
147            } else {
148                s = format!("0x{most_significant_word:x}");
149                for word in words_msb_to_lsb {
150                    write!(s, "_{word:08x}").unwrap();
151                }
152            }
153            Token::NumericLiteral(s)
154        };
155
156        self.out.tokens.push(literal_token);
157    }
158
159    fn operand(&mut self, operand_name: &'static str, kind: spec::OperandKind) {
160        if !operand_name.is_empty() {
161            self.out.tokens.push(Token::OperandName(operand_name));
162        }
163
164        let (name, def) = kind.name_and_def();
165
166        // FIXME(eddyb) should this be a hard error?
167        let emit_missing_error = |this: &mut Self| {
168            this.out.tokens.push(Token::Error(format!("/* missing {name} */")));
169        };
170
171        let mut maybe_get_enum_word = || match self.imms.next() {
172            Some(spv::Imm::Short(found_kind, word)) => {
173                assert_eq!(kind, found_kind);
174                Some(word)
175            }
176            Some(spv::Imm::LongStart(..) | spv::Imm::LongCont(..)) => unreachable!(),
177            None => None,
178        };
179
180        match def {
181            spec::OperandKindDef::BitEnum { empty_name, bits } => {
182                let word = match maybe_get_enum_word() {
183                    Some(word) => word,
184                    None => return emit_missing_error(self),
185                };
186
187                self.out.tokens.push(Token::OperandKindNamespacePrefix(name));
188                if word == 0 {
189                    self.out.tokens.push(Token::EnumerandName(empty_name));
190                } else if let Some(bit_idx) = spec::BitIdx::of_single_set_bit(word) {
191                    let (bit_name, bit_def) = bits.get_named(bit_idx).unwrap();
192                    self.out.tokens.push(Token::EnumerandName(bit_name));
193                    self.enumerant_params(bit_def);
194                } else {
195                    self.out.tokens.push(Token::Punctuation("{"));
196                    let mut first = true;
197                    for bit_idx in spec::BitIdx::of_all_set_bits(word) {
198                        if !first {
199                            self.out.tokens.push(Token::Punctuation(", "));
200                        }
201                        first = false;
202
203                        let (bit_name, bit_def) = bits.get_named(bit_idx).unwrap();
204                        self.out.tokens.push(Token::EnumerandName(bit_name));
205                        self.enumerant_params(bit_def);
206                    }
207                    self.out.tokens.push(Token::Punctuation("}"));
208                }
209            }
210            spec::OperandKindDef::ValueEnum { variants } => {
211                let word = match maybe_get_enum_word() {
212                    Some(word) => word,
213                    None => return emit_missing_error(self),
214                };
215
216                let (variant_name, variant_def) =
217                    variants.get_named(word.try_into().unwrap()).unwrap();
218                self.out.tokens.extend([
219                    Token::OperandKindNamespacePrefix(name),
220                    Token::EnumerandName(variant_name),
221                ]);
222                self.enumerant_params(variant_def);
223            }
224            spec::OperandKindDef::Id => match self.ids.next() {
225                Some(id) => {
226                    self.out.tokens.push(Token::Id(id));
227                }
228                None => emit_missing_error(self),
229            },
230            spec::OperandKindDef::Literal { .. } => {
231                // FIXME(eddyb) there's no reason to take the first word now,
232                // `self.literal(kind)` could do it itself.
233                match self.imms.next() {
234                    Some(
235                        spv::Imm::Short(found_kind, word) | spv::Imm::LongStart(found_kind, word),
236                    ) => {
237                        assert_eq!(kind, found_kind);
238                        self.literal(kind, word);
239                    }
240                    Some(spv::Imm::LongCont(..)) => unreachable!(),
241                    None => emit_missing_error(self),
242                }
243            }
244        }
245    }
246
247    fn inst_operands(mut self, opcode: spec::Opcode) -> impl Iterator<Item = TokensForOperand<ID>> {
248        opcode.def().all_operands_with_names().map_while(move |(mode, name_and_kind)| {
249            if mode == spec::OperandMode::Optional && self.is_exhausted() {
250                return None;
251            }
252            let (name, kind) = name_and_kind.name_and_kind();
253            self.operand(name, kind);
254            Some(mem::take(&mut self.out))
255        })
256    }
257}
258
259/// Print a single SPIR-V operand from only immediates, potentially composed of
260/// an enumerand with parameters (which consumes more immediates).
261pub fn operand_from_imms<T>(imms: impl IntoIterator<Item = spv::Imm>) -> TokensForOperand<T> {
262    let mut printer = OperandPrinter {
263        imms: imms.into_iter().peekable(),
264        ids: iter::empty().peekable(),
265        out: TokensForOperand::default(),
266    };
267    let &kind = match printer.imms.peek().unwrap() {
268        spv::Imm::Short(kind, _) | spv::Imm::LongStart(kind, _) => kind,
269        spv::Imm::LongCont(..) => unreachable!(),
270    };
271    printer.operand("", kind);
272    assert!(printer.imms.next().is_none());
273    printer.out
274}
275
276/// Group (ordered according to `opcode`) `imms` and `ids` into logical operands
277/// (i.e. long immediates are unflattened) and produce one [`TokensForOperand`] by
278/// printing each of them.
279pub fn inst_operands<ID>(
280    opcode: spec::Opcode,
281    imms: impl IntoIterator<Item = spv::Imm>,
282    ids: impl IntoIterator<Item = ID>,
283) -> impl Iterator<Item = TokensForOperand<ID>> {
284    OperandPrinter {
285        imms: imms.into_iter().peekable(),
286        ids: ids.into_iter().peekable(),
287        out: TokensForOperand::default(),
288    }
289    .inst_operands(opcode)
290}
spirt/spv/print.rs

spirt/spv/
print.rs
