wgsl_parser/decl/
func.rs

1use std::sync::Arc;
2
3use gramatika::{Parse, ParseStreamer, Span, Spanned};
4
5use crate::{
6	common::{AttributeList, TypeDecl},
7	parser::ErrorRecoveringParseStream,
8	stmt::BlockStmt,
9	token::{brace, keyword, operator, punct, Token, TokenKind},
10	ParseStream,
11};
12
13use super::Decl;
14
15#[derive(Clone, DebugLisp)]
16pub struct FunctionDecl {
17	pub attributes: Option<AttributeList>,
18	pub storage: Token,
19	pub storage_modifier: Option<Token>,
20	pub name: Token,
21	pub params: Arc<[Decl]>,
22	pub return_ty: Option<TypeDecl>,
23	pub body: BlockStmt,
24}
25
26#[derive(Clone, DebugLisp)]
27pub struct ParamDecl {
28	pub attributes: Option<AttributeList>,
29	pub name: Token,
30	pub ty: TypeDecl,
31}
32
33impl Parse for FunctionDecl {
34	type Stream = ParseStream;
35
36	fn parse(input: &mut Self::Stream) -> gramatika::Result<Self> {
37		let storage = input.consume(keyword![fn])?;
38		let storage_modifier = if input.check(operator![<]) {
39			input.consume(operator![<])?;
40			let modifier = input.consume_kind(TokenKind::Keyword)?;
41			input.consume(operator![>])?;
42
43			Some(modifier)
44		} else {
45			None
46		};
47		let name = input.consume_as(TokenKind::Ident, Token::function)?;
48		input.consume(brace!["("])?;
49
50		let params = input.parse_seq_with_finisher(
51			|input| !input.check(brace![")"]),
52			|input, param| {
53				if !input.check(brace![")"]) {
54					input.consume(punct![,])?;
55				}
56				Ok(Decl::Param(param))
57			},
58		)?;
59
60		input.consume(brace![")"])?;
61		let return_ty = if input.check(operator![->]) {
62			Some(input.parse::<TypeDecl>()?)
63		} else {
64			None
65		};
66
67		let body = input.parse::<BlockStmt>()?;
68
69		Ok(Self {
70			attributes: None,
71			storage,
72			storage_modifier,
73			name,
74			params: params.into(),
75			return_ty,
76			body,
77		})
78	}
79}
80
81impl Spanned for FunctionDecl {
82	fn span(&self) -> Span {
83		self.storage.span().through(self.body.brace_close.span())
84	}
85}
86
87impl Parse for ParamDecl {
88	type Stream = ParseStream;
89
90	fn parse(input: &mut Self::Stream) -> gramatika::Result<Self> {
91		let attributes = if input.check(punct!["@"]) {
92			Some(input.parse::<AttributeList>()?)
93		} else {
94			None
95		};
96		let name = input.consume_kind(TokenKind::Ident)?;
97		let ty = input.parse::<TypeDecl>()?;
98
99		Ok(Self {
100			attributes,
101			name,
102			ty,
103		})
104	}
105}
106
107impl Spanned for ParamDecl {
108	fn span(&self) -> Span {
109		let span_start = self
110			.attributes
111			.as_ref()
112			.and_then(|attrs| attrs.attributes.first().map(|attr| &attr.at_sign))
113			.unwrap_or(&self.name)
114			.span();
115
116		span_start.through(self.ty.span())
117	}
118}