From df24bc9945a8a5bf429d6fb2c4c5609975ba5f3c Mon Sep 17 00:00:00 2001
From: Quaternions <krakow20@gmail.com>
Date: Mon, 4 Sep 2023 17:34:31 -0700
Subject: [PATCH] collect referent instead of Source

---
 src/main.rs | 33 +++++++++++++++++++++------------
 1 file changed, 21 insertions(+), 12 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 34c88a5..65d54ba 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,15 +11,14 @@ fn class_is_a(class: &str, superclass: &str) -> bool {
     return false
 }
 
-fn recursive_collect_scripts(scripts: &mut std::vec::Vec<String>,dom: &rbx_dom_weak::WeakDom, instance: &rbx_dom_weak::Instance){
+fn recursive_collect_scripts(scripts: &mut std::vec::Vec<rbx_dom_weak::types::Ref>,dom: &rbx_dom_weak::WeakDom, instance: &rbx_dom_weak::Instance){
     for &referent in instance.children() {
-        let c = dom.get_by_ref(referent).unwrap();
-        if class_is_a(c.class.as_str(), "LuaSourceContainer") {
-            if let Some(rbx_dom_weak::types::Variant::String(s)) = c.properties.get("Source") {
-                scripts.push(s.clone());
+        if let Some(c) = dom.get_by_ref(referent) {
+            if class_is_a(c.class.as_str(), "LuaSourceContainer") {
+                scripts.push(c.referent());//copy ref
             }
+            recursive_collect_scripts(scripts,dom,c);
         }
-        recursive_collect_scripts(scripts,dom,c);
     }
 }
 
@@ -35,18 +34,28 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         allowed.insert(std::fs::read_to_string(entry?.path())?);
     }
 
-    let mut scripts = std::vec::Vec::<String>::new();
+    let mut scripts = std::vec::Vec::<rbx_dom_weak::types::Ref>::new();
     recursive_collect_scripts(&mut scripts, &dom, dom.root());
 
     //check scribb
     let mut any_failed=false;
-    for (i,script) in scripts.iter().enumerate() {
-        if allowed.contains(script) {
-            println!("pass");
+    for (i,&referent) in scripts.iter().enumerate() {
+        if let Some(script) = dom.get_by_ref(referent) {
+            if let Some(rbx_dom_weak::types::Variant::String(s)) = script.properties.get("Source") {
+                if allowed.contains(s) {
+                    println!("pass");
+                }else{
+                    println!("fail");
+                    any_failed=true;
+                    std::fs::write(format!("blocked/{}.lua",i),s)?;
+                }
+            }else{
+                println!("failed to get source");
+                any_failed=true;
+            }
         }else{
-            println!("fail");
+            println!("failed to deref script");
             any_failed=true;
-            std::fs::write(format!("blocked/{}.lua",i),script)?;
         }
     }
     if any_failed {