1use crate::locals::{skip_string, skip_trivia};
16use crate::{find_matching_delimiter, find_solidity_keyword, is_identifier, parse_identifier};
17
18#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct StructMember {
21 pub name: String,
22 pub ty: String,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq)]
27struct StructDecl {
28 name: String,
29 scope: Option<String>,
31 members: Vec<StructMember>,
32}
33
34#[derive(Debug, Clone, PartialEq, Eq)]
35struct EnumDecl {
36 name: String,
37 scope: Option<String>,
38 variants: Vec<String>,
39}
40
41#[derive(Debug, Clone, PartialEq, Eq)]
42struct ValueTypeDecl {
43 name: String,
44 scope: Option<String>,
45 underlying: String,
46}
47
48#[derive(Debug, Clone, PartialEq, Eq)]
50struct Scope {
51 name: String,
52 source_id: u64,
53 start: usize,
54 end: usize,
55 bases: Vec<String>,
57}
58
59#[derive(Debug, Clone, Default, PartialEq, Eq)]
61pub struct SourceTypes {
62 structs: Vec<StructDecl>,
63 enums: Vec<EnumDecl>,
64 value_types: Vec<ValueTypeDecl>,
65 scopes: Vec<Scope>,
66}
67
68impl SourceTypes {
69 #[must_use]
71 pub fn parse(source: &str) -> Self {
72 let mut types = Self::default();
73 types.add_source(0, source);
74 types
75 }
76
77 pub fn add_source(&mut self, source_id: u64, source: &str) {
79 let code = code_ranges(source);
80 let scopes = scan_scopes(source, source_id, &code);
81 let scope_of = |offset: usize| -> Option<String> {
82 scopes
83 .iter()
84 .filter(|scope| scope.start <= offset && offset < scope.end)
85 .min_by_key(|scope| scope.end - scope.start)
86 .map(|scope| scope.name.clone())
87 };
88 let mut cursor = 0;
89 while let Some(at) = find_keyword(source, &code, "struct", cursor) {
90 cursor = at + "struct".len();
91 let Some((name, open, close)) = braced_declaration(source, cursor) else {
92 continue;
93 };
94 self.structs.push(StructDecl {
95 name: name.to_owned(),
96 scope: scope_of(at),
97 members: parse_members(&source[open + 1..close]),
98 });
99 cursor = close + 1;
100 }
101 let mut cursor = 0;
102 while let Some(at) = find_keyword(source, &code, "enum", cursor) {
103 cursor = at + "enum".len();
104 let Some((name, open, close)) = braced_declaration(source, cursor) else {
105 continue;
106 };
107 let variants = source[open + 1..close]
108 .split(',')
109 .map(str::trim)
110 .filter(|variant| !variant.is_empty())
111 .map(str::to_owned)
112 .collect();
113 self.enums.push(EnumDecl {
114 name: name.to_owned(),
115 scope: scope_of(at),
116 variants,
117 });
118 cursor = close + 1;
119 }
120 let mut cursor = 0;
121 while let Some(at) = find_keyword(source, &code, "type", cursor) {
122 cursor = at + "type".len();
123 let Some((name, name_end)) = parse_identifier(source, skip_trivia(source, cursor))
124 else {
125 continue;
126 };
127 let after = skip_trivia(source, name_end);
128 let Some(("is", is_end)) = parse_identifier(source, after) else {
129 continue;
130 };
131 let underlying_start = skip_trivia(source, is_end);
132 let Some(end) = source[underlying_start..].find(';') else {
133 continue;
134 };
135 let underlying = source[underlying_start..underlying_start + end].trim();
136 if underlying.is_empty() {
137 continue;
138 }
139 self.value_types.push(ValueTypeDecl {
140 name: name.to_owned(),
141 scope: scope_of(at),
142 underlying: underlying.to_owned(),
143 });
144 cursor = underlying_start + end;
145 }
146 self.scopes.extend(scopes);
147 }
148
149 pub(crate) fn extend_from(&mut self, other: &Self) {
151 self.structs.extend_from_slice(&other.structs);
152 self.enums.extend_from_slice(&other.enums);
153 self.value_types.extend_from_slice(&other.value_types);
154 self.scopes.extend_from_slice(&other.scopes);
155 }
156
157 #[must_use]
159 pub fn is_empty(&self) -> bool {
160 self.structs.is_empty() && self.enums.is_empty() && self.value_types.is_empty()
161 }
162
163 #[must_use]
166 pub fn scope_at(&self, source_id: u64, offset: u64) -> Option<&str> {
167 let offset = usize::try_from(offset).ok()?;
168 self.scopes
169 .iter()
170 .filter(|scope| {
171 scope.source_id == source_id && scope.start <= offset && offset < scope.end
172 })
173 .min_by_key(|scope| scope.end - scope.start)
174 .map(|scope| scope.name.as_str())
175 }
176
177 #[must_use]
180 pub fn struct_members(&self, scope: Option<&str>, name: &str) -> Option<&[StructMember]> {
181 self.resolve(
182 self.structs.iter(),
183 |decl| (&decl.name, decl.scope.as_deref()),
184 scope,
185 name,
186 )
187 .map(|decl| decl.members.as_slice())
188 }
189
190 #[must_use]
192 pub fn enum_variants(&self, scope: Option<&str>, name: &str) -> Option<&[String]> {
193 self.resolve(
194 self.enums.iter(),
195 |decl| (&decl.name, decl.scope.as_deref()),
196 scope,
197 name,
198 )
199 .map(|decl| decl.variants.as_slice())
200 }
201
202 #[must_use]
204 pub fn underlying(&self, scope: Option<&str>, name: &str) -> Option<&str> {
205 self.resolve(
206 self.value_types.iter(),
207 |decl| (&decl.name, decl.scope.as_deref()),
208 scope,
209 name,
210 )
211 .map(|decl| decl.underlying.as_str())
212 }
213
214 #[must_use]
216 pub fn enum_literal(&self, scope: Option<&str>, path: &str) -> Option<u64> {
217 let (enum_name, variant) = path.rsplit_once('.')?;
218 let variants = self.enum_variants(scope, enum_name)?;
219 variants
220 .iter()
221 .position(|candidate| candidate == variant)
222 .map(|index| index as u64)
223 }
224
225 #[must_use]
227 pub fn is_struct(&self, scope: Option<&str>, name: &str) -> bool {
228 self.struct_members(scope, name).is_some()
229 }
230
231 #[must_use]
236 pub fn enum_variants_by_name(&self) -> std::collections::BTreeMap<String, Vec<String>> {
237 let mut names = std::collections::BTreeMap::new();
238 for decl in &self.enums {
239 if let Some(scope) = &decl.scope {
240 names
241 .entry(format!("{scope}.{}", decl.name))
242 .or_insert_with(|| decl.variants.clone());
243 }
244 names
245 .entry(decl.name.clone())
246 .or_insert_with(|| decl.variants.clone());
247 }
248 names
249 }
250
251 fn bases(&self, scope: &str) -> Vec<&str> {
253 let mut seen = Vec::<&str>::new();
254 let mut queue = self
255 .scopes
256 .iter()
257 .filter(|candidate| candidate.name == scope)
258 .flat_map(|candidate| candidate.bases.iter().map(String::as_str))
259 .collect::<std::collections::VecDeque<_>>();
260 while let Some(base) = queue.pop_front() {
261 if seen.contains(&base) {
262 continue;
263 }
264 seen.push(base);
265 queue.extend(
266 self.scopes
267 .iter()
268 .filter(|candidate| candidate.name == base)
269 .flat_map(|candidate| candidate.bases.iter().map(String::as_str)),
270 );
271 }
272 seen
273 }
274
275 fn resolve<'a, T>(
279 &'a self,
280 declarations: impl Iterator<Item = &'a T> + Clone,
281 key: impl Fn(&'a T) -> (&'a String, Option<&'a str>),
282 scope: Option<&str>,
283 name: &str,
284 ) -> Option<&'a T> {
285 let (qualifier, bare) = match name.rsplit_once('.') {
286 Some((qualifier, bare)) => (Some(qualifier), bare),
287 None => (None, name),
288 };
289 let candidates = declarations
290 .filter(|declaration| key(declaration).0 == bare)
291 .collect::<Vec<_>>();
292 if candidates.is_empty() {
293 return None;
294 }
295 let in_scope = |wanted: Option<&str>| {
296 candidates
297 .iter()
298 .copied()
299 .find(|declaration| key(declaration).1 == wanted)
300 };
301 if let Some(qualifier) = qualifier {
302 return in_scope(Some(qualifier)).or_else(|| candidates.first().copied());
305 }
306 if let Some(scope) = scope {
307 if let Some(found) = in_scope(Some(scope)) {
308 return Some(found);
309 }
310 for base in self.bases(scope) {
311 if let Some(found) = in_scope(Some(base)) {
312 return Some(found);
313 }
314 }
315 }
316 in_scope(None).or_else(|| candidates.first().copied())
317 }
318}
319
320fn code_ranges(source: &str) -> Vec<(usize, usize)> {
322 let bytes = source.as_bytes();
323 let mut ranges = Vec::new();
324 let mut index = 0;
325 let mut start = 0;
326 while index < bytes.len() {
327 let skipped = skip_trivia(source, index);
328 if skipped != index {
329 if start < index {
330 ranges.push((start, index));
331 }
332 index = skipped;
333 start = index;
334 continue;
335 }
336 match bytes[index] {
337 b'"' | b'\'' => {
338 if start < index {
339 ranges.push((start, index));
340 }
341 index = skip_string(source, index);
342 start = index;
343 }
344 _ => index += 1,
345 }
346 }
347 if start < bytes.len() {
348 ranges.push((start, bytes.len()));
349 }
350 ranges
351}
352
353fn find_keyword(
355 source: &str,
356 code: &[(usize, usize)],
357 keyword: &str,
358 start: usize,
359) -> Option<usize> {
360 let mut cursor = start;
361 while let Some(at) = find_solidity_keyword(source, keyword, cursor) {
362 let before = code.partition_point(|(from, _)| *from <= at);
365 if before
366 .checked_sub(1)
367 .is_some_and(|range| at < code[range].1)
368 {
369 return Some(at);
370 }
371 cursor = at + keyword.len();
372 }
373 None
374}
375
376fn scan_scopes(source: &str, source_id: u64, code: &[(usize, usize)]) -> Vec<Scope> {
378 let mut scopes = Vec::new();
379 for keyword in ["contract", "library", "interface"] {
380 let mut cursor = 0;
381 while let Some(at) = find_keyword(source, code, keyword, cursor) {
382 cursor = at + keyword.len();
383 let Some((name, name_end)) = parse_identifier(source, skip_trivia(source, cursor))
384 else {
385 continue;
386 };
387 let Some(open) = source[name_end..].find('{').map(|found| name_end + found) else {
390 continue;
391 };
392 let header = &source[name_end..open];
393 if header.contains(';') {
394 continue;
395 }
396 let Some(close) = find_matching_delimiter(source, open, b'{', b'}') else {
397 continue;
398 };
399 scopes.push(Scope {
400 name: name.to_owned(),
401 source_id,
402 start: open,
403 end: close,
404 bases: parse_bases(header),
405 });
406 cursor = open + 1;
407 }
408 }
409 scopes
410}
411
412fn parse_bases(header: &str) -> Vec<String> {
414 let trimmed = skip_trivia(header, 0);
415 let Some(rest) = header[trimmed..].strip_prefix("is") else {
416 return Vec::new();
417 };
418 if rest
419 .bytes()
420 .next()
421 .is_some_and(|byte| byte == b'_' || byte.is_ascii_alphanumeric())
422 {
423 return Vec::new();
424 }
425 let mut bases = Vec::new();
426 let mut depth = 0_i32;
427 let mut start = 0;
428 let bytes = rest.as_bytes();
429 for index in 0..=bytes.len() {
430 match bytes.get(index) {
431 Some(b'(') => {
432 if depth == 0 {
433 push_base(&mut bases, &rest[start..index]);
434 start = bytes.len();
435 }
436 depth += 1;
437 }
438 Some(b')') => depth -= 1,
439 Some(b',') | None if depth == 0 => {
440 if start < index {
441 push_base(&mut bases, &rest[start..index]);
442 }
443 start = index + 1;
444 }
445 _ => {}
446 }
447 }
448 bases
449}
450
451fn push_base(bases: &mut Vec<String>, text: &str) {
452 let name = text.trim();
453 if is_identifier(name) {
454 bases.push(name.to_owned());
455 }
456}
457
458fn braced_declaration(source: &str, after: usize) -> Option<(&str, usize, usize)> {
461 let (name, name_end) = parse_identifier(source, skip_trivia(source, after))?;
462 let open = skip_trivia(source, name_end);
463 if source.as_bytes().get(open) != Some(&b'{') {
464 return None;
465 }
466 let close = find_matching_delimiter(source, open, b'{', b'}')?;
467 Some((name, open, close))
468}
469
470fn parse_members(body: &str) -> Vec<StructMember> {
472 let mut members = Vec::new();
473 let mut index = 0;
474 while index < body.len() {
475 index = skip_trivia(body, index);
476 if index >= body.len() {
477 break;
478 }
479 let Some(end) = find_member_end(body, index) else {
480 break;
481 };
482 if let Some((ty, name)) = split_member(&body[index..end]) {
483 members.push(StructMember { name, ty });
484 }
485 index = end + 1;
486 }
487 members
488}
489
490fn split_member(text: &str) -> Option<(String, String)> {
492 let text = text.trim();
493 let name_start = text
494 .rfind(|character: char| !(character.is_ascii_alphanumeric() || character == '_'))
495 .map_or(0, |index| index + 1);
496 let name = &text[name_start..];
497 let ty = text[..name_start]
498 .split_whitespace()
499 .collect::<Vec<_>>()
500 .join(" ");
501 if !is_identifier(name) || ty.is_empty() {
502 return None;
503 }
504 Some((ty, name.to_owned()))
505}
506
507fn find_member_end(body: &str, index: usize) -> Option<usize> {
510 let bytes = body.as_bytes();
511 let mut depth = 0_i32;
512 for (offset, byte) in bytes.iter().enumerate().skip(index) {
513 match byte {
514 b'(' | b'[' => depth += 1,
515 b')' | b']' => depth -= 1,
516 b';' if depth <= 0 => return Some(offset),
517 _ => {}
518 }
519 }
520 None
521}
522
523#[cfg(test)]
524mod tests {
525 use super::{SourceTypes, StructMember};
526
527 const SOURCE: &str = "\
528// SPDX-License-Identifier: MIT
529pragma solidity ^0.8.20;
530
531type Price is uint128;
532
533struct Pair {
534 uint256 a;
535 uint256 b;
536}
537
538enum Level { Low, High }
539
540contract Shop {
541 enum Color { Red, Green, Blue }
542 struct Item {
543 uint256 id;
544 string name;
545 Color color;
546 uint256[] tags;
547 mapping(address => uint256) owners;
548 Pair[2] pairs;
549 }
550 type Money is uint256;
551 function f() public pure returns (uint256) { return 0; } // struct in a comment
552}
553
554library Math {
555 struct Item { uint8 x; }
556}
557
558abstract contract Base is Shop {
559 struct Item { uint16 y; }
560}
561
562contract Outlet is Base, Math(1) {
563 function g() public pure returns (uint256) { return 1; }
564}
565";
566
567 #[test]
568 fn finds_structs_enums_and_value_types_with_their_scopes() {
569 let types = SourceTypes::parse(SOURCE);
570 assert_eq!(
571 types.struct_members(None, "Pair"),
572 Some(
573 &[
574 StructMember {
575 name: "a".to_owned(),
576 ty: "uint256".to_owned()
577 },
578 StructMember {
579 name: "b".to_owned(),
580 ty: "uint256".to_owned()
581 },
582 ][..]
583 )
584 );
585 let item = types.struct_members(None, "Shop.Item").expect("Shop.Item");
586 assert_eq!(
587 item.iter()
588 .map(|member| (member.name.as_str(), member.ty.as_str()))
589 .collect::<Vec<_>>(),
590 [
591 ("id", "uint256"),
592 ("name", "string"),
593 ("color", "Color"),
594 ("tags", "uint256[]"),
595 ("owners", "mapping(address => uint256)"),
596 ("pairs", "Pair[2]"),
597 ]
598 );
599 assert_eq!(
600 types.struct_members(None, "Math.Item").map(<[_]>::len),
601 Some(1)
602 );
603 assert_eq!(
604 types
605 .enum_variants(None, "Color")
606 .map(|variants| variants.join(",")),
607 Some("Red,Green,Blue".to_owned())
608 );
609 assert_eq!(
610 types
611 .enum_variants(None, "Level")
612 .map(|variants| variants.join(",")),
613 Some("Low,High".to_owned())
614 );
615 assert_eq!(types.underlying(None, "Price"), Some("uint128"));
616 assert_eq!(types.underlying(None, "Shop.Money"), Some("uint256"));
617 assert_eq!(types.enum_literal(None, "Color.Green"), Some(1));
618 assert_eq!(types.enum_literal(None, "Shop.Color.Blue"), Some(2));
619 assert_eq!(types.enum_literal(None, "Color.Purple"), None);
620 assert_eq!(types.enum_literal(None, "Missing.Red"), None);
621 assert!(types.struct_members(None, "Nothing").is_none());
622 let names = types.enum_variants_by_name();
623 assert_eq!(names["Shop.Color"], ["Red", "Green", "Blue"]);
624 assert_eq!(names["Color"], ["Red", "Green", "Blue"]);
625 assert_eq!(names["Level"], ["Low", "High"]);
626 }
627
628 #[test]
629 fn a_bare_name_resolves_in_the_executing_contract_then_its_bases() {
630 let types = SourceTypes::parse(SOURCE);
631 let width = |scope: Option<&str>| {
633 types
634 .struct_members(scope, "Item")
635 .map(|members| members[0].ty.clone())
636 };
637 assert_eq!(width(Some("Shop")).as_deref(), Some("uint256"));
638 assert_eq!(width(Some("Math")).as_deref(), Some("uint8"));
639 assert_eq!(width(Some("Base")).as_deref(), Some("uint16"));
640 assert_eq!(width(Some("Outlet")).as_deref(), Some("uint16"));
642 assert_eq!(width(Some("Other")).as_deref(), Some("uint256"));
644 assert_eq!(width(None).as_deref(), Some("uint256"));
645 assert_eq!(
647 types.struct_members(Some("Math"), "Pair").map(<[_]>::len),
648 Some(2)
649 );
650 assert_eq!(types.enum_literal(Some("Outlet"), "Color.Blue"), Some(2));
652 let offset = SOURCE.find("function g()").expect("g") as u64;
654 assert_eq!(types.scope_at(0, offset), Some("Outlet"));
655 assert_eq!(types.scope_at(0, 0), None);
656 assert_eq!(types.scope_at(1, offset), None);
657 }
658
659 #[test]
660 fn ignores_keywords_inside_identifiers_and_comments() {
661 let types = SourceTypes::parse(
662 "contract C { uint256 mystruct; /* struct Hidden { uint256 x; } */ enum E { A } }",
663 );
664 assert!(types.struct_members(None, "Hidden").is_none());
665 assert_eq!(types.enum_variants(None, "E").map(<[_]>::len), Some(1));
666 }
667}