Pass spans from parser into AST

Author: aljazerzen

edb/_edgeql_parser.pyi

@@ -59,6 +59,8 @@ class CSTNode:
 class Production:
     id: int
     args: list[CSTNode]
+    start: int | None
+    end: int | None
 
 class Terminal:
     text: str

edb/common/parsing.py

@@ -99,6 +99,7 @@ def decorator(func: Any):
 
 
 class Nonterm(parsing.Nonterm):
+    span: Span
 
     def __init_subclass__(cls, *, is_internal=False, **kwargs):
         """Add docstrings to class and reduce functions
@@ -125,7 +126,6 @@ def __init_subclass__(cls, *, is_internal=False, **kwargs):
         for name, attr in cls.__dict__.items():
             if (name.startswith('reduce_') and
                     isinstance(attr, types.FunctionType)):
-                inline_index = getattr(attr, 'inline_index', None)
 
                 if attr.__doc__ is None:
                     tokens = name.split('_')
@@ -138,14 +138,11 @@ def __init_subclass__(cls, *, is_internal=False, **kwargs):
                     if prec is not None:
                         doc += ' [{}]'.format(prec)
 
+                    inline_index = getattr(attr, 'inline_index', None)
                     attr = lambda self, *args, meth=attr: meth(self, *args)
                     attr.__doc__ = doc
-
-                a = span.wrap_function_to_infer_spans(attr)
-
-                a.__doc__ = attr.__doc__
-                a.inline_index = inline_index
-                setattr(cls, name, a)
+                    attr.inline_index = inline_index
+                    setattr(cls, name, attr)
 
 
 class ListNonterm(Nonterm, is_internal=True):

edb/common/span.py

@@ -298,7 +298,9 @@ def generic_visit(self, node):
 class SpanValidator(ast.NodeVisitor):
     def generic_visit(self, node):
         if getattr(node, 'span', None) is None:
-            raise RuntimeError('node {} has no span'.format(node))
+            from edb.edgeql import ast as qlast
+            if not isinstance(node, (qlast.ShapeOperation, qlast.Options)):
+                raise RuntimeError('node {} has no span'.format(node))
         super().generic_visit(node)
 
 

edb/edgeql-parser/edgeql-parser-python/src/parser.rs

@@ -69,6 +69,10 @@ pub struct Production {
     id: usize,
     #[pyo3(get)]
     args: PyObject,
+    #[pyo3(get)]
+    start: Option<u64>,
+    #[pyo3(get)]
+    end: Option<u64>,
 }
 
 #[pyclass]
@@ -195,6 +199,8 @@ impl<'py> IntoPyObject<'py> for ParserCSTNode<'_> {
                     Production {
                         id: prod.id,
                         args: PyList::new(py, prod.args.iter().map(ParserCSTNode))?.into(),
+                        start: prod.span.map(|s| s.start),
+                        end: prod.span.map(|s| s.end),
                     },
                 )?),
                 terminal: None,

edb/edgeql-parser/src/parser/cst.rs

@@ -29,6 +29,7 @@ pub struct Terminal {
 pub struct Production<'a> {
     pub id: usize,
     pub args: &'a [CSTNode<'a>],
+    pub span: Option<Span>,
 
     /// When a production is inlined, its id is saved into the new production
     /// This is needed when matching CST nodes by production id.

edb/edgeql-parser/src/parser/mod.rs

@@ -394,6 +394,7 @@ impl<'s> Parser<'s> {
 
         let value = CSTNode::Production(Production {
             id: reduce.production_id,
+            span: get_span_of_nodes(args),
             args,
             inlined_ids: None,
         });
@@ -419,7 +420,7 @@ impl<'s> Parser<'s> {
                         Some(ctx.alloc_slice_and_push(&new_prod.inlined_ids, inlined_id));
                 }
 
-                extend_span(&mut value, span, ctx);
+                expand_span(&mut value, span, ctx);
             } else {
                 // place back
                 value = CSTNode::Production(production);
@@ -607,38 +608,54 @@ impl<'a> StackNode<'a> {
     }
 }
 
-fn get_span_of_nodes(args: &[CSTNode]) -> Option<Span> {
-    let start = args.iter().find_map(|x| match x {
+/// Returns the span of syntactically ordered nodes. Panic on empty nodes.
+fn get_span_of_nodes(nodes: &[CSTNode]) -> Option<Span> {
+    let start = nodes.iter().find_map(|x| match x {
         CSTNode::Terminal(t) => Some(t.span.start),
-        CSTNode::Production(p) => get_span_of_nodes(p.args).map(|x| x.start),
-        _ => None,
+        CSTNode::Production(p) => Some(p.span?.start),
+        CSTNode::Empty => panic!(),
     })?;
-    let end = args.iter().rev().find_map(|x| match x {
+    let end = nodes.iter().rev().find_map(|x| match x {
         CSTNode::Terminal(t) => Some(t.span.end),
-        CSTNode::Production(p) => get_span_of_nodes(p.args).map(|x| x.end),
-        _ => None,
+        CSTNode::Production(p) => Some(p.span?.end),
+        CSTNode::Empty => panic!(),
     })?;
     Some(Span { start, end })
 }
 
-fn extend_span<'a>(value: &mut CSTNode<'a>, span: Option<Span>, ctx: &'a Context) {
+fn expand_span<'a>(value: &mut CSTNode<'a>, span: Option<Span>, ctx: &'a Context) {
     let Some(span) = span else {
         return;
     };
 
-    let CSTNode::Terminal(terminal) = value else {
-        return;
+    let curr_span = match value {
+        CSTNode::Empty => panic!(),
+        CSTNode::Terminal(t) => &t.span,
+        CSTNode::Production(p) => {
+            if let Some(span) = &p.span {
+                span
+            } else {
+                p.span = Some(span);
+                return;
+            }
+        }
     };
 
-    let mut new_term = terminal.clone();
-
-    if span.start < new_term.span.start {
-        new_term.span.start = span.start;
-    }
-    if span.end > new_term.span.end {
-        new_term.span.end = span.end;
+    if curr_span.start <= span.start && span.end <= curr_span.end {
+        return;
     }
-    *terminal = ctx.alloc_terminal(new_term);
+
+    let span = span.extend(curr_span);
+
+    match value {
+        CSTNode::Empty => panic!(),
+        CSTNode::Terminal(t) => {
+            let mut new_term = t.clone();
+            new_term.span = span;
+            *t = ctx.alloc_terminal(new_term);
+        }
+        CSTNode::Production(p) => p.span = Some(span),
+    };
 }
 
 const PARSER_COUNT_MAX: usize = 10;

edb/edgeql-parser/src/position.rs

@@ -79,6 +79,13 @@ impl Span {
             end: right.end,
         }
     }
+
+    pub fn extend(self, other: &Span) -> Span {
+        Span {
+            start: u64::min(self.start, other.start),
+            end: u64::max(self.end, other.end),
+        }
+    }
 }
 
 fn new_lines_in_fragment(data: &[u8]) -> u64 {

edb/edgeql/ast.py

@@ -278,12 +278,14 @@ class Constant(BaseConstant):
     value: str
 
     @classmethod
-    def string(cls, value: str) -> Constant:
-        return Constant(kind=ConstantKind.STRING, value=value)
+    def string(cls, value: str, span: Span | None = None) -> Constant:
+        return Constant(kind=ConstantKind.STRING, value=value, span=span)
 
     @classmethod
-    def boolean(cls, b: bool) -> Constant:
-        return Constant(kind=ConstantKind.BOOLEAN, value=str(b).lower())
+    def boolean(cls, b: bool, span: Span | None = None) -> Constant:
+        return Constant(
+            kind=ConstantKind.BOOLEAN, value=str(b).lower(), span=span
+        )
 
     @classmethod
     def integer(cls, i: int) -> Constant:

edb/edgeql/parser/__init__.py

@@ -236,8 +236,23 @@ def _cst_to_ast(
             production_id = node.id
             non_term_type, method = productions[production_id]
             sym = non_term_type()
+
+            if node.start is not None and node.end is not None:
+                sym.span = parsing.Span(
+                    filename=filename,
+                    buffer=source.text(),
+                    start=node.start,
+                    end=node.end,
+                )
+            else:
+                sym.span = None
+
             method(sym, *args)
 
+            # span
+            if sym.span is not None and isinstance(sym.val, qlast.Base):
+                sym.val.span = sym.span
+
             # push into result stack
             result.append(sym)
 

edb/edgeql/parser/grammar/commondl.py

@@ -417,6 +417,8 @@ def reduce_USING_Identifier_EXPRESSION(self, _using, ident, _expression):
 
 
 class ProcessFunctionBlockMixin:
+    span: parsing.Span
+
     def _process_function_body(self, block, *, optional_using: bool=False):
         props: dict[str, typing.Any] = {}
 
@@ -483,6 +485,7 @@ def _process_function_body(self, block, *, optional_using: bool=False):
                 from_function=from_function,
                 from_expr=from_expr,
                 code=code,
+                span=self.span,
             )
 
             props['nativecode'] = nativecode