forge_doc/preprocessor/
contract_inheritance.rs

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