Optimize embedded dirs

3 files changed, 98 insertions, 56 deletions

diff --git a/rustlings-macros/src/lib.rs b/rustlings-macros/src/lib.rs
index fc2bcf1..4417a4f 100644
--- a/rustlings-macros/src/lib.rs
+++ b/rustlings-macros/src/lib.rs
@@ -25,14 +25,28 @@ pub fn include_files(_: TokenStream) -> TokenStream {
     let solution_files = exercises
         .iter()
         .map(|exercise| format!("../solutions/{}/{}.rs", exercise.dir, exercise.name));
-    let dirs = exercises.iter().map(|exercise| &exercise.dir);
-    let readmes = exercises
+
+    let mut dirs = Vec::with_capacity(32);
+    let mut dir_inds = vec![0; exercises.len()];
+
+    for (exercise, dir_ind) in exercises.iter().zip(&mut dir_inds) {
+        // The directory is often the last one inserted.
+        if let Some(ind) = dirs.iter().rev().position(|dir| *dir == exercise.dir) {
+            *dir_ind = dirs.len() - 1 - ind;
+            continue;
+        }
+
+        dirs.push(exercise.dir.as_str());
+        *dir_ind = dirs.len() - 1;
+    }
+
+    let readmes = dirs
         .iter()
-        .map(|exercise| format!("../exercises/{}/README.md", exercise.dir));
+        .map(|dir| format!("../exercises/{dir}/README.md"));
 
     quote! {
         EmbeddedFiles {
-            exercise_files: &[#(ExerciseFiles { exercise: include_bytes!(#exercise_files), solution: include_bytes!(#solution_files) }),*],
+            exercise_files: &[#(ExerciseFiles { exercise: include_bytes!(#exercise_files), solution: include_bytes!(#solution_files), dir_ind: #dir_inds }),*],
             exercise_dirs: &[#(ExerciseDir { name: #dirs, readme: include_bytes!(#readmes) }),*]
         }
     }
diff --git a/src/app_state.rs b/src/app_state.rs
index 8cb3e46..492be34 100644
--- a/src/app_state.rs
+++ b/src/app_state.rs
@@ -193,12 +193,12 @@ impl AppState {
         Ok(())
     }
 
-    pub fn set_current_exercise_ind(&mut self, ind: usize) -> Result<()> {
-        if ind >= self.exercises.len() {
+    pub fn set_current_exercise_ind(&mut self, exercise_ind: usize) -> Result<()> {
+        if exercise_ind >= self.exercises.len() {
             bail!(BAD_INDEX_ERR);
         }
 
-        self.current_exercise_ind = ind;
+        self.current_exercise_ind = exercise_ind;
 
         self.write()
     }
@@ -215,8 +215,11 @@ impl AppState {
         self.write()
     }
 
-    pub fn set_pending(&mut self, ind: usize) -> Result<()> {
-        let exercise = self.exercises.get_mut(ind).context(BAD_INDEX_ERR)?;
+    pub fn set_pending(&mut self, exercise_ind: usize) -> Result<()> {
+        let exercise = self
+            .exercises
+            .get_mut(exercise_ind)
+            .context(BAD_INDEX_ERR)?;
 
         if exercise.done {
             exercise.done = false;
@@ -229,16 +232,10 @@ impl AppState {
 
     // Official exercises: Dump the original file from the binary.
     // Third-party exercises: Reset the exercise file with `git stash`.
-    fn reset(&self, ind: usize, dir_name: Option<&str>, path: &str) -> Result<()> {
+    fn reset(&self, exercise_ind: usize, path: &str) -> Result<()> {
         if self.official_exercises {
             return EMBEDDED_FILES
-                .write_exercise_to_disk(
-                    ind,
-                    dir_name.context(
-                        "Official exercises must be nested in the `exercises` directory",
-                    )?,
-                    path,
-                )
+                .write_exercise_to_disk(exercise_ind, path)
                 .with_context(|| format!("Failed to reset the exercise {path}"));
         }
 
@@ -265,7 +262,7 @@ impl AppState {
     pub fn reset_current_exercise(&mut self) -> Result<&'static str> {
         self.set_pending(self.current_exercise_ind)?;
         let exercise = self.current_exercise();
-        self.reset(self.current_exercise_ind, exercise.dir, exercise.path)?;
+        self.reset(self.current_exercise_ind, exercise.path)?;
 
         Ok(exercise.path)
     }
@@ -277,7 +274,7 @@ impl AppState {
 
         self.set_pending(exercise_ind)?;
         let exercise = &self.exercises[exercise_ind];
-        self.reset(exercise_ind, exercise.dir, exercise.path)?;
+        self.reset(exercise_ind, exercise.path)?;
 
         Ok(exercise.path)
     }
@@ -315,18 +312,9 @@ impl AppState {
         let current_exercise = self.current_exercise();
 
         if self.official_exercises {
-            let dir_name = current_exercise
-                .dir
-                .context("Official exercises must be nested in the `exercises` directory")?;
-            let solution_path = format!("solutions/{dir_name}/{}.rs", current_exercise.name);
-
-            EMBEDDED_FILES.write_solution_to_disk(
-                self.current_exercise_ind,
-                dir_name,
-                &solution_path,
-            )?;
-
-            Ok(Some(solution_path))
+            EMBEDDED_FILES
+                .write_solution_to_disk(self.current_exercise_ind, current_exercise.name)
+                .map(Some)
         } else {
             let solution_path = if let Some(dir) = current_exercise.dir {
                 format!("solutions/{dir}/{}.rs", current_exercise.name)
diff --git a/src/embedded.rs b/src/embedded.rs
index a84e332..23c8d6e 100644
--- a/src/embedded.rs
+++ b/src/embedded.rs
@@ -1,4 +1,4 @@
-use anyhow::{bail, Context, Error, Result};
+use anyhow::{Context, Error, Result};
 use std::{
     fs::{create_dir, create_dir_all, OpenOptions},
     io::{self, Write},
@@ -34,6 +34,7 @@ impl WriteStrategy {
 struct ExerciseFiles {
     exercise: &'static [u8],
     solution: &'static [u8],
+    dir_ind: usize,
 }
 
 struct ExerciseDir {
@@ -43,11 +44,10 @@ struct ExerciseDir {
 
 impl ExerciseDir {
     fn init_on_disk(&self) -> Result<()> {
-        let path_prefix = "exercises/";
-        let readme_path_postfix = "/README.md";
-        let mut dir_path =
-            String::with_capacity(path_prefix.len() + self.name.len() + readme_path_postfix.len());
-        dir_path.push_str(path_prefix);
+        // 20 = 10 + 10
+        // exercises/ + /README.md
+        let mut dir_path = String::with_capacity(20 + self.name.len());
+        dir_path.push_str("exercises/");
         dir_path.push_str(self.name);
 
         if let Err(e) = create_dir(&dir_path) {
@@ -60,10 +60,9 @@ impl ExerciseDir {
             );
         }
 
-        let readme_path = {
-            dir_path.push_str(readme_path_postfix);
-            dir_path
-        };
+        let mut readme_path = dir_path;
+        readme_path.push_str("/README.md");
+
         WriteStrategy::Overwrite.write(&readme_path, self.readme)?;
 
         Ok(())
@@ -95,30 +94,71 @@ impl EmbeddedFiles {
         Ok(())
     }
 
-    pub fn write_exercise_to_disk(
-        &self,
-        exercise_ind: usize,
-        dir_name: &str,
-        path: &str,
-    ) -> Result<()> {
-        let Some(dir) = self.exercise_dirs.iter().find(|dir| dir.name == dir_name) else {
-            bail!("`{dir_name}` not found in the embedded directories");
-        };
+    pub fn write_exercise_to_disk(&self, exercise_ind: usize, path: &str) -> Result<()> {
+        let exercise_files = &EMBEDDED_FILES.exercise_files[exercise_ind];
+        let dir = &EMBEDDED_FILES.exercise_dirs[exercise_files.dir_ind];
 
         dir.init_on_disk()?;
-        WriteStrategy::Overwrite.write(path, self.exercise_files[exercise_ind].exercise)
+        WriteStrategy::Overwrite.write(path, exercise_files.exercise)
     }
 
+    // Write the solution file to disk and return its path.
     pub fn write_solution_to_disk(
         &self,
         exercise_ind: usize,
-        dir_name: &str,
-        path: &str,
-    ) -> Result<()> {
-        let dir_path = format!("solutions/{dir_name}");
+        exercise_name: &str,
+    ) -> Result<String> {
+        let exercise_files = &EMBEDDED_FILES.exercise_files[exercise_ind];
+        let dir = &EMBEDDED_FILES.exercise_dirs[exercise_files.dir_ind];
+
+        // 14 = 10 + 1 + 3
+        // solutions/ + / + .rs
+        let mut dir_path = String::with_capacity(14 + dir.name.len() + exercise_name.len());
+        dir_path.push_str("solutions/");
+        dir_path.push_str(dir.name);
         create_dir_all(&dir_path)
             .with_context(|| format!("Failed to create the directory {dir_path}"))?;
 
-        WriteStrategy::Overwrite.write(path, self.exercise_files[exercise_ind].solution)
+        let mut solution_path = dir_path;
+        solution_path.push('/');
+        solution_path.push_str(exercise_name);
+        solution_path.push_str(".rs");
+
+        WriteStrategy::Overwrite.write(&solution_path, exercise_files.solution)?;
+
+        Ok(solution_path)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use serde::Deserialize;
+
+    use super::*;
+
+    #[derive(Deserialize)]
+    struct ExerciseInfo {
+        dir: String,
+    }
+
+    #[derive(Deserialize)]
+    struct InfoFile {
+        exercises: Vec<ExerciseInfo>,
+    }
+
+    #[test]
+    fn dirs() {
+        let exercises = toml_edit::de::from_str::<InfoFile>(include_str!("../info.toml"))
+            .expect("Failed to parse `info.toml`")
+            .exercises;
+
+        assert_eq!(exercises.len(), EMBEDDED_FILES.exercise_files.len());
+
+        for (exercise, exercise_files) in exercises.iter().zip(EMBEDDED_FILES.exercise_files) {
+            assert_eq!(
+                exercise.dir,
+                EMBEDDED_FILES.exercise_dirs[exercise_files.dir_ind].name,
+            );
+        }
     }
 }