forge_doc/preprocessor/
contract_inheritance.rs

1use super::{Preprocessor, PreprocessorId};
2use crate::{Document, ParseSource, PreprocessorOutput, document::DocumentContent};
3use alloy_primitives::map::HashMap;
4use std::path::PathBuf;
5
6/// [ContractInheritance] preprocessor id.
7pub const CONTRACT_INHERITANCE_ID: PreprocessorId = PreprocessorId("contract_inheritance");
8
9/// The contract inheritance preprocessor.
10///
11/// It matches the documents with inner [`ParseSource::Contract`](crate::ParseSource) elements,
12/// iterates over their base contracts and attempts
13/// to link them with the paths of the other contract documents.
14///
15/// This preprocessor writes to [Document]'s context.
16#[derive(Debug, Default)]
17pub struct ContractInheritance {
18    /// Whether to capture inherited contracts from libraries.
19    pub include_libraries: bool,
20}
21
22impl Preprocessor for ContractInheritance {
23    fn id(&self) -> PreprocessorId {
24        CONTRACT_INHERITANCE_ID
25    }
26
27    fn preprocess(&self, documents: Vec<Document>) -> Result<Vec<Document>, eyre::Error> {
28        for document in &documents {
29            if let DocumentContent::Single(ref item) = document.content
30                && let ParseSource::Contract(ref contract) = item.source
31            {
32                let mut links = HashMap::default();
33
34                // Attempt to match bases to other contracts
35                for base in &contract.bases {
36                    let base_ident = base.ident.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.add_context(self.id(), PreprocessorOutput::ContractInheritance(links));
45                }
46            }
47        }
48
49        Ok(documents)
50    }
51}
52
53impl ContractInheritance {
54    fn try_link_base(&self, base: &str, documents: &Vec<Document>) -> Option<PathBuf> {
55        for candidate in documents {
56            if candidate.from_library && !self.include_libraries {
57                continue;
58            }
59            if let DocumentContent::Single(ref item) = candidate.content
60                && let ParseSource::Contract(ref contract) = item.source
61                && base == contract.name
62            {
63                return Some(candidate.relative_output_path().to_path_buf());
64            }
65        }
66        None
67    }
68}
forge_doc/preprocessor/contract_inheritance.rs

forge_doc/preprocessor/
contract_inheritance.rs
