1use 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#[derive(Serialize, Deserialize, Default)]
20pub struct MockProjectSkeleton {
21 pub files: Vec<MockFile>,
23 pub libraries: Vec<MockLib>,
25}
26
27impl MockProjectSkeleton {
28 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#[derive(Serialize)]
36pub struct MockProjectGenerator {
37 #[serde(skip)]
39 name_strategy: Box<dyn NamingStrategy + 'static>,
40
41 #[serde(flatten)]
42 inner: MockProjectSkeleton,
43}
44
45impl MockProjectGenerator {
46 pub fn new(settings: &MockProjectSettings) -> Self {
48 let mut mock = Self::default();
49 mock.populate(settings);
50 mock
51 }
52
53 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 gen.add_sources(edges.files().count());
71
72 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 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 pub fn into_inner(self) -> MockProjectSkeleton {
104 self.inner
105 }
106
107 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 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 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 pub fn random() -> Self {
162 let settings = MockProjectSettings::random();
163 let mut mock = Self::default();
164 mock.populate(&settings);
165 mock
166 }
167
168 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 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 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 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 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 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 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 file.emit_artifacts = false;
245 }
246 }
247 self
248 }
249
250 pub fn populate_imports(&mut self, settings: &MockProjectSettings) -> &mut Self {
252 let mut rng = rand::thread_rng();
253
254 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 pub fn get_file(&self, id: usize) -> &MockFile {
283 &self.inner.files[id]
284 }
285
286 pub fn file_ids(&self) -> impl Iterator<Item = usize> + '_ {
288 self.inner.files.iter().map(|f| f.id)
289 }
290
291 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 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 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 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 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 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
377trait NamingStrategy {
379 fn new_source_file_name(&mut self, id: usize) -> String;
381
382 fn new_lib_name(&mut self, id: usize) -> String;
384}
385
386#[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#[derive(Debug, Clone, Serialize, Deserialize)]
404pub struct MockFile {
405 pub id: usize,
407 pub name: String,
409 pub imports: BTreeSet<MockImport>,
411 pub lib_id: Option<usize>,
413 pub emit_artifacts: bool,
415}
416
417impl MockFile {
418 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 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 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 Internal(usize),
482 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#[derive(Debug, Clone, Serialize, Deserialize)]
498pub struct MockLib {
499 pub name: String,
501 pub id: usize,
503 pub offset: usize,
505 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#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
521pub struct MockProjectSettings {
522 pub num_sources: usize,
524 pub num_libs: usize,
526 pub num_lib_files: usize,
528 pub min_imports: usize,
530 pub max_imports: usize,
532 pub allow_no_artifacts_files: bool,
534}
535
536impl MockProjectSettings {
537 pub fn random() -> Self {
539 let mut rng = rand::thread_rng();
540 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 pub fn large() -> Self {
553 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 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
579struct NodesIter<'a> {
581 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 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}