ethers_solc/project_util/
mock.rs

1//! Helpers to generate mock projects
2
3use crate::{
4    error::Result, remappings::Remapping, resolver::GraphEdges, Graph, ProjectPathsConfig,
5    SolcError,
6};
7use rand::{
8    distributions::{Distribution, Uniform},
9    seq::SliceRandom,
10    Rng,
11};
12use serde::{Deserialize, Serialize};
13use std::{
14    collections::{BTreeSet, HashMap, HashSet, VecDeque},
15    path::{Path, PathBuf},
16};
17
18/// Represents the layout of a project
19#[derive(Serialize, Deserialize, Default)]
20pub struct MockProjectSkeleton {
21    /// all files for the project
22    pub files: Vec<MockFile>,
23    /// all libraries
24    pub libraries: Vec<MockLib>,
25}
26
27impl MockProjectSkeleton {
28    /// Returns a list of file ids the given file id imports.
29    pub fn imported_nodes(&self, from: usize) -> impl Iterator<Item = usize> + '_ {
30        self.files[from].imports.iter().map(|i| i.file_id())
31    }
32}
33
34/// Represents a virtual project
35#[derive(Serialize)]
36pub struct MockProjectGenerator {
37    /// how to name things
38    #[serde(skip)]
39    name_strategy: Box<dyn NamingStrategy + 'static>,
40
41    #[serde(flatten)]
42    inner: MockProjectSkeleton,
43}
44
45impl MockProjectGenerator {
46    /// Create a new project and populate it using the given settings
47    pub fn new(settings: &MockProjectSettings) -> Self {
48        let mut mock = Self::default();
49        mock.populate(settings);
50        mock
51    }
52
53    /// Create a skeleton of a real project
54    pub fn create(paths: &ProjectPathsConfig) -> Result<Self> {
55        fn get_libs(edges: &GraphEdges, lib_folder: &Path) -> Option<HashMap<PathBuf, Vec<usize>>> {
56            let mut libs: HashMap<_, Vec<_>> = HashMap::new();
57            for lib_file in edges.library_files() {
58                let component =
59                    edges.node_path(lib_file).strip_prefix(lib_folder).ok()?.components().next()?;
60                libs.entry(lib_folder.join(component)).or_default().push(lib_file);
61            }
62            Some(libs)
63        }
64
65        let graph = Graph::resolve(paths)?;
66        let mut gen = MockProjectGenerator::default();
67        let (_, edges) = graph.into_sources();
68
69        // add all files as source files
70        gen.add_sources(edges.files().count());
71
72        // stores libs and their files
73        let libs = get_libs(
74            &edges,
75            &paths.libraries.first().cloned().unwrap_or_else(|| paths.root.join("lib")),
76        )
77        .ok_or_else(|| SolcError::msg("Failed to detect libs"))?;
78
79        // mark all files as libs
80        for (lib_id, lib_files) in libs.into_values().enumerate() {
81            let lib_name = gen.name_strategy.new_lib_name(lib_id);
82            let offset = gen.inner.files.len();
83            let lib = MockLib { name: lib_name, id: lib_id, num_files: lib_files.len(), offset };
84            for lib_file in lib_files {
85                let file = &mut gen.inner.files[lib_file];
86                file.lib_id = Some(lib_id);
87                file.name = gen.name_strategy.new_lib_name(file.id);
88            }
89            gen.inner.libraries.push(lib);
90        }
91
92        for id in edges.files() {
93            for import in edges.imported_nodes(id).iter().copied() {
94                let import = gen.get_import(import);
95                gen.inner.files[id].imports.insert(import);
96            }
97        }
98
99        Ok(gen)
100    }
101
102    /// Consumes the type and returns the underlying skeleton
103    pub fn into_inner(self) -> MockProjectSkeleton {
104        self.inner
105    }
106
107    /// Generate all solidity files and write under the paths config
108    pub fn write_to(&self, paths: &ProjectPathsConfig, version: impl AsRef<str>) -> Result<()> {
109        let version = version.as_ref();
110        for file in self.inner.files.iter() {
111            let imports = self.get_imports(file.id);
112            let content = file.mock_content(version, imports.join("\n").as_str());
113            super::create_contract_file(file.target_path(self, paths), content)?;
114        }
115
116        Ok(())
117    }
118
119    fn get_imports(&self, file: usize) -> Vec<String> {
120        let file = &self.inner.files[file];
121        let mut imports = Vec::with_capacity(file.imports.len());
122
123        for import in file.imports.iter() {
124            match *import {
125                MockImport::Internal(f) => {
126                    imports.push(format!("import \"./{}.sol\";", self.inner.files[f].name));
127                }
128                MockImport::External(lib, f) => {
129                    imports.push(format!(
130                        "import \"{}/{}.sol\";",
131                        self.inner.libraries[lib].name, self.inner.files[f].name
132                    ));
133                }
134            }
135        }
136        imports
137    }
138
139    /// Returns all the remappings for the project for the given root path
140    pub fn remappings_at(&self, root: &Path) -> Vec<Remapping> {
141        self.inner
142            .libraries
143            .iter()
144            .map(|lib| {
145                let path = root.join("lib").join(&lib.name).join("src");
146                format!("{}/={}/", lib.name, path.display()).parse().unwrap()
147            })
148            .collect()
149    }
150
151    /// Returns all the remappings for the project
152    pub fn remappings(&self) -> Vec<Remapping> {
153        self.inner
154            .libraries
155            .iter()
156            .map(|lib| format!("{0}/=lib/{0}/src/", lib.name).parse().unwrap())
157            .collect()
158    }
159
160    /// Generates a random project with random settings
161    pub fn random() -> Self {
162        let settings = MockProjectSettings::random();
163        let mut mock = Self::default();
164        mock.populate(&settings);
165        mock
166    }
167
168    /// Adds sources and libraries and populates imports based on the settings
169    pub fn populate(&mut self, settings: &MockProjectSettings) -> &mut Self {
170        self.add_sources(settings.num_lib_files);
171        for _ in 0..settings.num_libs {
172            self.add_lib(settings.num_lib_files);
173        }
174        self.populate_imports(settings)
175    }
176
177    fn next_file_id(&self) -> usize {
178        self.inner.files.len()
179    }
180
181    fn next_lib_id(&self) -> usize {
182        self.inner.libraries.len()
183    }
184
185    /// Adds a new source file
186    pub fn add_source(&mut self) -> &mut Self {
187        let id = self.next_file_id();
188        let name = self.name_strategy.new_source_file_name(id);
189        let file =
190            MockFile { id, name, imports: Default::default(), lib_id: None, emit_artifacts: true };
191        self.inner.files.push(file);
192        self
193    }
194
195    /// Adds `num` new source files
196    pub fn add_sources(&mut self, num: usize) -> &mut Self {
197        for _ in 0..num {
198            self.add_source();
199        }
200        self
201    }
202
203    /// Adds a new lib file
204    pub fn add_lib_file(&mut self, lib_id: usize) -> &mut Self {
205        let id = self.next_file_id();
206        let name = self.name_strategy.new_source_file_name(id);
207        let file = MockFile {
208            id,
209            name,
210            imports: Default::default(),
211            lib_id: Some(lib_id),
212            emit_artifacts: true,
213        };
214        self.inner.files.push(file);
215        self
216    }
217
218    /// Adds `num` new source files
219    pub fn add_lib_files(&mut self, num: usize, lib_id: usize) -> &mut Self {
220        for _ in 0..num {
221            self.add_lib_file(lib_id);
222        }
223        self
224    }
225
226    /// Adds a new lib with the number of lib files
227    pub fn add_lib(&mut self, num_files: usize) -> &mut Self {
228        let lib_id = self.next_lib_id();
229        let lib_name = self.name_strategy.new_lib_name(lib_id);
230        let offset = self.inner.files.len();
231        self.add_lib_files(num_files, lib_id);
232        self.inner.libraries.push(MockLib { name: lib_name, id: lib_id, num_files, offset });
233        self
234    }
235
236    /// randomly assign empty file status so that mocked files don't emit artifacts
237    pub fn assign_empty_files(&mut self) -> &mut Self {
238        let mut rng = rand::thread_rng();
239        let die = Uniform::from(0..self.inner.files.len());
240        for file in self.inner.files.iter_mut() {
241            let throw = die.sample(&mut rng);
242            if throw == 0 {
243                // give it a 1 in num(files) chance that the file will be empty
244                file.emit_artifacts = false;
245            }
246        }
247        self
248    }
249
250    /// Populates the imports of the project
251    pub fn populate_imports(&mut self, settings: &MockProjectSettings) -> &mut Self {
252        let mut rng = rand::thread_rng();
253
254        // populate imports
255        for id in 0..self.inner.files.len() {
256            let imports = if let Some(lib) = self.inner.files[id].lib_id {
257                let num_imports = rng
258                    .gen_range(settings.min_imports..=settings.max_imports)
259                    .min(self.inner.libraries[lib].num_files.saturating_sub(1));
260                self.unique_imports_for_lib(&mut rng, lib, id, num_imports)
261            } else {
262                let num_imports = rng
263                    .gen_range(settings.min_imports..=settings.max_imports)
264                    .min(self.inner.files.len().saturating_sub(1));
265                self.unique_imports_for_source(&mut rng, id, num_imports)
266            };
267
268            self.inner.files[id].imports = imports;
269        }
270        self
271    }
272
273    fn get_import(&self, id: usize) -> MockImport {
274        if let Some(lib) = self.inner.files[id].lib_id {
275            MockImport::External(lib, id)
276        } else {
277            MockImport::Internal(id)
278        }
279    }
280
281    /// Returns the file for the given id
282    pub fn get_file(&self, id: usize) -> &MockFile {
283        &self.inner.files[id]
284    }
285
286    /// All file ids
287    pub fn file_ids(&self) -> impl Iterator<Item = usize> + '_ {
288        self.inner.files.iter().map(|f| f.id)
289    }
290
291    /// Returns an iterator over all file ids that are source files or imported by source files
292    ///
293    /// In other words, all files that are relevant in order to compile the project's source files.
294    pub fn used_file_ids(&self) -> impl Iterator<Item = usize> + '_ {
295        let mut file_ids = BTreeSet::new();
296        for file in self.internal_file_ids() {
297            file_ids.extend(NodesIter::new(file, &self.inner))
298        }
299        file_ids.into_iter()
300    }
301
302    /// All ids of internal files
303    pub fn internal_file_ids(&self) -> impl Iterator<Item = usize> + '_ {
304        self.inner.files.iter().filter(|f| !f.is_external()).map(|f| f.id)
305    }
306
307    /// All ids of external files
308    pub fn external_file_ids(&self) -> impl Iterator<Item = usize> + '_ {
309        self.inner.files.iter().filter(|f| f.is_external()).map(|f| f.id)
310    }
311
312    /// generates exactly `num` unique imports in the range of all files
313    ///
314    /// # Panics
315    ///
316    /// if `num` can't be satisfied because the range is too narrow
317    fn unique_imports_for_source<R: Rng + ?Sized>(
318        &self,
319        rng: &mut R,
320        id: usize,
321        num: usize,
322    ) -> BTreeSet<MockImport> {
323        assert!(self.inner.files.len() > num);
324        let mut imports: Vec<_> = (0..self.inner.files.len()).collect();
325        imports.shuffle(rng);
326        imports.into_iter().filter(|i| *i != id).map(|id| self.get_import(id)).take(num).collect()
327    }
328
329    /// Modifies the content of the given file
330    pub fn modify_file(
331        &self,
332        id: usize,
333        paths: &ProjectPathsConfig,
334        version: impl AsRef<str>,
335    ) -> Result<PathBuf> {
336        let file = &self.inner.files[id];
337        let target = file.target_path(self, paths);
338        let content = file.modified_content(version, self.get_imports(id).join("\n").as_str());
339        super::create_contract_file(target.clone(), content)?;
340
341        Ok(target)
342    }
343
344    /// generates exactly `num` unique imports in the range of a lib's files
345    ///
346    /// # Panics
347    ///
348    /// if `num` can't be satisfied because the range is too narrow
349    fn unique_imports_for_lib<R: Rng + ?Sized>(
350        &self,
351        rng: &mut R,
352        lib_id: usize,
353        id: usize,
354        num: usize,
355    ) -> BTreeSet<MockImport> {
356        let lib = &self.inner.libraries[lib_id];
357        assert!(lib.num_files > num);
358        let mut imports: Vec<_> = (lib.offset..(lib.offset + lib.len())).collect();
359        imports.shuffle(rng);
360        imports.into_iter().filter(|i| *i != id).map(|id| self.get_import(id)).take(num).collect()
361    }
362}
363
364#[allow(clippy::derivable_impls)]
365impl Default for MockProjectGenerator {
366    fn default() -> Self {
367        Self { name_strategy: Box::<SimpleNamingStrategy>::default(), inner: Default::default() }
368    }
369}
370
371impl From<MockProjectSkeleton> for MockProjectGenerator {
372    fn from(inner: MockProjectSkeleton) -> Self {
373        Self { inner, ..Default::default() }
374    }
375}
376
377/// Used to determine the names for elements
378trait NamingStrategy {
379    /// Return a new name for the given source file id
380    fn new_source_file_name(&mut self, id: usize) -> String;
381
382    /// Return a new name for the given lib id
383    fn new_lib_name(&mut self, id: usize) -> String;
384}
385
386/// A primitive naming that simply uses ids to create unique names
387#[derive(Debug, Clone, Copy, Default)]
388pub struct SimpleNamingStrategy {
389    _priv: (),
390}
391
392impl NamingStrategy for SimpleNamingStrategy {
393    fn new_source_file_name(&mut self, id: usize) -> String {
394        format!("SourceFile{id}")
395    }
396
397    fn new_lib_name(&mut self, id: usize) -> String {
398        format!("Lib{id}")
399    }
400}
401
402/// Skeleton of a mock source file
403#[derive(Debug, Clone, Serialize, Deserialize)]
404pub struct MockFile {
405    /// internal id of this file
406    pub id: usize,
407    /// The source name of this file
408    pub name: String,
409    /// all the imported files
410    pub imports: BTreeSet<MockImport>,
411    /// lib id if this file is part of a lib
412    pub lib_id: Option<usize>,
413    /// whether this file should emit artifacts
414    pub emit_artifacts: bool,
415}
416
417impl MockFile {
418    /// Returns `true` if this file is part of an external lib
419    pub fn is_external(&self) -> bool {
420        self.lib_id.is_some()
421    }
422
423    pub fn target_path(&self, gen: &MockProjectGenerator, paths: &ProjectPathsConfig) -> PathBuf {
424        let mut target = if let Some(lib) = self.lib_id {
425            paths.root.join("lib").join(&gen.inner.libraries[lib].name).join("src").join(&self.name)
426        } else {
427            paths.sources.join(&self.name)
428        };
429        target.set_extension("sol");
430
431        target
432    }
433
434    /// Returns the content to use for a modified file
435    ///
436    /// The content here is arbitrary, it should only differ from the mocked content
437    pub fn modified_content(&self, version: impl AsRef<str>, imports: &str) -> String {
438        format!(
439            r#"
440// SPDX-License-Identifier: UNLICENSED
441pragma solidity {};
442{}
443contract {} {{
444    function hello() public {{}}
445}}
446            "#,
447            version.as_ref(),
448            imports,
449            self.name
450        )
451    }
452
453    /// Returns a mocked content for the file
454    pub fn mock_content(&self, version: impl AsRef<str>, imports: &str) -> String {
455        let version = version.as_ref();
456        if self.emit_artifacts {
457            format!(
458                r#"
459// SPDX-License-Identifier: UNLICENSED
460pragma solidity {};
461{}
462contract {} {{}}
463            "#,
464                version, imports, self.name
465            )
466        } else {
467            format!(
468                r#"
469// SPDX-License-Identifier: UNLICENSED
470pragma solidity {version};
471{imports}
472            "#,
473            )
474        }
475    }
476}
477
478#[derive(Debug, Clone, Eq, PartialEq, Hash, Ord, PartialOrd, Serialize, Deserialize)]
479pub enum MockImport {
480    /// Import from the same project
481    Internal(usize),
482    /// external library import
483    /// (`lib id`, `file id`)
484    External(usize, usize),
485}
486
487impl MockImport {
488    pub fn file_id(&self) -> usize {
489        *match self {
490            MockImport::Internal(id) => id,
491            MockImport::External(_, id) => id,
492        }
493    }
494}
495
496/// Container of a mock lib
497#[derive(Debug, Clone, Serialize, Deserialize)]
498pub struct MockLib {
499    /// name of the lib, like `ds-test`
500    pub name: String,
501    /// internal id of this lib
502    pub id: usize,
503    /// offset in the total set of files
504    pub offset: usize,
505    /// number of files included in this lib
506    pub num_files: usize,
507}
508
509impl MockLib {
510    pub fn len(&self) -> usize {
511        self.num_files
512    }
513
514    pub fn is_empty(&self) -> bool {
515        self.len() == 0
516    }
517}
518
519/// Settings to use when generate a mock project
520#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
521pub struct MockProjectSettings {
522    /// number of source files to generate
523    pub num_sources: usize,
524    /// number of libraries to use
525    pub num_libs: usize,
526    /// how many lib files to generate per lib
527    pub num_lib_files: usize,
528    /// min amount of import statements a file can use
529    pub min_imports: usize,
530    /// max amount of import statements a file can use
531    pub max_imports: usize,
532    /// whether to also use files that don't emit artifacts
533    pub allow_no_artifacts_files: bool,
534}
535
536impl MockProjectSettings {
537    /// Generates a new instance with random settings within an arbitrary range
538    pub fn random() -> Self {
539        let mut rng = rand::thread_rng();
540        // arbitrary thresholds
541        MockProjectSettings {
542            num_sources: rng.gen_range(2..25),
543            num_libs: rng.gen_range(0..5),
544            num_lib_files: rng.gen_range(1..10),
545            min_imports: rng.gen_range(0..3),
546            max_imports: rng.gen_range(4..10),
547            allow_no_artifacts_files: true,
548        }
549    }
550
551    /// Generates settings for a large project
552    pub fn large() -> Self {
553        // arbitrary thresholds
554        MockProjectSettings {
555            num_sources: 35,
556            num_libs: 4,
557            num_lib_files: 15,
558            min_imports: 3,
559            max_imports: 12,
560            allow_no_artifacts_files: true,
561        }
562    }
563}
564
565impl Default for MockProjectSettings {
566    fn default() -> Self {
567        // these are arbitrary
568        Self {
569            num_sources: 20,
570            num_libs: 2,
571            num_lib_files: 10,
572            min_imports: 0,
573            max_imports: 5,
574            allow_no_artifacts_files: true,
575        }
576    }
577}
578
579/// An iterator over a node and its dependencies
580struct NodesIter<'a> {
581    /// stack of nodes
582    stack: VecDeque<usize>,
583    visited: HashSet<usize>,
584    skeleton: &'a MockProjectSkeleton,
585}
586
587impl<'a> NodesIter<'a> {
588    fn new(start: usize, skeleton: &'a MockProjectSkeleton) -> Self {
589        Self { stack: VecDeque::from([start]), visited: HashSet::new(), skeleton }
590    }
591}
592
593impl<'a> Iterator for NodesIter<'a> {
594    type Item = usize;
595    fn next(&mut self) -> Option<Self::Item> {
596        let file = self.stack.pop_front()?;
597
598        if self.visited.insert(file) {
599            // push the file's direct imports to the stack if we haven't visited it already
600            self.stack.extend(self.skeleton.imported_nodes(file));
601        }
602        Some(file)
603    }
604}
605
606#[cfg(test)]
607mod tests {
608    use super::*;
609
610    #[test]
611    fn can_generate_mock_project() {
612        let _ = MockProjectGenerator::random();
613    }
614}