From ec081c28e69ef33e172de3ec461e00ed069f3493 Mon Sep 17 00:00:00 2001
From: Sergei Zharinov <zharinov@users.noreply.github.com>
Date: Sat, 6 Dec 2025 13:09:40 -0300
Subject: [PATCH] fix: shape analysis crash on malformed Capture/Quantifier

---
 crates/plotnik-lib/src/query/invariants.rs   | 16 -----
 crates/plotnik-lib/src/query/shapes.rs       | 12 ++--
 crates/plotnik-lib/src/query/shapes_tests.rs | 65 ++++++++++++++++++++
 3 files changed, 72 insertions(+), 21 deletions(-)

diff --git a/crates/plotnik-lib/src/query/invariants.rs b/crates/plotnik-lib/src/query/invariants.rs
index f1db2cf8..5f6ce71f 100644
--- a/crates/plotnik-lib/src/query/invariants.rs
+++ b/crates/plotnik-lib/src/query/invariants.rs
@@ -34,22 +34,6 @@ pub fn ensure_both_branch_kinds<'a>(
     }
 }
 
-#[inline]
-pub fn ensure_capture_has_inner<T>(inner: Option<T>) -> T {
-    inner.expect(
-        "shape_cardinalities: Capture without inner expression \
-         (parser uses checkpoint, inner is guaranteed)",
-    )
-}
-
-#[inline]
-pub fn ensure_quantifier_has_inner<T>(inner: Option<T>) -> T {
-    inner.expect(
-        "shape_cardinalities: Quantifier without inner expression \
-         (parser uses checkpoint, inner is guaranteed)",
-    )
-}
-
 #[inline]
 pub fn ensure_ref_has_name<T>(name: Option<T>) -> T {
     name.expect(
diff --git a/crates/plotnik-lib/src/query/shapes.rs b/crates/plotnik-lib/src/query/shapes.rs
index 5a3fe840..8467a91a 100644
--- a/crates/plotnik-lib/src/query/shapes.rs
+++ b/crates/plotnik-lib/src/query/shapes.rs
@@ -8,9 +8,7 @@
 //! undefined refs, etc.).
 
 use super::Query;
-use super::invariants::{
-    ensure_capture_has_inner, ensure_quantifier_has_inner, ensure_ref_has_name,
-};
+use super::invariants::ensure_ref_has_name;
 use crate::diagnostics::DiagnosticKind;
 use crate::parser::{Expr, FieldExpr, Ref, SeqExpr, SyntaxNode};
 
@@ -55,12 +53,16 @@ impl Query<'_> {
             Expr::SeqExpr(seq) => self.seq_cardinality(seq),
 
             Expr::CapturedExpr(cap) => {
-                let inner = ensure_capture_has_inner(cap.inner());
+                let Some(inner) = cap.inner() else {
+                    return ShapeCardinality::Invalid;
+                };
                 self.get_or_compute(&inner)
             }
 
             Expr::QuantifiedExpr(q) => {
-                let inner = ensure_quantifier_has_inner(q.inner());
+                let Some(inner) = q.inner() else {
+                    return ShapeCardinality::Invalid;
+                };
                 self.get_or_compute(&inner)
             }
 
diff --git a/crates/plotnik-lib/src/query/shapes_tests.rs b/crates/plotnik-lib/src/query/shapes_tests.rs
index 9c132b68..392c99ae 100644
--- a/crates/plotnik-lib/src/query/shapes_tests.rs
+++ b/crates/plotnik-lib/src/query/shapes_tests.rs
@@ -413,3 +413,68 @@ fn invalid_ref_to_bodyless_def() {
         Ref⁻ X
     ");
 }
+
+#[test]
+fn invalid_capture_without_inner() {
+    // Error recovery: `extra` is invalid, but `@y` still creates a Capture node
+    let query = Query::try_from("(call extra @y)").unwrap();
+    assert!(!query.is_valid());
+    insta::assert_snapshot!(query.dump_with_cardinalities(), @r"
+    Root¹
+      Def¹
+        NamedNode¹ call
+          CapturedExpr⁻ @y
+    ");
+    insta::assert_snapshot!(query.dump_diagnostics(), @r"
+    error: bare identifier is not a valid expression; wrap in parentheses: `(identifier)`
+      |
+    1 | (call extra @y)
+      |       ^^^^^
+    ");
+}
+
+#[test]
+fn invalid_capture_without_inner_standalone() {
+    // Standalone capture without preceding expression
+    let query = Query::try_from("@x").unwrap();
+    assert!(!query.is_valid());
+    insta::assert_snapshot!(query.dump_diagnostics(), @r"
+    error: `@` must follow an expression to capture
+      |
+    1 | @x
+      | ^
+    ");
+}
+
+#[test]
+fn invalid_multiple_captures_with_error() {
+    let query = Query::try_from("(call (Undefined) @x extra @y)").unwrap();
+    assert!(!query.is_valid());
+    insta::assert_snapshot!(query.dump_with_cardinalities(), @r"
+    Root¹
+      Def¹
+        NamedNode¹ call
+          CapturedExpr⁻ @x
+            Ref⁻ Undefined
+          CapturedExpr⁻ @y
+    ");
+}
+
+#[test]
+fn invalid_quantifier_without_inner() {
+    // Error recovery: `extra` is invalid, but `*` still creates a Quantifier node
+    let query = Query::try_from("(foo extra*)").unwrap();
+    assert!(!query.is_valid());
+    insta::assert_snapshot!(query.dump_with_cardinalities(), @r"
+    Root¹
+      Def¹
+        NamedNode¹ foo
+          QuantifiedExpr⁻ *
+    ");
+    insta::assert_snapshot!(query.dump_diagnostics(), @r"
+    error: bare identifier is not a valid expression; wrap in parentheses: `(identifier)`
+      |
+    1 | (foo extra*)
+      |      ^^^^^
+    ");
+}