feat: implement get_args :))))

packages/plantarium/src/tree.rs

@@ -1,94 +1,95 @@
-pub fn get_arg<'a>(index: i32, args: &'a Vec<i32>) -> Option<&'a [i32]> {
-    let mut arg_index = -1;
-    let mut i: usize = 0;
-    let mut depth = 0;
+pub fn get_args<'a>(args: &'a Vec<i32>) -> Vec<&[i32]> {
+    let mut idx: usize = 0;
+    let mut depth = -1;
 
-    let mut arg_start_index = 0;
+    let mut arg_start_index = 2;
     let mut next_bracket_index = 0;
+    let mut last_bracket_index = 0;
 
-    while i < args.len() {
-        println!("------");
+    let mut out_args: Vec<&[i32]> = Vec::new();
 
-        let is_bracket = i == next_bracket_index;
-        let value = args[i];
+    let length = args.len();
+
+    while idx < length {
+        let is_bracket = idx == next_bracket_index;
+        let value = args[idx];
 
         // if we are at a bracket
         if is_bracket {
-            next_bracket_index = 1 + i + args[i + 1] as usize;
-            // if is opening
+            // if we are at the end of the args
+            if idx >= length - 1 {
+                break;
+            }
+
+            last_bracket_index = next_bracket_index;
+            next_bracket_index = 1 + idx + args[idx + 1] as usize;
+
             if value == 0 {
-                if i != 0 {
                 depth += 1;
+            } else if value == 1 {
+                depth -= 1;
             }
 
             if depth == 0 {
-                    i += 2;
-                    arg_index += 1;
+                // skip over the bracket encoding
+                idx += 2;
             } else {
-                    i = next_bracket_index;
-                }
-            // if closing
-            } else if value == 1 {
-                depth -= 1;
-                if depth > 1 {
-                    i = next_bracket_index;
-                } else {
-                    i += 2;
-                }
+                // skip to the next bracket if we are at depth > 0
+                idx = next_bracket_index - 1;
             }
+
             continue;
         }
 
-        arg_index += 1;
-
-        println!(
-            "i: {}, next: {}, v: {}, ai: {}",
-            i, next_bracket_index, args[i], arg_index
-        );
-
-        println!("is_bracket: {}", is_bracket);
-        println!("depth: {}", depth);
-
-        println!("arg_start_index: {}", arg_start_index);
-
-        if arg_index == index {
-            println!("------");
-            return Some(&args[arg_start_index..=i]);
+        // this is at the end of args where normally multiple ] are encoded
+        if depth < 0 {
+            break;
         }
 
-        i += 1;
+        // remove starting bracket from single numbers
+        if idx - arg_start_index < 3 && last_bracket_index == arg_start_index {
+            arg_start_index += 2;
         }
 
-    None
+        out_args.push(&args[arg_start_index..=idx]);
+
+        idx += 1;
+        arg_start_index = idx;
     }
 
-fn main() {
+    return out_args;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_example_input_a() {
         let input_a = vec![0, 4, 1, 2, 3, 0, 7, 1, 2, 4, 2, 4, 1, 1, 1, 1];
         // -> [1, 2, 3, [1, 2, 4, 2, 4]]
-    // 1: [1]
-    // 2: [2]
-    // 3: [3]
-    // 4: [1,2,4,2,4]
 
-    let input_b = vec![0, 2, 1, 0, 4, 4, 2, 4, 1, 2, 2, 0, 3, 2, 3, 1, 1, 1, 1];
-    // -> [1,[4,2,4], 2, [2,3]];
-    // 1: [1]
-    // 2: [4,2,4]
-    // 3: [2]
-    // 4: [2,3]
+        let args = get_args(&input_a);
 
-    let input = input_b;
-
-    //let first_arg = get_arg(0, &input);
-    //println!("-----> FIRST: {:?}", first_arg);
-
-    //let second_arg = get_arg(1, &input);
-    //println!("-----> SECOND: {:?}", second_arg);
-
-    //let third_arg = get_arg(2, &input);
-    //println!("-----> THIRD: {:?}", third_arg);
-
-    let fourth_arg = get_arg(3, &input);
-    println!("-----> FOURTH: {:?}", fourth_arg);
+        assert_eq!(args.len(), 4);
+        assert_eq!(args[0], [1]);
+        assert_eq!(args[1], [2]);
+        assert_eq!(args[2], [3]);
+        assert_eq!(args[3], [0, 7, 1, 2, 4, 2, 4, 1]);
     }
 
+    #[test]
+    fn test_example_input_b() {
+        let input_b = vec![0, 3, 7, 1, 0, 4, 4, 2, 4, 1, 2, 2, 0, 3, 2, 3, 1, 1, 1, 1];
+        // -> [1,[4,2,4], 2, [2,3]]
+
+        let args = get_args(&input_b);
+
+        assert_eq!(args.len(), 5);
+        assert_eq!(args[0], [7]);
+        assert_eq!(args[1], [1]);
+        assert_eq!(args[2], [0, 4, 4, 2, 4]);
+        assert_eq!(args[3], [2]);
+        assert_eq!(args[4], [0, 3, 2, 3]);
+    }
+}