Skip to main content

ptx_parser/parser/
module.rs

1use crate::{
2    alt, err, func,
3    lexer::PtxToken,
4    mapc, ok,
5    parser::{
6        ParseErrorKind, PtxParseError, PtxParser, PtxTokenStream, Span,
7        util::{
8            comma_p, directive_exact_p, identifier_p, optional, parse_u32_literal, sep_by, seq,
9            skip_first, skip_semicolon, string_literal_p, try_map, u32_p, u64_p,
10        },
11    },
12    seq_n,
13    r#type::{
14        AliasFunctionDirective, CodeLinkage, DataLinkage, DwarfDirective, EntryFunctionDirective,
15        FuncFunctionDirective, SectionDirective, module::*, variable::ModuleVariableDirective,
16    },
17};
18
19impl PtxParser for Module {
20    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
21        |stream| {
22            let (directives, span) = stream.try_with_span(|stream| {
23                let mut directives = Vec::new();
24                while !stream.is_at_end() {
25                    // Design note: module parsing must not use `many`, because
26                    // it treats a malformed directive like the end of the
27                    // sequence and hides the actionable parser error.
28                    let (directive, _) = ModuleDirective::parse()(stream)?;
29                    directives.push(directive);
30                }
31                Ok(directives)
32            })?;
33            Ok((Module { directives, span }, span))
34        }
35    }
36}
37
38impl PtxParser for ModuleDirective {
39    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
40        // Design note: `alt` retains the deepest diagnostic, so each module
41        // form has one parser and no parallel keyword-dispatch implementation.
42        alt!(
43            parse_module_variable(),
44            parse_entry_function(),
45            parse_func_function(),
46            parse_alias_function(),
47            parse_module_info(),
48            parse_module_debug()
49        )
50    }
51}
52
53fn parse_module_variable()
54-> impl Fn(&mut PtxTokenStream) -> Result<(ModuleDirective, Span), PtxParseError> {
55    mapc!(
56        seq(
57            optional(DataLinkage::parse()),
58            ModuleVariableDirective::parse(),
59        ),
60        ModuleDirective::ModuleVariable { linkage, directive }
61    )
62}
63
64fn parse_entry_function()
65-> impl Fn(&mut PtxTokenStream) -> Result<(ModuleDirective, Span), PtxParseError> {
66    mapc!(
67        seq(
68            optional(CodeLinkage::parse()),
69            EntryFunctionDirective::parse(),
70        ),
71        ModuleDirective::EntryFunction { linkage, directive }
72    )
73}
74
75fn parse_func_function()
76-> impl Fn(&mut PtxTokenStream) -> Result<(ModuleDirective, Span), PtxParseError> {
77    mapc!(
78        seq(
79            optional(CodeLinkage::parse()),
80            FuncFunctionDirective::parse(),
81        ),
82        ModuleDirective::FuncFunction { linkage, directive }
83    )
84}
85
86fn parse_alias_function()
87-> impl Fn(&mut PtxTokenStream) -> Result<(ModuleDirective, Span), PtxParseError> {
88    mapc!(
89        AliasFunctionDirective::parse(),
90        ModuleDirective::AliasFunction { directive }
91    )
92}
93
94fn parse_module_info()
95-> impl Fn(&mut PtxTokenStream) -> Result<(ModuleDirective, Span), PtxParseError> {
96    mapc!(
97        ModuleInfoDirectiveKind::parse(),
98        ModuleDirective::ModuleInfo { directive }
99    )
100}
101
102fn parse_module_debug()
103-> impl Fn(&mut PtxTokenStream) -> Result<(ModuleDirective, Span), PtxParseError> {
104    mapc!(
105        ModuleDebugDirective::parse(),
106        ModuleDirective::Debug { directive }
107    )
108}
109
110impl PtxParser for ModuleInfoDirectiveKind {
111    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
112        alt!(
113            mapc!(
114                VersionDirective::parse(),
115                ModuleInfoDirectiveKind::Version { directive }
116            ),
117            mapc!(
118                TargetDirective::parse(),
119                ModuleInfoDirectiveKind::Target { directive }
120            ),
121            mapc!(
122                AddressSizeDirective::parse(),
123                ModuleInfoDirectiveKind::AddressSize { directive }
124            )
125        )
126    }
127}
128
129impl PtxParser for VersionDirective {
130    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
131        try_map(
132            skip_first(directive_exact_p("version"), version_number_p()),
133            func!(|(major, minor)| { ok!(VersionDirective { major, minor }) }),
134        )
135    }
136}
137
138impl PtxParser for TargetDirective {
139    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
140        mapc!(
141            skip_first(
142                directive_exact_p("target"),
143                sep_by(TargetString::parse(), comma_p()),
144            ),
145            TargetDirective { entries }
146        )
147    }
148}
149
150impl PtxParser for TargetString {
151    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
152        // Parse target specifiers like "sm_80", "texmode_unified", etc.
153        try_map(
154            identifier_p(),
155            func!(|name| match name.as_str() {
156                "sm_120a" => ok!(TargetString::Sm120a),
157                "sm_120f" => ok!(TargetString::Sm120f),
158                "sm_120" => ok!(TargetString::Sm120),
159                "sm_121a" => ok!(TargetString::Sm121a),
160                "sm_121f" => ok!(TargetString::Sm121f),
161                "sm_121" => ok!(TargetString::Sm121),
162                "sm_110a" => ok!(TargetString::Sm110a),
163                "sm_110f" => ok!(TargetString::Sm110f),
164                "sm_110" => ok!(TargetString::Sm110),
165                "sm_100a" => ok!(TargetString::Sm100a),
166                "sm_100f" => ok!(TargetString::Sm100f),
167                "sm_100" => ok!(TargetString::Sm100),
168                "sm_101a" => ok!(TargetString::Sm101a),
169                "sm_101f" => ok!(TargetString::Sm101f),
170                "sm_101" => ok!(TargetString::Sm101),
171                "sm_103a" => ok!(TargetString::Sm103a),
172                "sm_103f" => ok!(TargetString::Sm103f),
173                "sm_103" => ok!(TargetString::Sm103),
174                "sm_90a" => ok!(TargetString::Sm90a),
175                "sm_90" => ok!(TargetString::Sm90),
176                "sm_80" => ok!(TargetString::Sm80),
177                "sm_86" => ok!(TargetString::Sm86),
178                "sm_87" => ok!(TargetString::Sm87),
179                "sm_88" => ok!(TargetString::Sm88),
180                "sm_89" => ok!(TargetString::Sm89),
181                "sm_70" => ok!(TargetString::Sm70),
182                "sm_72" => ok!(TargetString::Sm72),
183                "sm_75" => ok!(TargetString::Sm75),
184                "sm_60" => ok!(TargetString::Sm60),
185                "sm_61" => ok!(TargetString::Sm61),
186                "sm_62" => ok!(TargetString::Sm62),
187                "sm_50" => ok!(TargetString::Sm50),
188                "sm_52" => ok!(TargetString::Sm52),
189                "sm_53" => ok!(TargetString::Sm53),
190                "sm_30" => ok!(TargetString::Sm30),
191                "sm_32" => ok!(TargetString::Sm32),
192                "sm_35" => ok!(TargetString::Sm35),
193                "sm_37" => ok!(TargetString::Sm37),
194                "sm_20" => ok!(TargetString::Sm20),
195                "sm_10" => ok!(TargetString::Sm10),
196                "sm_11" => ok!(TargetString::Sm11),
197                "sm_12" => ok!(TargetString::Sm12),
198                "sm_13" => ok!(TargetString::Sm13),
199                "texmode_unified" => ok!(TargetString::TexmodeUnified),
200                "texmode_independent" => ok!(TargetString::TexmodeIndependent),
201                "debug" => ok!(TargetString::Debug),
202                "map_f64_to_f32" => ok!(TargetString::MapF64ToF32),
203                _ => err!(ParseErrorKind::InvalidLiteral(format!(
204                    "unknown target specifier: {}",
205                    name
206                ))),
207            }),
208        )
209    }
210}
211
212impl PtxParser for AddressSizeDirective {
213    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
214        mapc!(
215            skip_first(directive_exact_p("address_size"), AddressSize::parse()),
216            AddressSizeDirective { size }
217        )
218    }
219}
220
221impl PtxParser for ModuleDebugDirective {
222    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
223        alt!(
224            mapc!(
225                FileDirective::parse(),
226                ModuleDebugDirective::File { directive }
227            ),
228            mapc!(
229                SectionDirective::parse(),
230                ModuleDebugDirective::Section { directive }
231            ),
232            mapc!(
233                skip_semicolon(DwarfDirective::parse()),
234                ModuleDebugDirective::Dwarf { directive }
235            )
236        )
237    }
238}
239
240impl PtxParser for FileDirective {
241    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
242        try_map(
243            skip_first(
244                directive_exact_p("file"),
245                seq_n!(
246                    u32_p(),
247                    string_literal_p(),
248                    optional(skip_first(
249                        comma_p(),
250                        seq(u64_p(), skip_first(comma_p(), u64_p())),
251                    )),
252                ),
253            ),
254            |(index, path, maybe_timestamps), span| {
255                let (timestamp, file_size) = if let Some((ts, size)) = maybe_timestamps {
256                    (Some(ts), Some(size))
257                } else {
258                    (None, None)
259                };
260                ok!(FileDirective {
261                    index,
262                    path,
263                    timestamp,
264                    file_size,
265                })
266            },
267        )
268    }
269}
270
271impl PtxParser for AddressSize {
272    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
273        try_map(
274            u32_p(),
275            func!(|value| match value {
276                32 => ok!(AddressSize::Size32),
277                64 => ok!(AddressSize::Size64),
278                other => err!(ParseErrorKind::InvalidLiteral(format!(
279                    "invalid address size: {} (expected 32 or 64)",
280                    other
281                ))),
282            }),
283        )
284    }
285}
286
287/// Parser for version numbers - handles both Float("8.5") and separate tokens (8 . 5)
288fn version_number_p() -> impl Fn(&mut PtxTokenStream) -> Result<((u32, u32), Span), PtxParseError> {
289    |stream| {
290        let start_pos = stream.position().0;
291
292        // Try to parse as float first
293        if let Ok((token, span)) = stream.peek() {
294            if let PtxToken::Float(f) = token {
295                let version_str = f.clone();
296                stream.consume()?;
297                let end_pos = stream.position().0;
298                let full_span = Span::new(start_pos, end_pos);
299                let parts: Vec<&str> = version_str.split('.').collect();
300                let span = span.clone();
301                if parts.len() != 2 {
302                    return err!(ParseErrorKind::InvalidLiteral(format!(
303                        "expected version in format X.Y, got {}",
304                        version_str
305                    )));
306                }
307                let major = parse_u32_literal(parts[0], span)?;
308                let minor = parse_u32_literal(parts[1], span)?;
309                return Ok(((major, minor), full_span));
310            }
311        }
312
313        // Otherwise parse as integer.integer
314        let (major, _) = u32_p()(stream)?;
315        stream.expect(&PtxToken::Dot)?;
316        let (minor, _) = u32_p()(stream)?;
317
318        let end_pos = stream.position().0;
319        let span = Span::new(start_pos, end_pos);
320        Ok(((major, minor), span))
321    }
322}