forge_doc/preprocessor/
contract_inheritance.rs
1use super::{Preprocessor, PreprocessorId};
2use crate::{Document, ParseSource, PreprocessorOutput, document::DocumentContent};
3use alloy_primitives::map::HashMap;
4use forge_fmt::solang_ext::SafeUnwrap;
5use std::path::PathBuf;
6
7pub const CONTRACT_INHERITANCE_ID: PreprocessorId = PreprocessorId("contract_inheritance");
9
10#[derive(Debug, Default)]
18pub struct ContractInheritance {
19 pub include_libraries: bool,
21}
22
23impl Preprocessor for ContractInheritance {
24 fn id(&self) -> PreprocessorId {
25 CONTRACT_INHERITANCE_ID
26 }
27
28 fn preprocess(&self, documents: Vec<Document>) -> Result<Vec<Document>, eyre::Error> {
29 for document in &documents {
30 if let DocumentContent::Single(ref item) = document.content
31 && let ParseSource::Contract(ref contract) = item.source
32 {
33 let mut links = HashMap::default();
34
35 for base in &contract.base {
37 let base_ident = base.name.identifiers.last().unwrap().name.clone();
38 if let Some(linked) = self.try_link_base(&base_ident, &documents) {
39 links.insert(base_ident, linked);
40 }
41 }
42
43 if !links.is_empty() {
44 document.add_context(self.id(), PreprocessorOutput::ContractInheritance(links));
46 }
47 }
48 }
49
50 Ok(documents)
51 }
52}
53
54impl ContractInheritance {
55 fn try_link_base(&self, base: &str, documents: &Vec<Document>) -> Option<PathBuf> {
56 for candidate in documents {
57 if candidate.from_library && !self.include_libraries {
58 continue;
59 }
60 if let DocumentContent::Single(ref item) = candidate.content
61 && let ParseSource::Contract(ref contract) = item.source
62 && base == contract.name.safe_unwrap().name
63 {
64 return Some(candidate.target_path.clone());
65 }
66 }
67 None
68 }
69}