From 624d4c88a100a9f4245b2d531322bb5f58719147 Mon Sep 17 00:00:00 2001
From: Tienson Qin <tiensonqin@gmail.com>
Date: Wed, 4 Nov 2026 20:01:41 +0000
Subject: [PATCH] fix(query): preserve bound terms in same-entity patterns

---
 impl/query_where.ml             | 23 +++++++++++++----------
 test/test_query_engine_fixes.ml | 31 +++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/impl/query_where.ml b/impl/query_where.ml
index d3fb83a..fdfe3c4 100644
--- a/impl/query_where.ml
+++ b/impl/query_where.ml
@@ -1184,18 +1184,21 @@ end) = struct
       then
         None
       else
-        let value_var_patterns, constant_patterns, required_patterns =
+        let* value_var_patterns, constant_patterns, required_patterns =
           patterns
           |> List.fold_left
-               (fun (value_vars, constants, required) (_, attr, value_term) ->
-                 match value_term with
-                 | QVar value_var when value_var <> e_var ->
-                   ((value_var, attr) :: value_vars, constants, required)
-                 | QValue value -> (value_vars, (attr, value) :: constants, required)
-                 | QWildcard -> (value_vars, constants, attr :: required)
-                 | QVar _ | QEntity _ | QAttr _ | QIdent _ | QLookupRef _ | QSource _ ->
-                   (value_vars, constants, required))
-               ([], [], [])
+               (fun acc (_, attr, value_term) ->
+                 match acc, value_term with
+                 | Some (value_vars, constants, required), QVar value_var when value_var <> e_var ->
+                   Some ((value_var, attr) :: value_vars, constants, required)
+                 | Some (value_vars, constants, required), QValue value ->
+                   Some (value_vars, (attr, value) :: constants, required)
+                 | Some (value_vars, constants, required), QWildcard ->
+                   Some (value_vars, constants, attr :: required)
+                 | None, _
+                 | Some _, (QVar _ | QEntity _ | QAttr _ | QIdent _ | QLookupRef _ | QSource _) ->
+                   None)
+               (Some ([], [], []))
         in
         let duplicate_value_var =
           let seen = Hashtbl.create (List.length value_var_patterns) in
diff --git a/test/test_query_engine_fixes.ml b/test/test_query_engine_fixes.ml
index 5399624..85c3364 100644
--- a/test/test_query_engine_fixes.ml
+++ b/test/test_query_engine_fixes.ml
@@ -317,6 +317,36 @@ let test_collection_in_binding () =
     (q_string ~inputs:[ Arg_collection [ Result_value (Int64 1L); Result_value (Int64 3L) ] ] db
        "[:find ?e :in $ [?x ...] :where [?e :attr ?x]]")
 
+let test_singleton_ref_collection_keeps_pattern_constraints () =
+  let schema =
+    schema_of_edn_string
+      "{:block/tags {:db/valueType :db.type/ref :db/cardinality :db.cardinality/many}}"
+  in
+  let db =
+    empty_db ~schema ()
+    |> db_with_string
+         "[{:db/id 1 :block/uuid \"untagged\"}
+           {:db/id 2 :block/uuid \"other-tag\" :block/tags [100]}
+           {:db/id 3 :block/uuid \"card\" :block/tags [200]}]"
+  in
+  List.iter
+    (fun (label, tags, expected) ->
+      assert_rows label expected
+        (q_string db
+           "[:find [?b ...] :in $ [?t ...] ?now-inst-ms %
+             :where [?b :block/tags ?t] [?b :block/uuid]]"
+           ~inputs:
+             [ Arg_collection tags
+             ; Arg_scalar (Result_value (Int64 0L))
+             ; Arg_rules []
+             ]))
+    [ "singleton ref keeps the tag constraint", [ Result_entity 200 ], [ [ Result_entity 3 ] ]
+    ; "missing ref returns no entities", [ Result_entity 201 ], []
+    ; "multiple refs keep the tag constraint", [ Result_entity 200; Result_entity 201 ], [ [ Result_entity 3 ] ]
+    ; "ref values preserve the same result", [ Result_value (Ref 200) ], [ [ Result_entity 3 ] ]
+    ; "empty input returns no entities", [], []
+    ]
+
 (* Bug 2 extended: comparison predicates over :in-bound vars *)
 let test_comparison_predicates_over_in () =
   let db =
@@ -1068,6 +1098,7 @@ let () =
     ; "rule_joined_with_clauses", test_rule_joined_with_clauses
     ; "predicate_over_in_scalar", test_predicate_over_in_scalar
     ; "collection_in_binding", test_collection_in_binding
+    ; "singleton_ref_collection_keeps_pattern_constraints", test_singleton_ref_collection_keeps_pattern_constraints
     ; "comparison_predicates_over_in", test_comparison_predicates_over_in
     ; "predicate_over_collection_in", test_predicate_over_collection_in
     ; "entity_map_lookup_ref_earlier_tx_entity", test_entity_map_lookup_ref_earlier_tx_entity