forge_doc/preprocessor/
contract_inheritance.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
use super::{Preprocessor, PreprocessorId};
use crate::{document::DocumentContent, Document, ParseSource, PreprocessorOutput};
use alloy_primitives::map::HashMap;
use forge_fmt::solang_ext::SafeUnwrap;
use std::path::PathBuf;

/// [ContractInheritance] preprocessor id.
pub const CONTRACT_INHERITANCE_ID: PreprocessorId = PreprocessorId("contract_inheritance");

/// The contract inheritance preprocessor.
///
/// It matches the documents with inner [`ParseSource::Contract`](crate::ParseSource) elements,
/// iterates over their [Base](solang_parser::pt::Base)s and attempts
/// to link them with the paths of the other contract documents.
///
/// This preprocessor writes to [Document]'s context.
#[derive(Debug, Default)]
pub struct ContractInheritance {
    /// Whether to capture inherited contracts from libraries.
    pub include_libraries: bool,
}

impl Preprocessor for ContractInheritance {
    fn id(&self) -> PreprocessorId {
        CONTRACT_INHERITANCE_ID
    }

    fn preprocess(&self, documents: Vec<Document>) -> Result<Vec<Document>, eyre::Error> {
        for document in documents.iter() {
            if let DocumentContent::Single(ref item) = document.content {
                if let ParseSource::Contract(ref contract) = item.source {
                    let mut links = HashMap::default();

                    // Attempt to match bases to other contracts
                    for base in contract.base.iter() {
                        let base_ident = base.name.identifiers.last().unwrap().name.clone();
                        if let Some(linked) = self.try_link_base(&base_ident, &documents) {
                            links.insert(base_ident, linked);
                        }
                    }

                    if !links.is_empty() {
                        // Write to context
                        document
                            .add_context(self.id(), PreprocessorOutput::ContractInheritance(links));
                    }
                }
            }
        }

        Ok(documents)
    }
}

impl ContractInheritance {
    fn try_link_base(&self, base: &str, documents: &Vec<Document>) -> Option<PathBuf> {
        for candidate in documents {
            if candidate.from_library && !self.include_libraries {
                continue;
            }
            if let DocumentContent::Single(ref item) = candidate.content {
                if let ParseSource::Contract(ref contract) = item.source {
                    if base == contract.name.safe_unwrap().name {
                        return Some(candidate.target_path.clone())
                    }
                }
            }
        }
        None
    }
}