From a9252ac27f67d2a260c2b8df7a2ccf8045ce08d2 Mon Sep 17 00:00:00 2001
From: krakow10 <krakow20@gmail.com>
Date: Sun, 29 Sep 2024 17:26:57 -0700
Subject: [PATCH] rbx_reflection: Superclass Iterator (#448)

---
 rbx_reflection/CHANGELOG.md        |  1 +
 rbx_reflection/src/database.rs     | 27 +++++++++++++++++++++++++++
 rbx_reflection_database/src/lib.rs | 23 +++++++++++++++++++++++
 3 files changed, 51 insertions(+)

diff --git a/rbx_reflection/CHANGELOG.md b/rbx_reflection/CHANGELOG.md
index bf32ac241..b6f9c36fe 100644
--- a/rbx_reflection/CHANGELOG.md
+++ b/rbx_reflection/CHANGELOG.md
@@ -1,6 +1,7 @@
 # rbx_reflection Changelog
 
 ## Unreleased Changes
+* Added `ReflectionDatabase::superclasses_iter`
 
 ## 4.7.0 (2024-08-22)
 * Update to rbx_types 1.10
diff --git a/rbx_reflection/src/database.rs b/rbx_reflection/src/database.rs
index 2db81dbeb..6d006201b 100644
--- a/rbx_reflection/src/database.rs
+++ b/rbx_reflection/src/database.rs
@@ -30,6 +30,24 @@ pub struct ReflectionDatabase<'a> {
     pub enums: HashMap<Cow<'a, str>, EnumDescriptor<'a>>,
 }
 
+pub struct SuperClassIter<'a> {
+    database: &'a ReflectionDatabase<'a>,
+    descriptor: Option<&'a ClassDescriptor<'a>>,
+}
+impl<'a> SuperClassIter<'a> {
+    fn next_descriptor(&self) -> Option<&'a ClassDescriptor<'a>> {
+        let superclass = self.descriptor?.superclass.as_ref()?;
+        self.database.classes.get(superclass)
+    }
+}
+impl<'a> Iterator for SuperClassIter<'a> {
+    type Item = &'a ClassDescriptor<'a>;
+    fn next(&mut self) -> Option<Self::Item> {
+        let next_descriptor = self.next_descriptor();
+        std::mem::replace(&mut self.descriptor, next_descriptor)
+    }
+}
+
 impl<'a> ReflectionDatabase<'a> {
     /// Creates an empty `ReflectionDatabase` with a version number of 0.0.0.0.
     pub fn new() -> Self {
@@ -59,6 +77,15 @@ impl<'a> ReflectionDatabase<'a> {
         Some(list)
     }
 
+    /// Returns an iterator of superclasses for the provided ClassDescriptor. This
+    /// iterator will start with the provided class and end with `Instance`.
+    pub fn superclasses_iter(&'a self, descriptor: &'a ClassDescriptor<'a>) -> SuperClassIter {
+        SuperClassIter {
+            database: self,
+            descriptor: Some(descriptor),
+        }
+    }
+
     /// Finds the default value of a property given its name and a class that
     /// contains or inherits the property. Returns `Some(&Variant)` if a default
     /// value exists, None otherwise.
diff --git a/rbx_reflection_database/src/lib.rs b/rbx_reflection_database/src/lib.rs
index 5d9fa6f6f..2f0fa5c14 100644
--- a/rbx_reflection_database/src/lib.rs
+++ b/rbx_reflection_database/src/lib.rs
@@ -14,10 +14,33 @@ pub fn get() -> &'static ReflectionDatabase<'static> {
 
 #[cfg(test)]
 mod test {
+    use rbx_reflection::ClassDescriptor;
+
     use super::*;
 
     #[test]
     fn smoke_test() {
         let _database = get();
     }
+
+    #[test]
+    fn superclasses_iter_test() {
+        let database = get();
+        let part_class_descriptor = database.classes.get("Part");
+        let mut iter = database.superclasses_iter(part_class_descriptor.unwrap());
+        fn class_descriptor_eq(lhs: Option<&ClassDescriptor>, rhs: Option<&ClassDescriptor>) {
+            let eq = match (lhs, rhs) {
+                (Some(lhs), Some(rhs)) => lhs.name == rhs.name,
+                (None, None) => true,
+                _ => false,
+            };
+            assert!(eq, "{:?} != {:?}", lhs, rhs);
+        }
+        class_descriptor_eq(iter.next(), part_class_descriptor);
+        class_descriptor_eq(iter.next(), database.classes.get("FormFactorPart"));
+        class_descriptor_eq(iter.next(), database.classes.get("BasePart"));
+        class_descriptor_eq(iter.next(), database.classes.get("PVInstance"));
+        class_descriptor_eq(iter.next(), database.classes.get("Instance"));
+        class_descriptor_eq(iter.next(), None);
+    }
 }