forge_doc/preprocessor/
contract_inheritance.rs

Help
1use super::{Preprocessor, PreprocessorId};
2use crate::{document::DocumentContent, Document, ParseSource, PreprocessorOutput};
3use alloy_primitives::map::HashMap;
4use forge_fmt::solang_ext::SafeUnwrap;
5use std::path::PathBuf;
6
7/// [ContractInheritance] preprocessor id.
8pub const CONTRACT_INHERITANCE_ID: PreprocessorId = PreprocessorId("contract_inheritance");
9
10/// The contract inheritance preprocessor.
11///
12/// It matches the documents with inner [`ParseSource::Contract`](crate::ParseSource) elements,
13/// iterates over their [Base](solang_parser::pt::Base)s and attempts
14/// to link them with the paths of the other contract documents.
15///
16/// This preprocessor writes to [Document]'s context.
17#[derive(Debug, Default)]
18pub struct ContractInheritance {
19    /// Whether to capture inherited contracts from libraries.
20    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                if let ParseSource::Contract(ref contract) = item.source {
32                    let mut links = HashMap::default();
33
34                    // Attempt to match bases to other contracts
35                    for base in &contract.base {
36                        let base_ident = base.name.identifiers.last().unwrap().name.clone();
37                        if let Some(linked) = self.try_link_base(&base_ident, &documents) {
38                            links.insert(base_ident, linked);
39                        }
40                    }
41
42                    if !links.is_empty() {
43                        // Write to context
44                        document
45                            .add_context(self.id(), PreprocessorOutput::ContractInheritance(links));
46                    }
47                }
48            }
49        }
50
51        Ok(documents)
52    }
53}
54
55impl ContractInheritance {
56    fn try_link_base(&self, base: &str, documents: &Vec<Document>) -> Option<PathBuf> {
57        for candidate in documents {
58            if candidate.from_library && !self.include_libraries {
59                continue;
60            }
61            if let DocumentContent::Single(ref item) = candidate.content {
62                if let ParseSource::Contract(ref contract) = item.source {
63                    if base == contract.name.safe_unwrap().name {
64                        return Some(candidate.target_path.clone())
65                    }
66                }
67            }
68        }
69        None
70    }
71}
forge_doc/preprocessor/contract_inheritance.rs

forge_doc/preprocessor/
contract_inheritance.rs
