[LQP] Add provenance to LQP parser (#220)

Add provenance (source location) tracking to the LQP parser across all SDKs (Python, Go, Julia). Update validator to use it.

Author: nystrom

.github/workflows/ci.yml

@@ -59,7 +59,7 @@ jobs:
     strategy:
       fail-fast: true
       matrix:
-        julia-version: ["1.10", "1.11", "1.12"]
+        julia-version: ["1.10", "1.12"]
 
     steps:
     - uses: actions/checkout@v4

Makefile

@@ -3,6 +3,7 @@
 # Usage:
 #   make              Build protobuf bindings, parsers, and printers.
 #   make protobuf     Lint, check breaking changes, and generate protobuf code.
+#   make force-protobuf  Force-regenerate all protobuf code.
 #   make parsers      Regenerate Python, Julia, and Go parsers from the grammar.
 #   make parser-X     Regenerate a single parser (X = python, julia, go).
 #   make force-parsers      Force-regenerate all parsers.
@@ -75,7 +76,7 @@ JL_PROTO_GENERATED := \
 	$(JL_PROTO_DIR)/fragments_pb.jl \
 	$(JL_PROTO_DIR)/transactions_pb.jl
 
-.PHONY: all protobuf parsers parser-python parser-julia parser-go \
+.PHONY: all protobuf force-protobuf parsers parser-python parser-julia parser-go \
 	force-parsers force-parser-python force-parser-julia force-parser-go \
 	printers printer-python printer-julia printer-go \
 	force-printers force-printer-python force-printer-julia force-printer-go \
@@ -93,6 +94,23 @@ protobuf: $(PY_PROTO_GENERATED) $(GO_PROTO_GENERATED) $(JL_PROTO_GENERATED)
 touch-proto-generated:
 	touch $(PY_PROTO_GENERATED) $(GO_PROTO_GENERATED) $(JL_PROTO_GENERATED)
 
+force-protobuf:
+	buf lint
+	buf breaking --against ".git#branch=main,subdir=proto"
+	buf generate
+	mkdir -p $(PY_PROTO_DIR)
+	cp gen/python/relationalai/lqp/v1/*_pb2.py* $(PY_PROTO_DIR)/
+	for file in $(PY_PROTO_DIR)/*_pb2.py*; do \
+		sed 's/from relationalai\.lqp\.v1/from lqp\.proto\.v1/g' "$$file" > "$$file.tmp" && \
+		mv "$$file.tmp" "$$file"; \
+		sed 's/import relationalai\.lqp\.v1/import lqp\.proto\.v1/g' "$$file" > "$$file.tmp" && \
+		mv "$$file.tmp" "$$file"; \
+	done
+	mkdir -p $(GO_PROTO_DIR)
+	cp gen/go/relationalai/lqp/v1/*.pb.go $(GO_PROTO_DIR)/
+	rm -rf gen/python gen/go
+	cd sdks/julia && julia --project=LogicalQueryProtocol.jl generate_proto.jl
+
 $(PY_PROTO_GENERATED) $(GO_PROTO_GENERATED): $(PROTO_FILES)
 	buf lint
 	buf breaking --against ".git#branch=main,subdir=proto"

meta/src/meta/codegen_python.py

@@ -174,7 +174,7 @@ def gen_list_type(self, element_type: str) -> str:
         return f"list[{element_type}]"
 
     def gen_option_type(self, element_type: str) -> str:
-        return f"Optional[{element_type}]"
+        return f"{element_type} | None"
 
     def gen_dict_type(self, key_type: str, value_type: str) -> str:
         return f"dict[{key_type}, {value_type}]"

meta/src/meta/codegen_templates.py

@@ -106,6 +106,9 @@ class BuiltinTemplate:
     "format_bytes": BuiltinTemplate('"0x" + {0}.hex()'),
     "pp_dispatch": BuiltinTemplate("None", ["self.pprint_dispatch({0})"]),
     "get_at": BuiltinTemplate("{0}[{1}]"),
+    # Provenance tracking
+    "span_start": BuiltinTemplate("self.span_start()"),
+    "record_span": BuiltinTemplate("None", ["self.record_span({0}, {1})"]),
 }
 
 
@@ -191,6 +194,9 @@ class BuiltinTemplate:
     "format_bytes": BuiltinTemplate('"0x" * bytes2hex({0})'),
     "pp_dispatch": BuiltinTemplate("nothing", ["_pprint_dispatch(pp, {0})"]),
     "get_at": BuiltinTemplate("{0}[{1} + 1]"),
+    # Provenance tracking
+    "span_start": BuiltinTemplate("span_start(parser)"),
+    "record_span": BuiltinTemplate("nothing", ["record_span!(parser, {0}, {1})"]),
 }
 
 
@@ -280,6 +286,9 @@ class BuiltinTemplate:
     "format_bytes": BuiltinTemplate('fmt.Sprintf("0x%x", {0})'),
     "pp_dispatch": BuiltinTemplate("nil", ["p.pprintDispatch({0})"]),
     "get_at": BuiltinTemplate("{0}[{1}]"),
+    # Provenance tracking
+    "span_start": BuiltinTemplate("int64(p.spanStart())"),
+    "record_span": BuiltinTemplate("nil", ["p.recordSpan(int({0}), {1})"]),
 }
 
 __all__ = [

meta/src/meta/parser_gen.py

@@ -90,6 +90,7 @@
     Terminal,
 )
 from .grammar_utils import is_epsilon, rhs_elements
+from .proto_ast import ProtoMessage
 from .target import (
     Assign,
     BaseType,
@@ -100,6 +101,7 @@
     ListExpr,
     ListType,
     Lit,
+    MessageType,
     ParseNonterminal,
     ParseNonterminalDef,
     Seq,
@@ -132,28 +134,55 @@ class AmbiguousGrammarError(Exception):
 
 
 def generate_parse_functions(
-    grammar: Grammar, indent: str = ""
+    grammar: Grammar,
+    indent: str = "",
+    proto_messages: dict[tuple[str, str], ProtoMessage] | None = None,
 ) -> list[ParseNonterminalDef]:
     parser_methods = []
     reachable, _ = grammar.analysis.partition_nonterminals_by_reachability()
     for nt in reachable:
         rules = grammar.rules[nt]
-        method_code = _generate_parse_method(nt, rules, grammar, indent)
+        method_code = _generate_parse_method(nt, rules, grammar, indent, proto_messages)
         parser_methods.append(method_code)
     return parser_methods
 
 
+def _wrap_with_span(body: TargetExpr, return_type) -> TargetExpr:
+    """Wrap a nonterminal body with span_start/record_span."""
+    span_var = Var(gensym("span_start"), BaseType("Int64"))
+    result_var = Var(gensym("result"), return_type)
+    type_name = return_type.name if isinstance(return_type, MessageType) else ""
+    return Let(
+        span_var,
+        Call(make_builtin("span_start"), []),
+        Let(
+            result_var,
+            body,
+            Seq(
+                [
+                    Call(make_builtin("record_span"), [span_var, Lit(type_name)]),
+                    result_var,
+                ]
+            ),
+        ),
+    )
+
+
 def _generate_parse_method(
-    lhs: Nonterminal, rules: list[Rule], grammar: Grammar, indent: str = ""
+    lhs: Nonterminal,
+    rules: list[Rule],
+    grammar: Grammar,
+    indent: str = "",
+    proto_messages: dict[tuple[str, str], ProtoMessage] | None = None,
 ) -> ParseNonterminalDef:
-    """Generate parse method code as string (preserving existing logic)."""
+    """Generate parse method for a nonterminal with provenance tracking."""
     return_type = None
     rhs = None
     follow_set = FollowSet(grammar, lhs)
     if len(rules) == 1:
         rule = rules[0]
         rhs = _generate_parse_rhs_ir(
-            rule.rhs, grammar, follow_set, True, rule.constructor
+            rule.rhs, grammar, follow_set, True, rule.constructor, proto_messages
         )
         return_type = rule.constructor.return_type
     else:
@@ -171,7 +200,6 @@ def _generate_parse_method(
                 ],
             )
         for i, rule in enumerate(rules):
-            # Ensure the return type is the same for all actions for this nonterminal.
             assert return_type is None or return_type == rule.constructor.return_type, (
                 f"Return type mismatch at rule {i}: {return_type} != {rule.constructor.return_type}"
             )
@@ -183,12 +211,19 @@ def _generate_parse_method(
                     make_builtin("equal"), [Var(prediction, BaseType("Int64")), Lit(i)]
                 ),
                 _generate_parse_rhs_ir(
-                    rule.rhs, grammar, follow_set, True, rule.constructor
+                    rule.rhs,
+                    grammar,
+                    follow_set,
+                    True,
+                    rule.constructor,
+                    proto_messages,
                 ),
                 tail,
             )
         rhs = Let(Var(prediction, BaseType("Int64")), predictor, tail)
     assert return_type is not None
+    if isinstance(return_type, MessageType):
+        rhs = _wrap_with_span(rhs, return_type)
     return ParseNonterminalDef(lhs, [], return_type, rhs, indent)
 
 
@@ -377,30 +412,19 @@ def _generate_parse_rhs_ir(
     follow_set: TerminalSequenceSet,
     apply_action: bool = False,
     action: Lambda | None = None,
+    proto_messages: dict[tuple[str, str], ProtoMessage] | None = None,
 ) -> TargetExpr:
-    """Generate IR for parsing an RHS.
-
-    Args:
-        rhs: The RHS to parse
-        grammar: The grammar
-        follow_set: TerminalSequenceSet for computing follow lazily
-        apply_action: Whether to apply the semantic action
-        action: The semantic action to apply (required if apply_action is True)
-
-    Returns IR expression for leaf nodes (Literal, Terminal, Nonterminal).
-    Returns None for complex cases that still use string generation.
-    """
+    """Generate IR for parsing an RHS with provenance tracking."""
     if isinstance(rhs, Sequence):
         return _generate_parse_rhs_ir_sequence(
-            rhs, grammar, follow_set, apply_action, action
+            rhs, grammar, follow_set, apply_action, action, proto_messages
         )
     elif isinstance(rhs, LitTerminal):
         parse_expr = Call(make_builtin("consume_literal"), [Lit(rhs.name)])
         if apply_action and action:
             return Seq([parse_expr, apply_lambda(action, [])])
         return parse_expr
     elif isinstance(rhs, NamedTerminal):
-        # Use terminal's actual type for consume_terminal instead of generic Token
         from .target import FunctionType
 
         terminal_type = rhs.target_type()
@@ -427,14 +451,18 @@ def _generate_parse_rhs_ir(
     elif isinstance(rhs, Option):
         assert grammar is not None
         predictor = _build_option_predictor(grammar, rhs.rhs, follow_set)
-        parse_result = _generate_parse_rhs_ir(rhs.rhs, grammar, follow_set, False, None)
+        parse_result = _generate_parse_rhs_ir(
+            rhs.rhs, grammar, follow_set, False, None, proto_messages
+        )
         return IfElse(predictor, Call(make_builtin("some"), [parse_result]), Lit(None))
     elif isinstance(rhs, Star):
         assert grammar is not None
         xs = Var(gensym("xs"), ListType(rhs.rhs.target_type()))
         cond = Var(gensym("cond"), BaseType("Boolean"))
         predictor = _build_option_predictor(grammar, rhs.rhs, follow_set)
-        parse_item = _generate_parse_rhs_ir(rhs.rhs, grammar, follow_set, False, None)
+        parse_item = _generate_parse_rhs_ir(
+            rhs.rhs, grammar, follow_set, False, None, proto_messages
+        )
         item = Var(gensym("item"), rhs.rhs.target_type())
         loop_body = Seq(
             [
@@ -458,6 +486,7 @@ def _generate_parse_rhs_ir_sequence(
     follow_set: TerminalSequenceSet,
     apply_action: bool = False,
     action: Lambda | None = None,
+    proto_messages: dict[tuple[str, str], ProtoMessage] | None = None,
 ) -> TargetExpr:
     if is_epsilon(rhs):
         return Lit(None)
@@ -473,7 +502,9 @@ def _generate_parse_rhs_ir_sequence(
             follow_set_i = ConcatSet(first_following, follow_set)
         else:
             follow_set_i = follow_set
-        elem_ir = _generate_parse_rhs_ir(elem, grammar, follow_set_i, False, None)
+        elem_ir = _generate_parse_rhs_ir(
+            elem, grammar, follow_set_i, False, None, proto_messages
+        )
         if isinstance(elem, LitTerminal):
             exprs.append(elem_ir)
         else:
@@ -494,13 +525,10 @@ def _generate_parse_rhs_ir_sequence(
         lambda_call = apply_lambda(action, arg_vars)
         exprs.append(lambda_call)
     elif len(arg_vars) > 1:
-        # Multiple values - wrap in tuple
         exprs.append(Call(make_builtin("tuple"), arg_vars))
     elif len(arg_vars) == 1:
-        # Single value - return the variable
         exprs.append(arg_vars[0])
     else:
-        # no non-literal elements, return None
         return Lit(None)
 
     if len(exprs) == 1:

meta/src/meta/parser_gen_common.py

@@ -19,6 +19,7 @@ def generate_parser(
     codegen: CodeGenerator,
     template_path: Path,
     command_line: str | None = None,
+    proto_messages=None,
 ) -> str:
     """Generate a parser from a grammar using the given code generator and template.
 
@@ -27,11 +28,14 @@ def generate_parser(
         codegen: Language-specific code generator
         template_path: Path to the language-specific template file
         command_line: Optional command line string for the file header
+        proto_messages: Optional proto message definitions for provenance tracking
     """
     template = template_path.read_text()
     indent = codegen.parse_def_indent
 
-    defns = generate_parse_functions(grammar, indent=indent)
+    defns = generate_parse_functions(
+        grammar, indent=indent, proto_messages=proto_messages
+    )
     lines = []
     for defn in defns:
         lines.append("")

meta/src/meta/parser_gen_go.py

@@ -14,7 +14,9 @@ def generate_parser_go(
 ) -> str:
     """Generate LL(k) recursive-descent parser in Go."""
     codegen = GoCodeGenerator(proto_messages=proto_messages)
-    return generate_parser(grammar, codegen, _TEMPLATE_PATH, command_line)
+    return generate_parser(
+        grammar, codegen, _TEMPLATE_PATH, command_line, proto_messages
+    )
 
 
 __all__ = [

meta/src/meta/parser_gen_julia.py

@@ -15,7 +15,9 @@ def generate_parser_julia(
 ) -> str:
     """Generate LL(k) recursive-descent parser in Julia."""
     codegen = JuliaCodeGenerator(proto_messages=proto_messages, config=PARSER_CONFIG)
-    return generate_parser(grammar, codegen, _TEMPLATE_PATH, command_line)
+    return generate_parser(
+        grammar, codegen, _TEMPLATE_PATH, command_line, proto_messages
+    )
 
 
 __all__ = [

meta/src/meta/parser_gen_python.py

@@ -14,4 +14,6 @@ def generate_parser_python(
 ) -> str:
     """Generate LL(k) recursive-descent parser in Python."""
     codegen = PythonCodeGenerator(proto_messages=proto_messages)
-    return generate_parser(grammar, codegen, _TEMPLATE_PATH, command_line)
+    return generate_parser(
+        grammar, codegen, _TEMPLATE_PATH, command_line, proto_messages
+    )

meta/src/meta/target_builtins.py

@@ -249,6 +249,10 @@ def is_builtin(name: str) -> bool:
 # === Sequence indexing ===
 register_builtin("get_at", [SequenceType(T), INT64], T)
 
+# === Provenance tracking ===
+register_builtin("span_start", [], INT64)
+register_builtin("record_span", [INT64, STRING], VOID)
+
 
 # === Validation functions ===