Add method to return units matching the prefix (#252)

src/course_library.rs

@@ -21,7 +21,7 @@ use tantivy::{
     schema::{Field, Schema, STORED, TEXT},
     Index, IndexReader, IndexWriter, ReloadPolicy,
 };
-use ustr::{Ustr, UstrMap};
+use ustr::{Ustr, UstrMap, UstrSet};
 use walkdir::WalkDir;
 
 use crate::{
@@ -79,6 +79,10 @@ pub trait CourseLibrary {
     /// Returns the IDs of all exercises in the given course sorted alphabetically.
     fn get_all_exercise_ids(&self) -> Vec<Ustr>;
 
+    /// Returns the set of units whose ID starts with the given prefix and are of the given type.
+    /// If `unit_type` is `None`, then all unit types are considered.
+    fn get_matching_prefix(&self, prefix: &str, unit_type: Option<UnitType>) -> UstrSet;
+
     /// Returns the IDs of all the units which match the given query.
     fn search(&self, query: &str) -> Result<Vec<Ustr>, CourseLibraryError>;
 
@@ -693,6 +697,52 @@ impl CourseLibrary for LocalCourseLibrary {
         exercises
     }
 
+    fn get_matching_prefix(&self, prefix: &str, unit_type: Option<UnitType>) -> UstrSet {
+        match unit_type {
+            Some(UnitType::Course) => self
+                .course_map
+                .iter()
+                .filter_map(|(id, _)| {
+                    if id.starts_with(prefix) {
+                        Some(*id)
+                    } else {
+                        None
+                    }
+                })
+                .collect(),
+            Some(UnitType::Lesson) => self
+                .lesson_map
+                .iter()
+                .filter_map(|(id, _)| {
+                    if id.starts_with(prefix) {
+                        Some(*id)
+                    } else {
+                        None
+                    }
+                })
+                .collect(),
+            Some(UnitType::Exercise) => self
+                .exercise_map
+                .iter()
+                .filter_map(|(id, _)| {
+                    if id.starts_with(prefix) {
+                        Some(*id)
+                    } else {
+                        None
+                    }
+                })
+                .collect(),
+            None => self
+                .course_map
+                .keys()
+                .chain(self.lesson_map.keys())
+                .chain(self.exercise_map.keys())
+                .filter(|id| id.starts_with(prefix))
+                .cloned()
+                .collect(),
+        }
+    }
+
     fn search(&self, query: &str) -> Result<Vec<Ustr>, CourseLibraryError> {
         self.search_helper(query)
             .map_err(|e| CourseLibraryError::Search(query.into(), e))

src/lib.rs

@@ -303,6 +303,12 @@ impl CourseLibrary for Trane {
         self.course_library.read().get_all_exercise_ids()
     }
 
+    fn get_matching_prefix(&self, prefix: &str, unit_type: Option<UnitType>) -> UstrSet {
+        self.course_library
+            .read()
+            .get_matching_prefix(prefix, unit_type)
+    }
+
     fn search(&self, query: &str) -> Result<Vec<Ustr>, CourseLibraryError> {
         self.course_library.read().search(query)
     }

tests/basic_tests.rs

@@ -41,7 +41,7 @@ use trane::{
             ExerciseFilter, FilterOp, FilterType, KeyValueFilter, SessionPart, StudySession,
             StudySessionData, UnitFilter,
         },
-        MasteryScore, SchedulerOptions, UserPreferences,
+        MasteryScore, SchedulerOptions, UnitType, UserPreferences,
     },
     review_list::ReviewList,
     scheduler::ExerciseScheduler,
@@ -1619,6 +1619,77 @@ fn schedule_study_session() -> Result<()> {
     Ok(())
 }
 
+/// Verifies matching the courses with the given prefix.
+#[test]
+fn get_matching_courses() -> Result<()> {
+    // Initialize test course library.
+    let temp_dir = TempDir::new()?;
+    let trane = init_test_simulation(&temp_dir.path(), &BASIC_LIBRARY)?;
+
+    // The test will use the ID of course 0 as the prefix.
+    let prefix = TestId(0, None, None).to_ustr();
+
+    // Get all the courses that match the prefix.
+    let matching_courses = trane.get_matching_prefix(&prefix, Some(UnitType::Course));
+    assert_eq!(matching_courses.len(), 1);
+    assert!(matching_courses.contains(&prefix));
+    Ok(())
+}
+
+/// Verifies matching the lessons with the given prefix.
+#[test]
+fn get_matching_lessons() -> Result<()> {
+    // Initialize test course library.
+    let temp_dir = TempDir::new()?;
+    let trane = init_test_simulation(&temp_dir.path(), &BASIC_LIBRARY)?;
+
+    // The test will use the ID of lesson 0::0 as the prefix.
+    let prefix = TestId(0, Some(0), None).to_ustr();
+
+    // Get all the lessons that match the prefix.
+    let matching_lessons = trane.get_matching_prefix(&prefix, Some(UnitType::Lesson));
+    assert_eq!(matching_lessons.len(), 1);
+    assert!(matching_lessons.contains(&prefix));
+    Ok(())
+}
+
+/// Verifies matching the exercises with the given prefix.
+#[test]
+fn get_matching_exercises() -> Result<()> {
+    // Initialize test course library.
+    let temp_dir = TempDir::new()?;
+    let trane = init_test_simulation(&temp_dir.path(), &BASIC_LIBRARY)?;
+
+    // The test will use the ID of exercise 0::0::0 as the prefix.
+    let prefix = TestId(0, Some(0), Some(0)).to_ustr();
+
+    // Get all the exercises that match the prefix.
+    let matching_exercises = trane.get_matching_prefix(&prefix, Some(UnitType::Exercise));
+    assert_eq!(matching_exercises.len(), 1);
+    assert!(matching_exercises.contains(&prefix));
+    Ok(())
+}
+
+/// Verifies matching all units with the given prefix.
+#[test]
+fn get_matching_units() -> Result<()> {
+    // Initialize test course library.
+    let temp_dir = TempDir::new()?;
+    let trane = init_test_simulation(&temp_dir.path(), &BASIC_LIBRARY)?;
+
+    // The test will use the ID of course 0 as the prefix.
+    let prefix = TestId(0, None, None).to_ustr();
+
+    // Get all the units that match the prefix. The 23 units are the course, the two lessons, and
+    // the ten exercises in each lesson.
+    let matching_units = trane.get_matching_prefix(&prefix, None);
+    assert_eq!(matching_units.len(), 23);
+    assert!(matching_units.contains(&prefix));
+    assert!(matching_units.contains(&TestId(0, Some(0), None).to_ustr()));
+    assert!(matching_units.contains(&TestId(0, Some(0), Some(0)).to_ustr()));
+    Ok(())
+}
+
 /// Verifies searching for courses in the course library works.
 #[test]
 fn course_library_search_courses() -> Result<()> {