1#[derive(Debug, Clone, PartialEq)]
6pub enum DirectiveKind {
7 IncludeSymbol,
9 Witness,
11 LeanCheck,
13}
14
15impl DirectiveKind {
16 fn label(&self) -> &'static str {
17 match self {
18 Self::IncludeSymbol => "include-symbol",
19 Self::Witness => "witness",
20 Self::LeanCheck => "lean-check",
21 }
22 }
23
24 fn from_info(info: &str) -> Option<Self> {
25 if info.starts_with("include-symbol") {
26 Some(Self::IncludeSymbol)
27 } else if info.starts_with("witness") {
28 Some(Self::Witness)
29 } else if info.starts_with("lean-check") {
30 Some(Self::LeanCheck)
31 } else {
32 None
33 }
34 }
35}
36
37#[derive(Debug, Clone, PartialEq)]
39pub struct Directive {
40 pub kind: DirectiveKind,
42 pub attrs: Vec<(String, String)>,
44 pub byte_range: (usize, usize),
47}
48
49impl Directive {
50 pub fn attr(&self, key: &str) -> Option<&str> {
52 self.attrs
53 .iter()
54 .find(|(k, _)| k == key)
55 .map(|(_, v)| v.as_str())
56 }
57
58 pub fn require(&self, key: &str) -> Result<&str, crate::CompilerError> {
60 self.attr(key)
61 .ok_or_else(|| crate::CompilerError::Attribute {
62 directive: self.kind.label(),
63 reason: format!("missing required attribute `{key}`"),
64 })
65 }
66}
67
68#[derive(Debug)]
71pub struct Scanned<'a> {
72 source: &'a str,
73 blocks: Vec<Block>,
74}
75
76#[derive(Debug)]
78pub enum Block {
79 Text(usize, usize),
81 Directive(Directive),
83}
84
85impl<'a> Scanned<'a> {
86 pub fn blocks(&self) -> &[Block] {
88 &self.blocks
89 }
90
91 pub fn source(&self) -> &'a str {
93 self.source
94 }
95
96 pub fn range(block: &Block) -> (usize, usize) {
98 match block {
99 Block::Text(start, end) => (*start, *end),
100 Block::Directive(d) => d.byte_range,
101 }
102 }
103
104 pub fn substitute<F>(&self, mut resolve: F) -> Result<String, crate::CompilerError>
108 where
109 F: FnMut(&Directive) -> Result<String, crate::CompilerError>,
110 {
111 let mut out = String::with_capacity(self.source.len());
112 let mut cursor = 0usize;
113 for block in &self.blocks {
114 let (start, end) = Self::range(block);
115 out.push_str(self.source.get(cursor..start).unwrap_or_default());
116 match block {
117 Block::Text(..) => {
118 out.push_str(self.source.get(start..end).unwrap_or_default());
119 }
120 Block::Directive(d) => {
121 out.push_str(&resolve(d)?);
122 out.push_str("\n\n");
123 }
124 }
125 cursor = end;
126 }
127 out.push_str(self.source.get(cursor..).unwrap_or_default());
128 Ok(out)
129 }
130}
131
132pub fn scan(source: &str) -> Scanned<'_> {
135 let mut blocks: Vec<Block> = Vec::new();
136 let bytes = source.as_bytes();
137 let mut cursor = 0usize;
138 let mut text_start = 0usize;
139
140 while cursor < bytes.len() {
141 let line_end = source[cursor..]
142 .find('\n')
143 .map(|p| cursor + p)
144 .unwrap_or(bytes.len());
145 let line = source.get(cursor..line_end).unwrap_or_default();
146 let trimmed = line.trim_start();
147 if let Some(rest) = trimmed.strip_prefix("```") {
148 let info = rest.trim();
149 if let Some(kind) = DirectiveKind::from_info(info).filter(|_| !info.contains('`')) {
153 if cursor > text_start {
155 blocks.push(Block::Text(text_start, cursor));
156 }
157 let mut close = None;
159 let mut scan = line_end + 1;
160 while scan <= bytes.len() {
161 let e = source[scan..]
162 .find('\n')
163 .map(|p| scan + p)
164 .unwrap_or(bytes.len());
165 if source.get(scan..e).unwrap_or_default().trim() == "```" {
166 close = Some(e);
167 break;
168 }
169 scan = e + 1;
170 }
171 let end = close.unwrap_or(bytes.len());
172 blocks.push(Block::Directive(Directive {
173 kind,
174 attrs: parse_attrs(info),
175 byte_range: (cursor, end),
176 }));
177 cursor = end;
178 text_start = end;
179 continue;
180 }
181 }
182 cursor = line_end + 1;
183 }
184 if text_start < bytes.len() {
185 blocks.push(Block::Text(text_start, bytes.len()));
186 }
187 Scanned { source, blocks }
188}
189
190fn parse_attrs(info: &str) -> Vec<(String, String)> {
194 let mut attrs = Vec::new();
195 let rest = match info.find(char::is_whitespace) {
196 Some(pos) => &info[pos..],
197 None => "",
198 };
199 let mut tokens = rest.split_whitespace();
200 while let Some(token) = tokens.next() {
201 let Some(eq) = token.find('=') else {
202 continue;
203 };
204 let Some(key) = token.get(..eq) else {
205 continue;
206 };
207 let Some(mut value) = token.get(eq + 1..).map(str::to_string) else {
208 continue;
209 };
210 if value.starts_with('"') && !value.ends_with('"') {
211 for cont in tokens.by_ref() {
213 value.push(' ');
214 value.push_str(cont);
215 if cont.ends_with('"') {
216 break;
217 }
218 }
219 }
220 let value = value.trim_matches('"').to_string();
221 if !attrs.iter().any(|(k, _)| k == key) {
222 attrs.push((key.to_string(), value));
223 }
224 }
225 attrs
226}
227
228#[cfg(test)]
229#[allow(clippy::indexing_slicing)]
232mod tests {
233 use super::*;
234
235 #[test]
236 fn finds_directive_fences_and_keeps_other_fences() {
237 let md = "intro\n\n```include-symbol file=\"a.rs\" symbol=\"x\"\ncode here\n```\n\n\
238 ```rust\nfn plain() {}\n```\n\n```witness wasm=\"m.wasm\" fuel=1000\nbody\n```\n";
239 let scanned = scan(md);
240 let directives: Vec<&Directive> = scanned
241 .blocks()
242 .iter()
243 .filter_map(|b| match b {
244 Block::Directive(d) => Some(d),
245 _ => None,
246 })
247 .collect();
248 assert_eq!(directives.len(), 2);
249 assert_eq!(directives[0].kind, DirectiveKind::IncludeSymbol);
250 assert_eq!(
251 directives[0].attr("file"),
252 Some("a.rs"),
253 "quoted value unquoted"
254 );
255 assert_eq!(directives[1].kind, DirectiveKind::Witness);
256 assert_eq!(directives[1].attr("fuel"), Some("1000"));
257 }
258
259 #[test]
260 fn quoted_values_with_spaces_stay_one_value() {
261 let attrs = parse_attrs("witness assert_stdout=\"hello world\" fuel=5");
262 assert_eq!(
263 attrs,
264 vec![
265 ("assert_stdout".to_string(), "hello world".to_string()),
266 ("fuel".to_string(), "5".to_string()),
267 ]
268 );
269 }
270}