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
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                && let ParseSource::Contract(ref contract) = item.source
32            {
33                let mut links = HashMap::default();
34
35                // Attempt to match bases to other contracts
36                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                    // Write to context
45                    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}