@ape-egg/vibe 1.3.2 → 1.6.0

package/compiler/src/compiler/manifest_builder.rs

@@ -4,20 +4,23 @@ use markup5ever_rcdom::{RcDom, NodeData, Handle};
 use regex::Regex;
 use serde::{Serialize, Serializer};
 use serde_json::Value;
-use std::collections::HashMap;
+use std::collections::{HashMap, BTreeMap};
+use crate::compiler::iteration_optimizer::build_iteration_optimizations;
 
 pub struct ManifestBuilder {
     binding_regex: Regex,
+    name_binding_fix_regex: Regex,
 }
 
 impl ManifestBuilder {
     pub fn new() -> Self {
         Self {
             binding_regex: Regex::new(r"@\[((?:[^\[\]]|\[[^\]]*\])+)\]").unwrap(),
+            name_binding_fix_regex: Regex::new(r#"(@\[[^\]]+\])="+"#).unwrap(),
         }
     }
 
-    pub fn build_from_html(&self, html: &str, _state: &Value) -> Result<ManifestNode, String> {
+    pub fn build_from_html(&self, html: &str, _state: &Value, iterations_as_is: bool) -> Result<ManifestNode, String> {
         // Parse HTML
         let dom = parse_document(RcDom::default(), Default::default())
             .from_utf8()
@@ -25,11 +28,59 @@ impl ManifestBuilder {
             .map_err(|e| format!("Failed to parse HTML: {:?}", e))?;
 
         // Walk tree starting from document root
-        let root_node = self.walk_node(&dom.document, &mut 0);
+        let mut root_node = self.walk_node(&dom.document, &mut 0);
+
+        // Inject compiled batch functions directly into iteration nodes (unless disabled)
+        if !iterations_as_is {
+            let compiled_functions = build_iteration_optimizations(html);
+            if let Some(optimizations) = compiled_functions {
+                self.inject_compiled_functions(&mut root_node, &optimizations.iterations);
+            }
+        }
 
         Ok(root_node)
     }
 
+    /// Recursively inject compiled batch functions into iteration nodes
+    fn inject_compiled_functions(&self, node: &mut ManifestNode, compiled_map: &HashMap<String, crate::compiler::iteration_optimizer::CompiledIteration>) {
+        // Check if this is an iteration node with a template
+        if node.node_type.as_deref() == Some("iteration") {
+            if let Some(ref mut compiled_data) = node.compiled {
+                if let Some(ref restoration) = compiled_data.restoration {
+                    if let Some(ref template) = restoration.template {
+                        // Generate hash for this template (same logic as optimizer)
+                        let template_hash = self.generate_template_hash(template);
+
+                        // Look up compiled function
+                        if let Some(compiled) = compiled_map.get(&template_hash) {
+                            compiled_data.iterations = Some(IterationData {
+                                batch_fn: Some(compiled.batch_fn.clone()),
+                                item_alias: Some(compiled.item_alias.clone()),
+                                index_alias: Some(compiled.index_alias.clone()),
+                            });
+                        }
+                    }
+                }
+            }
+        }
+
+        // Recursively process children
+        for child in node.children.values_mut() {
+            self.inject_compiled_functions(child, compiled_map);
+        }
+    }
+
+    /// Generate stable hash for template (matches iteration_optimizer logic)
+    fn generate_template_hash(&self, template: &str) -> String {
+        use std::collections::hash_map::DefaultHasher;
+        use std::hash::{Hash, Hasher};
+
+        let mut hasher = DefaultHasher::new();
+        template.trim().hash(&mut hasher);
+        let hash = hasher.finish();
+        format!("iter_{:x}", hash)
+    }
+
     fn walk_node(&self, node: &Handle, _counter: &mut usize) -> ManifestNode {
         self.walk_node_impl(node, &node.children.borrow())
     }
@@ -38,9 +89,13 @@ impl ManifestBuilder {
         let mut manifest_node = ManifestNode {
             element: None,
             parsed: vec![],
-            children: HashMap::new(),
-            hyperspeed_restoration: None,
+            children: BTreeMap::new(),
+            compiled: None,
             node_type: None,
+            attributes: None,
+            name_bindings: None,
+            meta: None,
+            runtime: None,
         };
 
         match &node.data {
@@ -50,8 +105,11 @@ impl ManifestBuilder {
                 // Check for @[...] bindings
                 if self.binding_regex.is_match(&text) {
                     let parsed = self.split_by_bindings(&text);
-                    manifest_node.hyperspeed_restoration = Some(RestorationData {
-                        parsed: Some(parsed),
+                    manifest_node.compiled = Some(CompiledData {
+                        restoration: Some(RestorationData {
+                            parsed: Some(parsed),
+                            ..Default::default()
+                        }),
                         ..Default::default()
                     });
                 }
@@ -60,20 +118,48 @@ impl ManifestBuilder {
             NodeData::Element { name, attrs, .. } => {
                 let _tag_name = name.local.to_string();
 
-                // Check attribute bindings
-                let mut attr_bindings = HashMap::new();
-                for attr in attrs.borrow().iter() {
+                // Check attribute bindings (both value and name bindings)
+                let mut attr_bindings = BTreeMap::new();
+                let mut name_bindings: Vec<Option<String>> = Vec::new();
+                let borrowed_attrs = attrs.borrow();
+
+                for (idx, attr) in borrowed_attrs.iter().enumerate() {
                     let attr_name = attr.name.local.to_string();
                     let attr_value = attr.value.to_string();
 
-                    if self.binding_regex.is_match(&attr_value) {
+                    // Check for name bindings (binding in attribute NAME)
+                    if self.binding_regex.is_match(&attr_name) {
+                        // Extend vec to include this index
+                        while name_bindings.len() <= idx {
+                            name_bindings.push(None);
+                        }
+                        name_bindings[idx] = Some(attr_name.clone());
+                    }
+                    // Check for attribute value bindings
+                    else if self.binding_regex.is_match(&attr_value) {
                         attr_bindings.insert(attr_name, attr_value);
                     }
                 }
 
+                // Store if we have any bindings
+                let has_name_bindings = name_bindings.iter().any(|b| b.is_some());
+
+                // Set runtime-compatible fields at node level
                 if !attr_bindings.is_empty() {
-                    manifest_node.hyperspeed_restoration = Some(RestorationData {
-                        attributes: Some(attr_bindings),
+                    manifest_node.attributes = Some(attr_bindings.clone());
+                }
+                if has_name_bindings {
+                    manifest_node.name_bindings = Some(name_bindings.clone());
+                }
+
+                // Also store in compiled.restoration for restoration purposes
+                if !attr_bindings.is_empty() || has_name_bindings {
+                    manifest_node.compiled = Some(CompiledData {
+                        restoration: Some(RestorationData {
+                            attributes: if !attr_bindings.is_empty() { Some(attr_bindings) } else { None },
+                            name_bindings: if has_name_bindings { Some(name_bindings) } else { None },
+                            ..Default::default()
+                        }),
                         ..Default::default()
                     });
                 }
@@ -87,6 +173,21 @@ impl ManifestBuilder {
                 if trimmed.starts_with("each ") {
                     manifest_node.node_type = Some("iteration".to_string());
 
+                    // Extract expression: "each items as item" -> "items as item"
+                    let expression = trimmed.strip_prefix("each ").unwrap_or("").to_string();
+
+                    // Parse expression to extract parts: "items as item, index" or "items as item"
+                    let parts: Vec<&str> = expression.split(" as ").collect();
+                    let array_path = parts.get(0).unwrap_or(&"").trim().to_string();
+                    let alias_part = parts.get(1).unwrap_or(&"").trim();
+                    let (item_alias, index_alias) = if let Some(comma_pos) = alias_part.find(',') {
+                        let item = alias_part[..comma_pos].trim().to_string();
+                        let index = alias_part[comma_pos + 1..].trim().to_string();
+                        (item, index)
+                    } else {
+                        (alias_part.to_string(), "index".to_string())
+                    };
+
                     // Find node index in siblings
                     let node_idx = siblings.iter().position(|n| std::ptr::eq(n as *const _, node as *const _));
 
@@ -97,8 +198,26 @@ impl ManifestBuilder {
                         String::new()
                     };
 
-                    manifest_node.hyperspeed_restoration = Some(RestorationData {
-                        template: Some(template),
+                    // Create minimal meta structure - runtime will populate template after parsing restored DOM
+                    manifest_node.meta = Some(serde_json::json!({
+                        "arrayPath": array_path,
+                        "itemAlias": item_alias,
+                        "indexAlias": index_alias
+                    }));
+
+                    // Runtime state
+                    manifest_node.runtime = Some(serde_json::json!({
+                        "instances": [],
+                        "templateRemoved": false
+                    }));
+
+                    // Compiled restoration data
+                    manifest_node.compiled = Some(CompiledData {
+                        restoration: Some(RestorationData {
+                            template: Some(template),
+                            expression: Some(expression),
+                            ..Default::default()
+                        }),
                         ..Default::default()
                     });
                 }
@@ -107,6 +226,7 @@ impl ManifestBuilder {
                 if trimmed.starts_with("if ") {
                     manifest_node.node_type = Some("conditional".to_string());
 
+                    let expression = trimmed.strip_prefix("if ").unwrap_or("").to_string();
                     let node_idx = siblings.iter().position(|n| std::ptr::eq(n as *const _, node as *const _));
 
                     let template = if let Some(idx) = node_idx {
@@ -115,8 +235,24 @@ impl ManifestBuilder {
                         String::new()
                     };
 
-                    manifest_node.hyperspeed_restoration = Some(RestorationData {
-                        template: Some(template),
+                    // Create minimal meta structure - runtime will populate branches after parsing restored DOM
+                    manifest_node.meta = Some(serde_json::json!({
+                        "expression": expression
+                    }));
+
+                    // Runtime state
+                    manifest_node.runtime = Some(serde_json::json!({
+                        "activeBranch": null,
+                        "activeInstance": null,
+                        "templateRemoved": false
+                    }));
+
+                    // Compiled restoration data
+                    manifest_node.compiled = Some(CompiledData {
+                        restoration: Some(RestorationData {
+                            template: Some(template),
+                            ..Default::default()
+                        }),
                         ..Default::default()
                     });
                 }
@@ -304,6 +440,11 @@ impl ManifestBuilder {
             }
         }
 
+        // Fix name bindings: Remove ="" added by html5ever serializer
+        // Name bindings like <icon @[section.icon]> get serialized as <icon @[section.icon]="">
+        // which is invalid HTML that browsers reject
+        html = self.name_binding_fix_regex.replace_all(&html, "$1").to_string();
+
         html
     }
 }
@@ -312,13 +453,26 @@ impl ManifestBuilder {
 pub struct ManifestNode {
     element: Option<()>,  // Always null in static manifest
     parsed: Vec<String>,
-    children: HashMap<String, ManifestNode>,
+    children: BTreeMap<String, ManifestNode>,
 
     #[allow(dead_code)]
-    hyperspeed_restoration: Option<RestorationData>,
+    compiled: Option<CompiledData>,
 
     #[allow(dead_code)]
     node_type: Option<String>,
+
+    // Runtime-compatible fields
+    #[allow(dead_code)]
+    attributes: Option<BTreeMap<String, String>>,
+
+    #[allow(dead_code)]
+    name_bindings: Option<Vec<Option<String>>>,
+
+    #[allow(dead_code)]
+    meta: Option<Value>,  // For conditionals and iterations
+
+    #[allow(dead_code)]
+    runtime: Option<Value>,  // Runtime state
 }
 
 // Custom Serialize implementation to control output format
@@ -340,9 +494,29 @@ impl Serialize for ManifestNode {
         // Always include children
         map.serialize_entry("children", &self.children)?;
 
-        // Optional: hyperspeedRestoration
-        if let Some(ref restoration) = self.hyperspeed_restoration {
-            map.serialize_entry("hyperspeedRestoration", restoration)?;
+        // Optional: attributes (runtime field)
+        if let Some(ref attributes) = self.attributes {
+            map.serialize_entry("attributes", attributes)?;
+        }
+
+        // Optional: nameBindings (runtime field)
+        if let Some(ref name_bindings) = self.name_bindings {
+            map.serialize_entry("nameBindings", name_bindings)?;
+        }
+
+        // Optional: meta (runtime field for conditionals/iterations)
+        if let Some(ref meta) = self.meta {
+            map.serialize_entry("meta", meta)?;
+        }
+
+        // Optional: runtime (runtime state)
+        if let Some(ref runtime) = self.runtime {
+            map.serialize_entry("runtime", runtime)?;
+        }
+
+        // Optional: compiled (contains restoration and iterations data)
+        if let Some(ref compiled) = self.compiled {
+            map.serialize_entry("compiled", compiled)?;
         }
 
         // Optional: type
@@ -360,10 +534,37 @@ pub struct RestorationData {
     pub parsed: Option<Vec<String>>,
 
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub attributes: Option<HashMap<String, String>>,
+    pub attributes: Option<BTreeMap<String, String>>,
 
     #[serde(skip_serializing_if = "Option::is_none")]
     pub template: Option<String>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub expression: Option<String>,
+
+    #[serde(skip_serializing_if = "Option::is_none", rename = "nameBindings")]
+    pub name_bindings: Option<Vec<Option<String>>>,
+}
+
+#[derive(Debug, Serialize, Default)]
+pub struct IterationData {
+    #[serde(skip_serializing_if = "Option::is_none", rename = "batchFn")]
+    pub batch_fn: Option<String>,
+
+    #[serde(skip_serializing_if = "Option::is_none", rename = "itemAlias")]
+    pub item_alias: Option<String>,
+
+    #[serde(skip_serializing_if = "Option::is_none", rename = "indexAlias")]
+    pub index_alias: Option<String>,
+}
+
+#[derive(Debug, Serialize, Default)]
+pub struct CompiledData {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub restoration: Option<RestorationData>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub iterations: Option<IterationData>,
 }
 
 #[cfg(test)]
@@ -376,7 +577,7 @@ mod tests {
         let html = r#"<div>Hello @[name]</div>"#;
         let state = json!({ "name": "World" });
         let builder = ManifestBuilder::new();
-        let manifest = builder.build_from_html(html, &state).unwrap();
+        let manifest = builder.build_from_html(html, &state, false).unwrap();
 
         // Should have children
         assert!(!manifest.children.is_empty());
@@ -387,7 +588,7 @@ mod tests {
         let html = r#"<input value="@[firstName]">"#;
         let state = json!({ "firstName": "John" });
         let builder = ManifestBuilder::new();
-        let manifest = builder.build_from_html(html, &state).unwrap();
+        let manifest = builder.build_from_html(html, &state, false).unwrap();
 
         // Should have children
         assert!(!manifest.children.is_empty());
@@ -398,7 +599,7 @@ mod tests {
         let html = r#"<body><!-- each items as item --><div>@[item]</div><!-- /each --></body>"#;
         let state = json!({ "items": [1, 2, 3] });
         let builder = ManifestBuilder::new();
-        let manifest = builder.build_from_html(html, &state).unwrap();
+        let manifest = builder.build_from_html(html, &state, false).unwrap();
 
         // Serialize to JSON to inspect
         let json = serde_json::to_string_pretty(&manifest).unwrap();
@@ -429,4 +630,28 @@ mod tests {
         let result = builder.split_by_bindings("Hello World");
         assert_eq!(result, vec!["Hello World"]);
     }
+
+    #[test]
+    fn build_multiple_name_bindings() {
+        let html = r#"<div test @[theme] @[size]></div>"#;
+        let state = json!({ "theme": "light", "size": "large" });
+        let builder = ManifestBuilder::new();
+        let manifest = builder.build_from_html(html, &state, false).unwrap();
+
+        // Navigate to div
+        let html_node = manifest.children.get("html_0").unwrap();
+        let body_node = html_node.children.get("body_1").unwrap();
+        let div_node = body_node.children.get("div_0").unwrap();
+
+        // Check name bindings (sparse array indexed by position)
+        let compiled = div_node.compiled.as_ref().unwrap();
+        let restoration = compiled.restoration.as_ref().unwrap();
+        let name_bindings = restoration.name_bindings.as_ref().unwrap();
+
+        // Attributes: [0] test, [1] @[theme], [2] @[size]
+        assert_eq!(name_bindings.len(), 3);
+        assert_eq!(name_bindings[0], None);  // "test" is not a binding
+        assert_eq!(name_bindings[1], Some("@[theme]".to_string()));
+        assert_eq!(name_bindings[2], Some("@[size]".to_string()));
+    }
 }

package/compiler/src/compiler/mod.rs

@@ -2,7 +2,11 @@ pub mod compile;
 mod state_extractor;
 mod manifest_builder;
 mod value_stamper;
+mod iteration_optimizer;
+mod js_analyzer;
+mod component_tagger;
+pub mod watcher;
 
 pub use compile::Compiler;
 #[allow(unused_imports)]
-pub use compile::CompileStats;
+pub use compile::{CompileStats, ManifestStats};

package/compiler/src/compiler/state_extractor.rs

@@ -1,41 +1,128 @@
 use regex::Regex;
 use serde_json::{Value, Map};
+use std::path::PathBuf;
+use crate::compiler::js_analyzer::JsAnalyzer;
 
 pub struct StateExtractor;
 
 impl StateExtractor {
     /// Extract state from all vibe() and component() calls in HTML
-    pub fn extract_from_html(html: &str) -> Result<Value, String> {
+    pub fn extract_from_html(html: &str, base_path: &PathBuf) -> Result<Value, String> {
+        // Try modern approach: parse script blocks with JsAnalyzer
+        if let Some(state) = Self::extract_from_scripts(html, base_path) {
+            return Ok(state);
+        }
+
+        // Fallback: legacy regex-based extraction for inline object literals
+        Self::extract_with_regex(html)
+    }
+
+    /// Extract state from <script> blocks using JavaScript AST analysis
+    fn extract_from_scripts(html: &str, base_path: &PathBuf) -> Option<Value> {
+        let script_regex = Regex::new(r#"(?s)<script[^>]*>(.*?)</script>"#).unwrap();
+        let mut analyzer = JsAnalyzer::new(base_path.clone());
+
         let mut merged_state = Map::new();
 
-        // Find vibe({ ... }) and state({ ... }) patterns
-        let vibe_regex = Regex::new(
-            r"(?:vibe|state)\s*\(\s*\{([^}]+(?:\{[^}]*\}[^}]*)*)\}\s*[,)]"
-        ).map_err(|e| format!("Failed to compile vibe regex: {}", e))?;
+        for cap in script_regex.captures_iter(html) {
+            let script = cap.get(1)?.as_str();
 
-        for caps in vibe_regex.captures_iter(html) {
-            let obj_literal = &caps[1];
-            let state = Self::parse_object_literal(obj_literal)?;
-            Self::merge_into(&mut merged_state, state);
+            // Parse and extract state with AST analysis
+            if let Some(Value::Object(state)) = analyzer.extract_state(script) {
+                for (k, v) in state {
+                    merged_state.insert(k, v);
+                }
+            }
         }
 
-        // Find component({ ... }) patterns
-        let component_regex = Regex::new(
-            r"component\s*\(\s*\{([^}]+(?:\{[^}]*\}[^}]*)*)\}\s*[,)]"
-        ).map_err(|e| format!("Failed to compile component regex: {}", e))?;
+        if merged_state.is_empty() {
+            None
+        } else {
+            Some(Value::Object(merged_state))
+        }
+    }
 
-        for caps in component_regex.captures_iter(html) {
-            let obj_literal = &caps[1];
-            let component_state = Self::parse_object_literal(obj_literal)?;
+    /// Legacy regex-based extraction (fallback)
+    fn extract_with_regex(html: &str) -> Result<Value, String> {
+        let mut merged_state = Map::new();
+
+        // Find vibe({ ... }) and state({ ... }) patterns using brace counting
+        let start_regex = Regex::new(r"(?:vibe|state)\s*\(\s*\{").unwrap();
 
-            // Component state needs namespacing with component ID
-            // For now, merge directly (components already resolved)
-            Self::merge_into(&mut merged_state, component_state);
+        for start_match in start_regex.find_iter(html) {
+            if let Some(obj_literal) = Self::extract_balanced_object(&html[start_match.end()..]) {
+                match Self::parse_object_literal(&obj_literal) {
+                    Ok(state) => Self::merge_into(&mut merged_state, state),
+                    Err(_) => continue, // Skip unparseable state
+                }
+            }
+        }
+
+        // Find component({ ... }) patterns
+        let component_start_regex = Regex::new(r"component\s*\(\s*\{").unwrap();
+
+        for start_match in component_start_regex.find_iter(html) {
+            if let Some(obj_literal) = Self::extract_balanced_object(&html[start_match.end()..]) {
+                match Self::parse_object_literal(&obj_literal) {
+                    Ok(component_state) => Self::merge_into(&mut merged_state, component_state),
+                    Err(_) => continue,
+                }
+            }
         }
 
         Ok(Value::Object(merged_state))
     }
 
+    /// Extract balanced object literal content (everything between { and matching })
+    /// Input should start right after the opening {
+    fn extract_balanced_object(s: &str) -> Option<String> {
+        let mut depth = 1;
+        let mut in_string = false;
+        let mut escape_next = false;
+        let mut quote_char = '\0';
+        let mut result = String::new();
+
+        for ch in s.chars() {
+            if escape_next {
+                result.push(ch);
+                escape_next = false;
+                continue;
+            }
+
+            if ch == '\\' {
+                result.push(ch);
+                escape_next = true;
+                continue;
+            }
+
+            if !in_string {
+                if ch == '"' || ch == '\'' || ch == '`' {
+                    in_string = true;
+                    quote_char = ch;
+                    result.push(ch);
+                } else if ch == '{' || ch == '[' {
+                    depth += 1;
+                    result.push(ch);
+                } else if ch == '}' || ch == ']' {
+                    depth -= 1;
+                    if depth == 0 {
+                        return Some(result);
+                    }
+                    result.push(ch);
+                } else {
+                    result.push(ch);
+                }
+            } else {
+                result.push(ch);
+                if ch == quote_char {
+                    in_string = false;
+                }
+            }
+        }
+
+        None
+    }
+
     /// Parse JavaScript object literal to JSON Value
     fn parse_object_literal(js: &str) -> Result<Map<String, Value>, String> {
         // 1. Normalize to valid JSON
@@ -84,14 +171,14 @@ mod tests {
     #[test]
     fn extract_simple_state() {
         let html = r#"<script>vibe({ count: 0 })</script>"#;
-        let state = StateExtractor::extract_from_html(html).unwrap();
+        let state = StateExtractor::extract_from_html(html, &PathBuf::from(".")).unwrap();
         assert_eq!(state["count"], 0);
     }
 
     #[test]
     fn extract_nested_state() {
         let html = r#"vibe({ user: { name: 'John', age: 30 } })"#;
-        let state = StateExtractor::extract_from_html(html).unwrap();
+        let state = StateExtractor::extract_from_html(html, &PathBuf::from(".")).unwrap();
         assert_eq!(state["user"]["name"], "John");
         assert_eq!(state["user"]["age"], 30);
     }
@@ -99,7 +186,7 @@ mod tests {
     #[test]
     fn extract_array_state() {
         let html = r#"vibe({ items: [1, 2, 3] })"#;
-        let state = StateExtractor::extract_from_html(html).unwrap();
+        let state = StateExtractor::extract_from_html(html, &PathBuf::from(".")).unwrap();
         assert_eq!(state["items"][0], 1);
         assert_eq!(state["items"][1], 2);
         assert_eq!(state["items"][2], 3);
@@ -111,7 +198,7 @@ mod tests {
             <script>vibe({ count: 0 })</script>
             <script>vibe({ name: 'Test' })</script>
         "#;
-        let state = StateExtractor::extract_from_html(html).unwrap();
+        let state = StateExtractor::extract_from_html(html, &PathBuf::from(".")).unwrap();
         assert_eq!(state["count"], 0);
         assert_eq!(state["name"], "Test");
     }
@@ -119,14 +206,14 @@ mod tests {
     #[test]
     fn extract_with_trailing_comma() {
         let html = r#"vibe({ count: 0, })"#;
-        let state = StateExtractor::extract_from_html(html).unwrap();
+        let state = StateExtractor::extract_from_html(html, &PathBuf::from(".")).unwrap();
         assert_eq!(state["count"], 0);
     }
 
     #[test]
     fn extract_component_state() {
         let html = r#"component({ title: 'Hello' })"#;
-        let state = StateExtractor::extract_from_html(html).unwrap();
+        let state = StateExtractor::extract_from_html(html, &PathBuf::from(".")).unwrap();
         assert_eq!(state["title"], "Hello");
     }
 
@@ -145,4 +232,32 @@ mod tests {
         let normalized = StateExtractor::normalize_js_to_json(js);
         assert!(normalized.contains(r#""John""#));
     }
+
+    #[test]
+    fn extract_array_of_objects() {
+        let html = r#"vibe({
+            categories: [
+                { name: 'Fruits', items: ['Apple', 'Banana'] },
+                { name: 'Veggies', items: ['Carrot'] }
+            ]
+        })"#;
+        let state = StateExtractor::extract_from_html(html, &PathBuf::from(".")).unwrap();
+        assert_eq!(state["categories"][0]["name"], "Fruits");
+        assert_eq!(state["categories"][0]["items"][0], "Apple");
+        assert_eq!(state["categories"][1]["name"], "Veggies");
+    }
+
+    #[test]
+    fn extract_balanced_object_simple() {
+        let input = "count: 0 }";
+        let result = StateExtractor::extract_balanced_object(input).unwrap();
+        assert_eq!(result, "count: 0 ");
+    }
+
+    #[test]
+    fn extract_balanced_object_nested() {
+        let input = "user: { name: 'John', items: [1, 2, 3] } }";
+        let result = StateExtractor::extract_balanced_object(input).unwrap();
+        assert_eq!(result, "user: { name: 'John', items: [1, 2, 3] } ");
+    }
 }