Skip to main content

substrait_explain/parser/
structural.rs

1//! Parser for the structural part of the Substrait file format.
2//!
3//! This is the overall parser for parsing the text format. It is responsible
4//! for tracking which section of the file we are currently parsing, and parsing
5//! each line separately.
6
7use std::fmt;
8
9use pest::iterators::Pair;
10use substrait::proto::extensions::AdvancedExtension;
11use substrait::proto::{
12    AggregateRel, CrossRel, FetchRel, FilterRel, JoinRel, Plan, PlanRel, ProjectRel, ReadRel, Rel,
13    RelRoot, SortRel, Version, plan_rel,
14};
15
16use crate::extensions::any::Any;
17use crate::extensions::{AddendumKind, ExtensionRegistry, SimpleExtensions, simple};
18use crate::parser::chunks::ChunkCursor;
19use crate::parser::common::{MessageParseError, ParsePair, ScopedParsePair};
20use crate::parser::errors::{ParseContext, ParseError, ParseResult};
21use crate::parser::expressions::Name;
22use crate::parser::extensions::{
23    AddendumInvocation, ExtensionInvocation, ExtensionParseError, ExtensionParser,
24};
25use crate::parser::relations::{
26    ExtensionReadRel, RelationParsingContext, VirtualReadRel, parse_set_relation_pair,
27};
28use crate::parser::{ErrorKind, ExpressionParser, RelationParsePair, Rule, unwrap_single_pair};
29
30pub const PLAN_HEADER: &str = "=== Plan";
31pub const VERSION_HEADER: &str = "=== Version";
32
33/// Represents an input line, trimmed of leading two-space indents and final
34/// whitespace. Contains the number of indents and the trimmed line.
35#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub struct IndentedLine<'a>(pub usize, pub &'a str);
37
38impl<'a> From<&'a str> for IndentedLine<'a> {
39    fn from(line: &'a str) -> Self {
40        let line = line.trim_end();
41        let mut spaces = 0;
42        for c in line.chars() {
43            if c == ' ' {
44                spaces += 1;
45            } else {
46                break;
47            }
48        }
49
50        let indents = spaces / 2;
51
52        let (_, trimmed) = line.split_at(indents * 2);
53
54        IndentedLine(indents, trimmed)
55    }
56}
57
58/// A `+`-prefixed addendum line attached to a relation node. The `pair` holds
59/// the grammar rule directly (already unwrapped from the outer `planNode`).
60#[derive(Debug, Clone)]
61pub struct Addendum<'a> {
62    pub pair: Pair<'a, Rule>, // Rule::addendum
63    pub line_no: i64,
64}
65
66/// A relation node in the plan tree, before conversion to a Substrait proto.
67#[derive(Debug, Clone)]
68pub struct RelationNode<'a> {
69    pub pair: Pair<'a, Rule>,
70    pub line_no: i64,
71    pub addenda: Vec<Addendum<'a>>,
72    pub children: Vec<RelationNode<'a>>,
73}
74
75impl<'a> RelationNode<'a> {
76    pub fn context(&self) -> ParseContext {
77        ParseContext {
78            line_no: self.line_no,
79            line: self.pair.as_str().to_string(),
80        }
81    }
82}
83
84/// A parsed plan line: either a relation or a `+`-prefixed addendum line.
85///
86/// Classification happens at construction time by inspecting the inner grammar
87/// rule, so downstream code can use standard Rust pattern matching rather than
88/// runtime rule inspection.
89#[derive(Debug, Clone)]
90pub enum LineNode<'a> {
91    Relation(RelationNode<'a>),
92    Addendum(Addendum<'a>),
93}
94
95impl<'a> LineNode<'a> {
96    pub fn parse(line: &'a str, line_no: i64) -> Result<Self, ParseError> {
97        let mut pairs: pest::iterators::Pairs<'a, Rule> =
98            <ExpressionParser as pest::Parser<Rule>>::parse(Rule::planNode, line).map_err(|e| {
99                ParseError::Plan(
100                    ParseContext {
101                        line_no,
102                        line: line.to_string(),
103                    },
104                    MessageParseError::new("planNode", ErrorKind::InvalidValue, Box::new(e)),
105                )
106            })?;
107
108        let outer = pairs.next().unwrap();
109        assert!(pairs.next().is_none()); // Should be exactly one pair
110        let inner = unwrap_single_pair(outer);
111
112        Ok(match inner.as_rule() {
113            Rule::addendum => LineNode::Addendum(Addendum {
114                pair: inner,
115                line_no,
116            }),
117            _ => LineNode::Relation(RelationNode {
118                pair: inner,
119                line_no,
120                addenda: Vec::new(),
121                children: Vec::new(),
122            }),
123        })
124    }
125
126    /// Parse the line as a top-level relation at depth 0 (either root_relation or regular relation)
127    pub fn parse_root(line: &'a str, line_no: i64) -> Result<Self, ParseError> {
128        let mut pairs: pest::iterators::Pairs<'a, Rule> = <ExpressionParser as pest::Parser<
129            Rule,
130        >>::parse(
131            Rule::top_level_relation, line
132        )
133        .map_err(|e| {
134            ParseError::Plan(
135                ParseContext::new(line_no, line.to_string()),
136                MessageParseError::new("top_level_relation", ErrorKind::Syntax, Box::new(e)),
137            )
138        })?;
139
140        let outer = pairs.next().unwrap();
141        assert!(pairs.next().is_none());
142
143        // top_level_relation is either root_relation or planNode.
144        // If planNode, unwrap one more level to obtain the specific relation rule.
145        let inner = unwrap_single_pair(outer);
146        let pair = if inner.as_rule() == Rule::planNode {
147            unwrap_single_pair(inner)
148        } else {
149            inner // root_relation
150        };
151
152        // planNode can include addenda (+Enh:, +Opt:, +Ext:); surface them so
153        // TreeBuilder::add_line can produce the appropriate depth-0 error.
154        if pair.as_rule() == Rule::addendum {
155            return Ok(LineNode::Addendum(Addendum { pair, line_no }));
156        }
157
158        Ok(LineNode::Relation(RelationNode {
159            pair,
160            line_no,
161            addenda: Vec::new(),
162            children: Vec::new(),
163        }))
164    }
165}
166
167#[derive(Copy, Clone, Debug, PartialEq, Eq)]
168pub enum State {
169    // The initial state, before we have parsed any lines.
170    Initial,
171    // The version section, after parsing the `=== Version` header. Accepts the
172    // optional indented `producer:` / `git_hash:` detail lines.
173    Version,
174    // The extensions section, after parsing the header and any other Extension lines.
175    Extensions,
176    // The plan section, after parsing the header and any other Plan lines.
177    Plan,
178}
179
180impl State {
181    /// Position of this section in the required document order (`Version`,
182    /// `Extensions`, `Plan`). Sections must appear in strictly increasing
183    /// order, so a `===` header may only transition to a state with a higher
184    /// rank than the current one.
185    fn section_rank(&self) -> u8 {
186        match self {
187            State::Initial => 0,
188            State::Version => 1,
189            State::Extensions => 2,
190            State::Plan => 3,
191        }
192    }
193}
194
195impl fmt::Display for State {
196    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
197        write!(f, "{self:?}")
198    }
199}
200
201// An in-progress tree builder, building the tree of relations.
202#[derive(Debug, Clone, Default)]
203pub struct TreeBuilder<'a> {
204    // Current tree of nodes being built. These have been successfully parsed
205    // into Pest pairs, but have not yet been converted to substrait plans.
206    current: Option<RelationNode<'a>>,
207    // Completed trees that have been built.
208    completed: Vec<RelationNode<'a>>,
209}
210
211impl<'a> TreeBuilder<'a> {
212    /// Traverse down the tree, always taking the last child at each level, until reaching the specified depth.
213    pub fn get_at_depth(&mut self, depth: usize) -> Option<&mut RelationNode<'a>> {
214        let mut node = self.current.as_mut()?;
215        for _ in 0..depth {
216            node = node.children.last_mut()?;
217        }
218        Some(node)
219    }
220
221    pub fn add_line(&mut self, depth: usize, node: LineNode<'a>) -> Result<(), ParseError> {
222        match node {
223            LineNode::Relation(rel_node) => {
224                if depth == 0 {
225                    if let Some(prev) = self.current.take() {
226                        self.completed.push(prev);
227                    }
228                    self.current = Some(rel_node);
229                    return Ok(());
230                }
231
232                let parent = match self.get_at_depth(depth - 1) {
233                    None => {
234                        return Err(ParseError::Plan(
235                            rel_node.context(),
236                            MessageParseError::invalid(
237                                "relation",
238                                rel_node.pair.as_span(),
239                                format!("No parent found for depth {depth}"),
240                            ),
241                        ));
242                    }
243                    Some(parent) => parent,
244                };
245
246                parent.children.push(rel_node);
247            }
248            LineNode::Addendum(addendum) => {
249                let context =
250                    ParseContext::new(addendum.line_no, addendum.pair.as_str().to_string());
251                if depth == 0 {
252                    return Err(ParseError::ValidationError(
253                        context,
254                        "addenda (+ Enh: / + Opt: / + Ext:) cannot appear at the top level"
255                            .to_string(),
256                    ));
257                }
258
259                let parent = match self.get_at_depth(depth - 1) {
260                    None => {
261                        return Err(ParseError::ValidationError(
262                            context,
263                            format!("no parent found for addendum at depth {depth}"),
264                        ));
265                    }
266                    Some(parent) => parent,
267                };
268
269                if !parent.children.is_empty() {
270                    return Err(ParseError::ValidationError(
271                        context,
272                        "addenda (+ Enh: / + Opt: / + Ext:) must appear before child relations, \
273                         not after"
274                            .to_string(),
275                    ));
276                }
277
278                parent.addenda.push(addendum);
279            }
280        }
281        Ok(())
282    }
283
284    /// End of input - move any remaining nodes from stack to completed and
285    /// return any trees in progress. Resets the builder to its initial state
286    /// (empty)
287    /// Move any remaining nodes from stack to completed
288    pub fn finish(&mut self) -> Vec<RelationNode<'a>> {
289        if let Some(node) = self.current.take() {
290            self.completed.push(node);
291        }
292        std::mem::take(&mut self.completed)
293    }
294}
295
296/// A child relation that has already been converted to a
297///  proto `Rel`, paired with its own output field count.
298struct ParsedChild {
299    rel: Rel,
300    field_count: usize,
301}
302
303/// Intermediate state for relation parsing: the structural tree data
304/// (children, addenda) has been parsed, but the relation's own grammar pair
305/// hasn't been converted to a protobuf relation yet.
306struct RelationContext<'a> {
307    pair: Pair<'a, Rule>,
308    line_no: i64,
309    children: Vec<ParsedChild>,
310    addenda: Addenda<'a>,
311}
312
313impl<'a> RelationContext<'a> {
314    /// Total input field count across all children.
315    fn input_field_count(&self) -> usize {
316        self.children.iter().map(|c| c.field_count).sum()
317    }
318
319    /// Per-child field counts, in order. Needed only by `Set`, to validate
320    /// all inputs share the same width
321    fn child_field_counts(&self) -> impl Iterator<Item = usize> + '_ {
322        self.children.iter().map(|c| c.field_count)
323    }
324
325    /// `ParseContext` for error reporting
326    fn as_parse_context(&self) -> ParseContext {
327        ParseContext::new(self.line_no, self.pair.as_str().to_string())
328    }
329}
330
331/// Discards field counts, keeping only the child relations.
332fn into_rels(children: Vec<ParsedChild>) -> Vec<Rel> {
333    children.into_iter().map(|c| c.rel).collect()
334}
335
336/// A parsed addendum line plus enough source location to build later errors.
337#[derive(Debug, Clone)]
338struct ParsedAddendum<'a> {
339    line_no: i64,
340    line: &'a str,
341    invocation: AddendumInvocation,
342}
343
344impl<'a> ParsedAddendum<'a> {
345    fn parse(extensions: &SimpleExtensions, addendum: Addendum<'a>) -> Result<Self, ParseError> {
346        let line_no = addendum.line_no;
347        let line = addendum.pair.as_str();
348        let invocation = AddendumInvocation::parse_pair(extensions, addendum.pair)
349            .map_err(|e| ParseError::Plan(ParseContext::new(line_no, line.to_string()), e))?;
350        Ok(Self {
351            line_no,
352            line,
353            invocation,
354        })
355    }
356
357    fn context(&self) -> ParseContext {
358        ParseContext::new(self.line_no, self.line.to_string())
359    }
360
361    fn relation_context<'b>(
362        &'b self,
363        registry: &'b ExtensionRegistry,
364    ) -> RelationParsingContext<'b> {
365        RelationParsingContext {
366            registry,
367            line_no: self.line_no,
368            line: self.line,
369        }
370    }
371    fn resolve_detail(&self, registry: &ExtensionRegistry) -> Result<Any, ParseError> {
372        self.relation_context(registry).resolve_addendum_detail(
373            self.invocation.kind,
374            &self.invocation.name,
375            &self.invocation.args,
376        )
377    }
378}
379
380/// Parsed `+` lines attached to a relation.
381#[derive(Debug, Clone, Default)]
382struct Addenda<'a> {
383    items: Vec<ParsedAddendum<'a>>,
384}
385
386impl<'a> Addenda<'a> {
387    fn parse(
388        extensions: &SimpleExtensions,
389        addenda: Vec<Addendum<'a>>,
390    ) -> Result<Self, ParseError> {
391        let items = addenda
392            .into_iter()
393            .map(|addendum| ParsedAddendum::parse(extensions, addendum))
394            .collect::<Result<Vec<_>, ParseError>>()?;
395        Ok(Self { items })
396    }
397
398    fn first(&self) -> Option<&ParsedAddendum<'a>> {
399        self.items.first()
400    }
401
402    fn reject_all(&self, message: &'static str) -> Result<(), ParseError> {
403        if let Some(addendum) = self.first() {
404            return Err(ParseError::ValidationError(
405                addendum.context(),
406                message.to_string(),
407            ));
408        }
409        Ok(())
410    }
411
412    fn into_standard_advanced_extension(
413        self,
414        registry: &ExtensionRegistry,
415    ) -> Result<Option<AdvancedExtension>, ParseError> {
416        let mut enhancement = None;
417        let mut optimizations = Vec::new();
418
419        for addendum in self.items {
420            match addendum.invocation.kind {
421                AddendumKind::Enhancement => {
422                    if enhancement.is_some() {
423                        return Err(ParseError::ValidationError(
424                            addendum.context(),
425                            "at most one enhancement per relation is allowed".to_string(),
426                        ));
427                    }
428                    enhancement = Some(addendum.resolve_detail(registry)?.into());
429                }
430                AddendumKind::Optimization => {
431                    optimizations.push(addendum.resolve_detail(registry)?.into());
432                }
433                AddendumKind::ExtensionTable => {
434                    return Err(ParseError::ValidationError(
435                        addendum.context(),
436                        "+ Ext addenda can only be used with Read:Extension".to_string(),
437                    ));
438                }
439            }
440        }
441
442        if enhancement.is_none() && optimizations.is_empty() {
443            return Ok(None);
444        }
445
446        Ok(Some(AdvancedExtension {
447            enhancement,
448            optimization: optimizations,
449        }))
450    }
451
452    fn into_extension_read_parts(
453        self,
454        registry: &ExtensionRegistry,
455        relation_context: ParseContext,
456    ) -> Result<(Any, Option<AdvancedExtension>), ParseError> {
457        let mut extension_table = None;
458        let mut advanced_addenda = Vec::new();
459
460        for addendum in self.items {
461            match addendum.invocation.kind {
462                AddendumKind::ExtensionTable => {
463                    if extension_table.is_some() {
464                        return Err(ParseError::ValidationError(
465                            addendum.context(),
466                            "Read:Extension allows exactly one + Ext addendum".to_string(),
467                        ));
468                    }
469                    extension_table = Some(addendum);
470                }
471                AddendumKind::Enhancement | AddendumKind::Optimization => {
472                    advanced_addenda.push(addendum);
473                }
474            }
475        }
476
477        let extension_table = extension_table.ok_or_else(|| {
478            ParseError::ValidationError(
479                relation_context,
480                "Read:Extension requires exactly one + Ext addendum".to_string(),
481            )
482        })?;
483
484        let detail = extension_table.resolve_detail(registry)?;
485        let advanced_extension = Addenda {
486            items: advanced_addenda,
487        }
488        .into_standard_advanced_extension(registry)?;
489
490        Ok((detail, advanced_extension))
491    }
492}
493
494// Relation parsing component - handles converting LineNodes to Relations
495#[derive(Debug, Clone, Default)]
496pub struct RelationParser<'a> {
497    tree: TreeBuilder<'a>,
498}
499
500impl<'a> RelationParser<'a> {
501    /// Dispatch by grammar rule after validating addenda constraints.
502    /// Standard relations go through [`parse_rel`](Self::parse_rel);
503    /// extension relations go through
504    /// [`parse_extension_relation`](Self::parse_extension_relation).
505    fn parse_relation(
506        &self,
507        extensions: &SimpleExtensions,
508        registry: &ExtensionRegistry,
509        ctx: RelationContext,
510    ) -> Result<(Rel, usize), ParseError> {
511        match ctx.pair.as_rule() {
512            Rule::extension_read_relation => {
513                self.parse_extension_read_relation(extensions, registry, ctx)
514            }
515            Rule::virtual_read_relation => {
516                self.parse_rel::<VirtualReadRel>(extensions, registry, ctx)
517            }
518            Rule::read_relation => self.parse_rel::<ReadRel>(extensions, registry, ctx),
519            Rule::filter_relation => self.parse_rel::<FilterRel>(extensions, registry, ctx),
520            Rule::project_relation => self.parse_rel::<ProjectRel>(extensions, registry, ctx),
521            Rule::aggregate_relation => self.parse_rel::<AggregateRel>(extensions, registry, ctx),
522            Rule::sort_relation => self.parse_rel::<SortRel>(extensions, registry, ctx),
523            Rule::fetch_relation => self.parse_rel::<FetchRel>(extensions, registry, ctx),
524            Rule::join_relation => self.parse_rel::<JoinRel>(extensions, registry, ctx),
525            Rule::set_relation => self.parse_set_relation(registry, ctx),
526            Rule::cross_relation => self.parse_rel::<CrossRel>(extensions, registry, ctx),
527            Rule::extension_relation => self.parse_extension_relation(extensions, registry, ctx),
528            _ => unreachable!("unhandled relation rule: {:?}", ctx.pair.as_rule()),
529        }
530    }
531
532    /// Generic bridge between [`parse_relation`](Self::parse_relation) and
533    /// the [`RelationParsePair`] trait: wraps `MessageParseError` with line
534    /// context and calls [`into_rel`](RelationParsePair::into_rel) to apply
535    /// addenda and produce the final [`Rel`].
536    fn parse_rel<T: RelationParsePair>(
537        &self,
538        extensions: &SimpleExtensions,
539        registry: &ExtensionRegistry,
540        ctx: RelationContext,
541    ) -> Result<(Rel, usize), ParseError> {
542        let RelationContext {
543            pair,
544            line_no,
545            children,
546            addenda,
547        } = ctx;
548        assert_eq!(pair.as_rule(), T::rule());
549        let line = pair.as_str();
550        let advanced_extension = addenda.into_standard_advanced_extension(registry)?;
551        let input_field_count: usize = children.iter().map(|c| c.field_count).sum();
552        let input_children = into_rels(children);
553
554        // TODO: pass `ParsedChildren` into `parse_pair_with_context`, instead
555        // of this overall 'input_field_count'; for some relations (e.g. a LEFT
556        // SEMI JOIN), the field count sum is not the correct count.
557        match T::parse_pair_with_context(extensions, pair, input_children, input_field_count) {
558            Ok((parsed, count)) => Ok((parsed.into_rel(advanced_extension), count)),
559            Err(e) => Err(ParseError::Plan(
560                ParseContext::new(line_no, line.to_string()),
561                e,
562            )),
563        }
564    }
565
566    /// Handle `Set` relations separately from [`parse_rel`](Self::parse_rel)
567    /// because validating that every input shares the same output width
568    /// needs each child's field count individually, not their sum.
569    fn parse_set_relation(
570        &self,
571        registry: &ExtensionRegistry,
572        ctx: RelationContext,
573    ) -> Result<(Rel, usize), ParseError> {
574        assert_eq!(ctx.pair.as_rule(), Rule::set_relation);
575        let context = ctx.as_parse_context();
576        let child_field_counts: Vec<usize> = ctx.child_field_counts().collect();
577        let RelationContext {
578            pair,
579            children,
580            addenda,
581            ..
582        } = ctx;
583        let advanced_extension = addenda.into_standard_advanced_extension(registry)?;
584        let input_children = into_rels(children);
585
586        parse_set_relation_pair(
587            pair,
588            input_children,
589            &child_field_counts,
590            advanced_extension,
591        )
592        .map_err(|e| ParseError::Plan(context, e))
593    }
594
595    /// Handle extension relations separately from [`parse_rel`](Self::parse_rel)
596    /// because they need registry lookups that [`RelationParsePair`] doesn't
597    /// support.
598    fn parse_extension_relation(
599        &self,
600        extensions: &SimpleExtensions,
601        registry: &ExtensionRegistry,
602        ctx: RelationContext,
603    ) -> Result<(Rel, usize), ParseError> {
604        assert_eq!(ctx.pair.as_rule(), Rule::extension_relation);
605        let line_no = ctx.line_no;
606        let line = ctx.pair.as_str().to_string();
607        let pair_span = ctx.pair.as_span();
608
609        ctx.addenda
610            .reject_all("extension relations do not support addenda (+ Enh / + Opt / + Ext)")?;
611
612        let ExtensionInvocation {
613            relation_kind,
614            name,
615            args: extension_args,
616        } = ExtensionInvocation::parse_pair(extensions, ctx.pair.clone())
617            .map_err(|e| ParseError::Plan(ParseContext::new(line_no, line.clone()), e))?;
618
619        let child_count = ctx.children.len();
620        relation_kind
621            .validate_child_count(child_count)
622            .map_err(|e| {
623                ParseError::Plan(
624                    ParseContext::new(line_no, line.to_string()),
625                    MessageParseError::invalid("extension_relation", pair_span, e),
626                )
627            })?;
628
629        let context = RelationParsingContext {
630            registry,
631            line_no,
632            line: &line,
633        };
634
635        let detail = context.resolve_extension_detail(&name, &extension_args)?;
636        let output_column_count = extension_args.output_columns.len();
637
638        let rel = relation_kind.create_rel(detail, into_rels(ctx.children));
639
640        Ok((rel, output_column_count))
641    }
642
643    /// Parse `Read:Extension[...]`, whose table detail is supplied by exactly
644    /// one `+ Ext:Name[...]` addendum.
645    fn parse_extension_read_relation(
646        &self,
647        extensions: &SimpleExtensions,
648        registry: &ExtensionRegistry,
649        ctx: RelationContext,
650    ) -> Result<(Rel, usize), ParseError> {
651        assert_eq!(ctx.pair.as_rule(), Rule::extension_read_relation);
652        let context = ParseContext::new(ctx.line_no, ctx.pair.as_str().to_string());
653        let input_field_count = ctx.input_field_count();
654        let (detail, advanced_extension) = ctx
655            .addenda
656            .into_extension_read_parts(registry, context.clone())?;
657
658        ExtensionReadRel::parse_pair_with_detail(
659            extensions,
660            ctx.pair,
661            into_rels(ctx.children),
662            input_field_count,
663            detail,
664            advanced_extension,
665        )
666        .map_err(|e| ParseError::Plan(context, e))
667    }
668
669    /// Walk the relation tree depth-first, converting structural types
670    /// (children, addenda) into proto types via [`RelationContext`].
671    /// Delegates grammar-rule-specific work to
672    /// [`parse_relation`](Self::parse_relation).
673    fn build_rel(
674        &self,
675        extensions: &SimpleExtensions,
676        registry: &ExtensionRegistry,
677        node: RelationNode,
678    ) -> Result<(Rel, usize), ParseError> {
679        let mut children: Vec<ParsedChild> = Vec::new();
680        for child in node.children {
681            let (rel, field_count) = self.build_rel(extensions, registry, child)?;
682            children.push(ParsedChild { rel, field_count });
683        }
684
685        let addenda = Addenda::parse(extensions, node.addenda)?;
686
687        self.parse_relation(
688            extensions,
689            registry,
690            RelationContext {
691                pair: node.pair,
692                line_no: node.line_no,
693                children,
694                addenda,
695            },
696        )
697    }
698
699    /// Build a tree of relations.
700    fn build_plan_rel(
701        &self,
702        extensions: &SimpleExtensions,
703        registry: &ExtensionRegistry,
704        node: RelationNode,
705    ) -> Result<PlanRel, ParseError> {
706        // Plain relations are allowed as root relations; they just don't have names.
707        if node.pair.as_rule() != Rule::root_relation {
708            let (rel, _) = self.build_rel(extensions, registry, node)?;
709            return Ok(PlanRel {
710                rel_type: Some(plan_rel::RelType::Rel(rel)),
711            });
712        }
713
714        // Root relations don't support addenda — reject rather than silently discard.
715        if !node.addenda.is_empty() {
716            let first = &node.addenda[0];
717            let context = ParseContext::new(first.line_no, first.pair.as_str().to_string());
718            return Err(ParseError::ValidationError(
719                context,
720                "addenda (+ Enh: / + Opt: / + Ext:) are not supported on Root relations"
721                    .to_string(),
722            ));
723        }
724
725        // Named root relation.
726        let context = node.context();
727        let span = node.pair.as_span();
728
729        // Parse the column names
730        let column_names_pair = unwrap_single_pair(node.pair);
731        assert_eq!(column_names_pair.as_rule(), Rule::root_name_list);
732
733        let names: Vec<String> = column_names_pair
734            .into_inner()
735            .map(|name_pair| {
736                assert_eq!(name_pair.as_rule(), Rule::name);
737                Name::parse_pair(name_pair).0
738            })
739            .collect();
740
741        let mut children = node.children;
742        let child = match children.len() {
743            1 => {
744                let (rel, _) = self.build_rel(extensions, registry, children.pop().unwrap())?;
745                rel
746            }
747            n => {
748                return Err(ParseError::Plan(
749                    context,
750                    MessageParseError::invalid(
751                        "root_relation",
752                        span,
753                        format!("Root relation must have exactly one child, found {n}"),
754                    ),
755                ));
756            }
757        };
758
759        Ok(PlanRel {
760            rel_type: Some(plan_rel::RelType::Root(RelRoot {
761                names,
762                input: Some(child),
763            })),
764        })
765    }
766
767    /// Build all the trees.
768    fn build(
769        mut self,
770        extensions: &SimpleExtensions,
771        registry: &ExtensionRegistry,
772    ) -> Result<Vec<PlanRel>, ParseError> {
773        let nodes = self.tree.finish();
774        nodes
775            .into_iter()
776            .map(|n| self.build_plan_rel(extensions, registry, n))
777            .collect::<Result<Vec<PlanRel>, ParseError>>()
778    }
779}
780
781/// A parser for Substrait query plans in text format.
782///
783/// The `Parser` converts human-readable Substrait text format into Substrait
784/// protobuf plans. It handles both the extensions section (which defines
785/// functions, types, etc.) and the plan section (which defines the actual query
786/// structure).
787///
788/// ## Usage
789///
790/// The simplest entry point is the static `parse()` method:
791///
792/// ```rust
793/// use substrait_explain::Parser;
794///
795/// let plan_text = r#"
796/// === Plan
797/// Root[c, d]
798///   Project[$1, 42]
799///     Read[schema.table => a:i64, b:string?]
800/// "#;
801///
802/// let plan = Parser::parse(plan_text).unwrap();
803/// ```
804///
805/// ## Input Format
806///
807/// The parser expects input in the following format:
808///
809/// ```text
810/// === Extensions
811/// URNs:
812///   @  1: https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml
813/// Functions:
814///   # 10 @  1: add
815/// === Plan
816/// Root[columns]
817///   Relation[arguments => columns]
818///     ChildRelation[arguments => columns]
819/// ```
820///
821/// - **Extensions section** (optional): Defines URNs and function/type declarations
822/// - **Plan section** (required): Defines the query structure with indented relations
823///
824/// ## Error Handling
825///
826/// The parser provides detailed error information including:
827/// - Line number where the error occurred
828/// - The actual line content that failed to parse
829/// - Specific error type and description
830///
831/// ```rust
832/// use substrait_explain::Parser;
833///
834/// let invalid_plan = r#"
835/// === Plan
836/// InvalidRelation[invalid syntax]
837/// "#;
838///
839/// match Parser::parse(invalid_plan) {
840///     Ok(plan) => println!("Successfully parsed"),
841///     Err(e) => eprintln!("Parse error: {}", e),
842/// }
843/// ```
844///
845/// ## Supported Relations
846///
847/// The parser supports all standard Substrait relations:
848/// - `Read[table => columns]` - Read from a table
849/// - `Project[expressions]` - Project columns/expressions
850/// - `Filter[condition => columns]` - Filter rows
851/// - `Root[columns]` - Root relation with output columns
852/// - And more...
853///
854/// ## Extensions Support
855///
856/// The parser fully supports Substrait Simple Extensions, allowing you to:
857/// - Define custom functions with URNs and anchors
858/// - Reference functions by name in expressions
859/// - Use custom types and type variations
860///
861/// ```rust
862/// use substrait_explain::Parser;
863///
864/// let plan_with_extensions = r#"
865/// === Extensions
866/// URNs:
867///   @  1: https://example.com/functions.yaml
868/// Functions:
869///   ## 10 @  1: my_custom_function
870/// === Plan
871/// Root[result]
872///   Project[my_custom_function($0, $1):i32]
873///     Read[table => col1:i32, col2:i32]
874/// "#;
875///
876/// let plan = Parser::parse(plan_with_extensions).unwrap();
877/// ```
878///
879/// ## Performance
880///
881/// The parser is designed for efficiency:
882/// - Single-pass parsing with minimal allocations
883/// - Early error detection and reporting
884/// - Memory-efficient tree building
885///
886/// ## Thread Safety
887///
888/// `Parser` instances are not thread-safe and should not be shared between threads.
889/// However, the static `parse()` method is safe to call from multiple threads.
890#[derive(Debug)]
891pub struct Parser<'a> {
892    line_no: i64,
893    state: State,
894    /// Cursor over the remaining input, advanced one chunk at a time by
895    /// [`next_chunk`](Self::next_chunk). `None` before parsing starts and once
896    /// the input is exhausted.
897    cursor: Option<ChunkCursor<'a>>,
898    /// The plan version from an optional `=== Version` section
899    version: Option<Version>,
900    extension_parser: ExtensionParser,
901    extension_registry: ExtensionRegistry,
902    relation_parser: RelationParser<'a>,
903}
904impl<'a> Default for Parser<'a> {
905    fn default() -> Self {
906        Self::new()
907    }
908}
909
910impl<'a> Parser<'a> {
911    /// Parse a Substrait plan from text format.
912    ///
913    /// This is the main entry point for parsing.
914    ///
915    /// The input should be in the Substrait text format, which consists of:
916    /// - An optional extensions section starting with "=== Extensions"
917    /// - A plan section starting with "=== Plan"
918    /// - Indented relation definitions
919    ///
920    /// # Examples
921    ///
922    /// Simple parsing:
923    /// ```rust
924    /// use substrait_explain::Parser;
925    ///
926    /// let plan_text = r#"
927    /// === Plan
928    /// Root[result]
929    ///   Read[table => col:i32]
930    /// "#;
931    ///
932    /// let plan = Parser::parse(plan_text).unwrap();
933    /// assert_eq!(plan.relations.len(), 1);
934    /// ```
935    ///
936    /// # Errors
937    ///
938    /// Returns a [`ParseError`] if the input cannot be parsed.
939    pub fn parse(input: &str) -> ParseResult {
940        Self::new().parse_plan(input)
941    }
942
943    /// Create a new parser with default configuration.
944    pub fn new() -> Self {
945        Self {
946            line_no: 1,
947            state: State::Initial,
948            cursor: None,
949            version: None,
950            extension_parser: ExtensionParser::default(),
951            extension_registry: ExtensionRegistry::new(),
952            relation_parser: RelationParser::default(),
953        }
954    }
955
956    /// Configure the parser to use the specified extension registry.
957    pub fn with_extension_registry(mut self, registry: ExtensionRegistry) -> Self {
958        self.extension_registry = registry;
959        self
960    }
961
962    /// Parse a Substrait plan with the current parser configuration.
963    pub fn parse_plan(mut self, input: &'a str) -> ParseResult {
964        self.cursor = ChunkCursor::new(input, 1);
965        while self.cursor.is_some() {
966            let (chunk, line_no) = self.next_chunk();
967
968            if chunk.trim().is_empty() {
969                continue;
970            }
971
972            self.line_no = line_no;
973            // TODO: multi-line chunk errors report a confusing location. The
974            // outer message uses the chunk's first line, while the Pest span
975            // inside counts lines relative to the chunk, so neither points at
976            // the true absolute line. E.g. a bad row on the 2nd line of a chunk
977            // starting at line 3 produces:
978            //   Error parsing plan on line 3: '...': ...
979            //    --> 2:8
980            // We should map the Pest-relative line back to an absolute line.
981            self.parse_line(chunk)?;
982        }
983
984        let plan = self.build_plan()?;
985        Ok(plan)
986    }
987
988    /// Group the next chunk out of the cursor: always its first physical line,
989    /// plus any `- ` continuation lines indented one level deeper while we are
990    /// in the plan section.
991    ///
992    /// For consistency, chunks always end with no newline.
993    fn next_chunk(&mut self) -> (&'a str, i64) {
994        let mut c = self
995            .cursor
996            .take()
997            .expect("next_chunk called with no cursor");
998
999        // Every chunk contains at least its first physical line.
1000        let first = c
1001            .peek_line()
1002            .expect("a non-exhausted cursor always yields a line");
1003        c.merge(first);
1004
1005        if matches!(self.state, State::Plan) {
1006            let base = IndentedLine::from(first.as_str()).0;
1007            while let Some(line) = c.peek_line() {
1008                let IndentedLine(depth, body) = IndentedLine::from(line.as_str());
1009                if depth == base + 1 && body.starts_with("- ") {
1010                    c.merge(line);
1011                } else {
1012                    break;
1013                }
1014            }
1015        }
1016
1017        let line_no = c.start_line_no();
1018        let (chunk, rest) = c.next();
1019        self.cursor = rest;
1020        (chunk.trim_end_matches(['\r', '\n']), line_no)
1021    }
1022
1023    /// Parse a single line of input.
1024    fn parse_line(&mut self, line: &'a str) -> Result<(), ParseError> {
1025        let indented_line = IndentedLine::from(line);
1026        let line_no = self.line_no;
1027        let ctx = || ParseContext {
1028            line_no,
1029            line: line.to_string(),
1030        };
1031
1032        // A `===`-prefixed line at indent 0 is always a section header: it
1033        // transitions us into the next section. Sections must appear in
1034        // strictly increasing order (`Version`, `Extensions`, `Plan`,
1035        // mirroring the Substrait `Plan` protobuf field order), so a header
1036        // cannot repeat a section or move backwards.
1037        if let IndentedLine(0, l) = indented_line
1038            && l.starts_with("===")
1039        {
1040            let next = self.parse_section_header(l)?;
1041            if next.section_rank() <= self.state.section_rank() {
1042                return Err(ParseError::ValidationError(
1043                    ctx(),
1044                    format!(
1045                        "section {next} is out of order after {current}; sections must appear \
1046                         in the order Version, Extensions, Plan",
1047                        current = self.state,
1048                    ),
1049                ));
1050            }
1051            self.state = next;
1052            return Ok(());
1053        }
1054
1055        match self.state {
1056            State::Initial => match indented_line {
1057                IndentedLine(0, l) if l.trim().is_empty() => Ok(()),
1058                IndentedLine(n, l) => Err(ParseError::Initial(
1059                    ParseContext::new(n as i64, l.to_string()),
1060                    MessageParseError::invalid(
1061                        "initial",
1062                        pest::Span::new(l, 0, l.len()).expect("Invalid span?!"),
1063                        format!("Unknown initial line: {l:?}"),
1064                    ),
1065                )),
1066            },
1067            State::Version => self.parse_version(indented_line),
1068            State::Extensions => self
1069                .parse_extensions(indented_line)
1070                .map_err(|e| ParseError::Extension(ctx(), e)),
1071            State::Plan => {
1072                let IndentedLine(depth, line_str) = indented_line;
1073
1074                // Parse the line
1075                let node = if depth == 0 {
1076                    LineNode::parse_root(line_str, line_no)?
1077                } else {
1078                    LineNode::parse(line_str, line_no)?
1079                };
1080
1081                self.relation_parser.tree.add_line(depth, node)
1082            }
1083        }
1084    }
1085
1086    /// Parse a `===`-prefixed section header line (at indent 0), returning
1087    /// the [`State`] to transition into. Also performs any header-specific
1088    /// parsing, e.g. storing the plan [`Version`] from a `=== Version`
1089    /// header.
1090    fn parse_section_header(&mut self, line: &str) -> Result<State, ParseError> {
1091        if line == simple::EXTENSIONS_HEADER {
1092            return Ok(State::Extensions);
1093        }
1094        if line == PLAN_HEADER {
1095            return Ok(State::Plan);
1096        }
1097        if let Some(rest) = line.strip_prefix(VERSION_HEADER)
1098            && (rest.is_empty() || rest.starts_with(' '))
1099        {
1100            self.version = Some(self.parse_version_header(line)?);
1101            return Ok(State::Version);
1102        }
1103
1104        Err(ParseError::Initial(
1105            ParseContext::new(self.line_no, line.to_string()),
1106            MessageParseError::invalid(
1107                "initial",
1108                pest::Span::new(line, 0, line.len()).expect("Invalid span?!"),
1109                format!("Unknown section header: {line:?}"),
1110            ),
1111        ))
1112    }
1113
1114    /// Parse the `=== Version <major>.<minor>.<patch>` header line and store
1115    /// the resulting [`Version`], leaving `producer` / `git_hash` empty
1116    fn parse_version_header(&self, line: &str) -> Result<Version, ParseError> {
1117        let ctx = || ParseContext::new(self.line_no, line.to_string());
1118        let rest = line
1119            .strip_prefix(VERSION_HEADER)
1120            .expect("version header prefix checked by caller")
1121            .trim();
1122
1123        let mut numbers = rest.split('.');
1124        let mut next_number = |field: &str| -> Result<u32, ParseError> {
1125            let part = numbers.next().filter(|p| !p.is_empty()).ok_or_else(|| {
1126                ParseError::ValidationError(
1127                    ctx(),
1128                    format!("version header is missing the {field} number, expected '{VERSION_HEADER} <major>.<minor>.<patch>'"),
1129                )
1130            })?;
1131            part.parse::<u32>().map_err(|e| {
1132                ParseError::ValidationError(
1133                    ctx(),
1134                    format!("invalid {field} version number {part:?}: {e}"),
1135                )
1136            })
1137        };
1138
1139        let major_number = next_number("major")?;
1140        let minor_number = next_number("minor")?;
1141        let patch_number = next_number("patch")?;
1142        if numbers.next().is_some() {
1143            return Err(ParseError::ValidationError(
1144                ctx(),
1145                format!(
1146                    "version {rest:?} has too many components, expected '<major>.<minor>.<patch>'"
1147                ),
1148            ));
1149        }
1150
1151        Ok(Version {
1152            major_number,
1153            minor_number,
1154            patch_number,
1155            ..Default::default()
1156        })
1157    }
1158
1159    /// Parses a single line from the version section: an indented
1160    /// `producer:` / `git_hash:` detail line. Section headers are
1161    /// intercepted by [`parse_line`](Self::parse_line) before reaching here.
1162    fn parse_version(&mut self, line: IndentedLine) -> Result<(), ParseError> {
1163        match line {
1164            IndentedLine(0, l) if l.trim().is_empty() => Ok(()),
1165            IndentedLine(1, l) => self.parse_version_detail(l),
1166            IndentedLine(_, l) => Err(ParseError::ValidationError(
1167                ParseContext::new(self.line_no, l.to_string()),
1168                format!(
1169                    "unexpected line in version section: {l:?}; expected an indented 'producer:' \
1170                     or 'git_hash:' line, or a section header"
1171                ),
1172            )),
1173        }
1174    }
1175
1176    /// Parse an indented `producer:` / `git_hash:` detail line, mutating the
1177    /// [`Version`] set by [`parse_version_header`](Self::parse_version_header).
1178    fn parse_version_detail(&mut self, line: &str) -> Result<(), ParseError> {
1179        let ctx = || ParseContext::new(self.line_no, line.to_string());
1180        let (key, value) = line.split_once(':').ok_or_else(|| {
1181            ParseError::ValidationError(
1182                ctx(),
1183                format!("expected 'producer:' or 'git_hash:' in version section, found {line:?}"),
1184            )
1185        })?;
1186        let value = value.trim().to_string();
1187        let version = self
1188            .version
1189            .as_mut()
1190            .expect("version is set on entry to the version section");
1191        match key.trim() {
1192            "producer" => version.producer = value,
1193            "git_hash" => version.git_hash = value,
1194            other => {
1195                return Err(ParseError::ValidationError(
1196                    ctx(),
1197                    format!("unknown version field {other:?}, expected 'producer' or 'git_hash'"),
1198                ));
1199            }
1200        }
1201        Ok(())
1202    }
1203
1204    /// Parse a single line from the extensions section of the input.
1205    /// Section headers are intercepted by [`parse_line`](Self::parse_line)
1206    /// before reaching here.
1207    fn parse_extensions(&mut self, line: IndentedLine<'_>) -> Result<(), ExtensionParseError> {
1208        self.extension_parser.parse_line(line)
1209    }
1210
1211    /// Build the plan from the parser state with warning collection.
1212    fn build_plan(self) -> Result<Plan, ParseError> {
1213        let Parser {
1214            version,
1215            relation_parser,
1216            extension_parser,
1217            extension_registry,
1218            ..
1219        } = self;
1220
1221        let extensions = extension_parser.extensions();
1222
1223        // Parse the tree into relations
1224        let root_relations = relation_parser.build(extensions, &extension_registry)?;
1225
1226        // Build the final plan
1227        Ok(Plan {
1228            version,
1229            extension_urns: extensions.to_extension_urns(),
1230            extensions: extensions.to_extension_declarations(),
1231            relations: root_relations,
1232            ..Default::default()
1233        })
1234    }
1235}
1236
1237#[cfg(test)]
1238mod tests {
1239    use substrait::proto::extensions::simple_extension_declaration::MappingType;
1240    use substrait::proto::rel::RelType;
1241
1242    use super::*;
1243    use crate::extensions::simple::ExtensionKind;
1244    use crate::parser::extensions::ExpectedExtensionLine;
1245
1246    #[test]
1247    fn test_parse_basic_block() {
1248        let mut expected_extensions = SimpleExtensions::new();
1249        expected_extensions
1250            .add_extension_urn("/urn/common".to_string(), 1)
1251            .unwrap();
1252        expected_extensions
1253            .add_extension_urn("/urn/specific_funcs".to_string(), 2)
1254            .unwrap();
1255        expected_extensions
1256            .add_extension(ExtensionKind::Function, 1, 10, "func_a".to_string())
1257            .unwrap();
1258        expected_extensions
1259            .add_extension(ExtensionKind::Function, 2, 11, "func_b_special".to_string())
1260            .unwrap();
1261        expected_extensions
1262            .add_extension(ExtensionKind::Type, 1, 20, "SomeType".to_string())
1263            .unwrap();
1264        expected_extensions
1265            .add_extension(ExtensionKind::TypeVariation, 2, 30, "VarX".to_string())
1266            .unwrap();
1267
1268        let mut parser = ExtensionParser::default();
1269        let input_block = r#"
1270URNs:
1271  @  1: /urn/common
1272  @  2: /urn/specific_funcs
1273Functions:
1274  # 10 @  1: func_a
1275  # 11 @  2: func_b_special
1276Types:
1277  # 20 @  1: SomeType
1278Type Variations:
1279  # 30 @  2: VarX
1280"#;
1281
1282        for line_str in input_block.trim().lines() {
1283            parser
1284                .parse_line(IndentedLine::from(line_str))
1285                .unwrap_or_else(|e| panic!("Failed to parse line \'{line_str}\': {e:?}"));
1286        }
1287
1288        assert_eq!(*parser.extensions(), expected_extensions);
1289
1290        let extensions_str = parser.extensions().to_string("  ");
1291        // The writer adds the header; the ExtensionParser does not parse the
1292        // header, so we add it here for comparison.
1293        let expected_str = format!(
1294            "{}\n{}",
1295            simple::EXTENSIONS_HEADER,
1296            input_block.trim_start()
1297        );
1298        assert_eq!(extensions_str.trim(), expected_str.trim());
1299        // Check final state after all lines are processed.
1300        // The last significant line in input_block is a TypeVariation declaration.
1301        assert_eq!(
1302            parser.state(),
1303            ExpectedExtensionLine::ExtensionDeclarations(ExtensionKind::TypeVariation)
1304        );
1305
1306        // Check that a subsequent blank line correctly resets state to Extensions.
1307        parser.parse_line(IndentedLine(0, "")).unwrap();
1308        assert_eq!(parser.state(), ExpectedExtensionLine::Extensions);
1309    }
1310
1311    /// Test that we can parse a larger extensions block and it matches the input.
1312    #[test]
1313    fn test_parse_complete_extension_block() {
1314        let mut parser = ExtensionParser::default();
1315        let input_block = r#"
1316URNs:
1317  @  1: /urn/common
1318  @  2: /urn/specific_funcs
1319  @  3: /urn/types_lib
1320  @  4: /urn/variations_lib
1321Functions:
1322  # 10 @  1: func_a
1323  # 11 @  2: func_b_special
1324  # 12 @  1: func_c_common
1325Types:
1326  # 20 @  1: CommonType
1327  # 21 @  3: LibraryType
1328  # 22 @  1: AnotherCommonType
1329Type Variations:
1330  # 30 @  4: VarX
1331  # 31 @  4: VarY
1332"#;
1333
1334        for line_str in input_block.trim().lines() {
1335            parser
1336                .parse_line(IndentedLine::from(line_str))
1337                .unwrap_or_else(|e| panic!("Failed to parse line \'{line_str}\': {e:?}"));
1338        }
1339
1340        let extensions_str = parser.extensions().to_string("  ");
1341        // The writer adds the header; the ExtensionParser does not parse the
1342        // header, so we add it here for comparison.
1343        let expected_str = format!(
1344            "{}\n{}",
1345            simple::EXTENSIONS_HEADER,
1346            input_block.trim_start()
1347        );
1348        assert_eq!(extensions_str.trim(), expected_str.trim());
1349    }
1350
1351    #[test]
1352    fn test_parse_relation_tree() {
1353        // Example plan with a Project, a Filter, and a Read, nested by indentation
1354        let plan = r#"=== Plan
1355Project[$0, $1, 42, 84]
1356  Filter[$2 => $0, $1]
1357    Read[my.table => a:i32, b:string?, c:boolean]
1358"#;
1359        let mut parser = Parser::default();
1360        for line in plan.lines() {
1361            parser.parse_line(line).unwrap();
1362        }
1363
1364        // Complete the current tree to convert it to relations
1365        let plan = parser.build_plan().unwrap();
1366
1367        let root_rel = &plan.relations[0].rel_type;
1368        let first_rel = match root_rel {
1369            Some(plan_rel::RelType::Rel(rel)) => rel,
1370            _ => panic!("Expected Rel type, got {root_rel:?}"),
1371        };
1372        // Root should be Project
1373        let project = match &first_rel.rel_type {
1374            Some(RelType::Project(p)) => p,
1375            other => panic!("Expected Project at root, got {other:?}"),
1376        };
1377
1378        // Check that Project has Filter as input
1379        assert!(project.input.is_some());
1380        let filter_input = project.input.as_ref().unwrap();
1381
1382        // Check that Filter has Read as input
1383        match &filter_input.rel_type {
1384            Some(RelType::Filter(_)) => {
1385                match &filter_input.rel_type {
1386                    Some(RelType::Filter(filter)) => {
1387                        assert!(filter.input.is_some());
1388                        let read_input = filter.input.as_ref().unwrap();
1389
1390                        // Check that Read has no input (it's a leaf)
1391                        match &read_input.rel_type {
1392                            Some(RelType::Read(_)) => {}
1393                            other => panic!("Expected Read relation, got {other:?}"),
1394                        }
1395                    }
1396                    other => panic!("Expected Filter relation, got {other:?}"),
1397                }
1398            }
1399            other => panic!("Expected Filter relation, got {other:?}"),
1400        }
1401    }
1402
1403    #[test]
1404    fn test_parse_root_relation() {
1405        // Test a plan with a Root relation
1406        let plan = r#"=== Plan
1407Root[result]
1408  Project[$0, $1]
1409    Read[my.table => a:i32, b:string?]
1410"#;
1411        let mut parser = Parser::default();
1412        for line in plan.lines() {
1413            parser.parse_line(line).unwrap();
1414        }
1415
1416        let plan = parser.build_plan().unwrap();
1417
1418        // Check that we have exactly one relation
1419        assert_eq!(plan.relations.len(), 1);
1420
1421        let root_rel = &plan.relations[0].rel_type;
1422        let rel_root = match root_rel {
1423            Some(plan_rel::RelType::Root(rel_root)) => rel_root,
1424            other => panic!("Expected Root type, got {other:?}"),
1425        };
1426
1427        // Check that the root has the correct name
1428        assert_eq!(rel_root.names, vec!["result"]);
1429
1430        // Check that the root has a Project as input
1431        let project_input = match &rel_root.input {
1432            Some(rel) => rel,
1433            None => panic!("Root should have an input"),
1434        };
1435
1436        let project = match &project_input.rel_type {
1437            Some(RelType::Project(p)) => p,
1438            other => panic!("Expected Project as root input, got {other:?}"),
1439        };
1440
1441        // Check that Project has Read as input
1442        let read_input = match &project.input {
1443            Some(rel) => rel,
1444            None => panic!("Project should have an input"),
1445        };
1446
1447        match &read_input.rel_type {
1448            Some(RelType::Read(_)) => {}
1449            other => panic!("Expected Read relation, got {other:?}"),
1450        }
1451    }
1452
1453    #[test]
1454    fn test_parse_root_relation_no_names() {
1455        // Test a plan with a Root relation with no names
1456        let plan = r#"=== Plan
1457Root[]
1458  Project[$0, $1]
1459    Read[my.table => a:i32, b:string?]
1460"#;
1461        let mut parser = Parser::default();
1462        for line in plan.lines() {
1463            parser.parse_line(line).unwrap();
1464        }
1465
1466        let plan = parser.build_plan().unwrap();
1467
1468        let root_rel = &plan.relations[0].rel_type;
1469        let rel_root = match root_rel {
1470            Some(plan_rel::RelType::Root(rel_root)) => rel_root,
1471            other => panic!("Expected Root type, got {other:?}"),
1472        };
1473
1474        // Check that the root has no names
1475        assert_eq!(rel_root.names, Vec::<String>::new());
1476    }
1477
1478    #[test]
1479    fn test_parse_full_plan() {
1480        // Test a complete Substrait plan with extensions and relations
1481        let input = r#"
1482=== Extensions
1483URNs:
1484  @  1: /urn/common
1485  @  2: /urn/specific_funcs
1486Functions:
1487  # 10 @  1: func_a
1488  # 11 @  2: func_b_special
1489Types:
1490  # 20 @  1: SomeType
1491Type Variations:
1492  # 30 @  2: VarX
1493
1494=== Plan
1495Project[$0, $1, 42, 84]
1496  Filter[$2 => $0, $1]
1497    Read[my.table => a:i32, b:string?, c:boolean]
1498"#;
1499
1500        let plan = Parser::parse(input).unwrap();
1501
1502        // Verify the plan structure
1503        assert_eq!(plan.extension_urns.len(), 2);
1504        assert_eq!(plan.extensions.len(), 4);
1505        assert_eq!(plan.relations.len(), 1);
1506
1507        // Verify extension URIs
1508        let urn1 = &plan.extension_urns[0];
1509        assert_eq!(urn1.extension_urn_anchor, 1);
1510        assert_eq!(urn1.urn, "/urn/common");
1511
1512        let urn2 = &plan.extension_urns[1];
1513        assert_eq!(urn2.extension_urn_anchor, 2);
1514        assert_eq!(urn2.urn, "/urn/specific_funcs");
1515
1516        // Verify extensions
1517        let func1 = &plan.extensions[0];
1518        match &func1.mapping_type {
1519            Some(MappingType::ExtensionFunction(f)) => {
1520                assert_eq!(f.function_anchor, 10);
1521                assert_eq!(f.extension_urn_reference, 1);
1522                assert_eq!(f.name, "func_a");
1523            }
1524            other => panic!("Expected ExtensionFunction, got {other:?}"),
1525        }
1526
1527        let func2 = &plan.extensions[1];
1528        match &func2.mapping_type {
1529            Some(MappingType::ExtensionFunction(f)) => {
1530                assert_eq!(f.function_anchor, 11);
1531                assert_eq!(f.extension_urn_reference, 2);
1532                assert_eq!(f.name, "func_b_special");
1533            }
1534            other => panic!("Expected ExtensionFunction, got {other:?}"),
1535        }
1536
1537        let type1 = &plan.extensions[2];
1538        match &type1.mapping_type {
1539            Some(MappingType::ExtensionType(t)) => {
1540                assert_eq!(t.type_anchor, 20);
1541                assert_eq!(t.extension_urn_reference, 1);
1542                assert_eq!(t.name, "SomeType");
1543            }
1544            other => panic!("Expected ExtensionType, got {other:?}"),
1545        }
1546
1547        let var1 = &plan.extensions[3];
1548        match &var1.mapping_type {
1549            Some(MappingType::ExtensionTypeVariation(v)) => {
1550                assert_eq!(v.type_variation_anchor, 30);
1551                assert_eq!(v.extension_urn_reference, 2);
1552                assert_eq!(v.name, "VarX");
1553            }
1554            other => panic!("Expected ExtensionTypeVariation, got {other:?}"),
1555        }
1556
1557        // Verify the relation tree structure
1558        let root_rel = &plan.relations[0];
1559        match &root_rel.rel_type {
1560            Some(plan_rel::RelType::Rel(rel)) => {
1561                match &rel.rel_type {
1562                    Some(RelType::Project(project)) => {
1563                        // Verify Project relation
1564                        assert_eq!(project.expressions.len(), 2); // 42 and 84
1565                        assert!(project.input.is_some()); // Should have Filter as input
1566
1567                        // Check the Filter input
1568                        let filter_input = project.input.as_ref().unwrap();
1569                        match &filter_input.rel_type {
1570                            Some(RelType::Filter(filter)) => {
1571                                assert!(filter.input.is_some()); // Should have Read as input
1572
1573                                // Check the Read input
1574                                let read_input = filter.input.as_ref().unwrap();
1575                                match &read_input.rel_type {
1576                                    Some(RelType::Read(read)) => {
1577                                        // Verify Read relation
1578                                        let schema = read.base_schema.as_ref().unwrap();
1579                                        assert_eq!(schema.names.len(), 3);
1580                                        assert_eq!(schema.names[0], "a");
1581                                        assert_eq!(schema.names[1], "b");
1582                                        assert_eq!(schema.names[2], "c");
1583
1584                                        let struct_ = schema.r#struct.as_ref().unwrap();
1585                                        assert_eq!(struct_.types.len(), 3);
1586                                    }
1587                                    other => panic!("Expected Read relation, got {other:?}"),
1588                                }
1589                            }
1590                            other => panic!("Expected Filter relation, got {other:?}"),
1591                        }
1592                    }
1593                    other => panic!("Expected Project relation, got {other:?}"),
1594                }
1595            }
1596            other => panic!("Expected Rel type, got {other:?}"),
1597        }
1598    }
1599
1600    /// Sections must appear in the order `Version`, `Extensions`, `Plan`
1601    /// (each optional except `Plan`); a header out of that order is rejected.
1602    #[test]
1603    fn test_section_headers_out_of_order_rejected() {
1604        let input = r#"
1605=== Plan
1606Root[result]
1607  Read[t => a:i32]
1608
1609=== Extensions
1610URNs:
1611  @  1: /urn/common
1612"#;
1613        let err = Parser::parse(input).unwrap_err();
1614        let msg = err.to_string();
1615        assert!(
1616            msg.contains("out of order"),
1617            "expected an out-of-order error, got: {msg}"
1618        );
1619    }
1620
1621    #[test]
1622    fn test_repeated_section_header_rejected() {
1623        let input = r#"
1624=== Extensions
1625URNs:
1626  @  1: /urn/common
1627
1628=== Extensions
1629URNs:
1630  @  2: /urn/other
1631
1632=== Plan
1633Root[result]
1634  Read[t => a:i32]
1635"#;
1636        let err = Parser::parse(input).unwrap_err();
1637        let msg = err.to_string();
1638        assert!(
1639            msg.contains("out of order"),
1640            "expected an out-of-order error, got: {msg}"
1641        );
1642    }
1643
1644    #[test]
1645    fn test_version_after_plan_rejected() {
1646        let input = r#"
1647=== Plan
1648Root[result]
1649  Read[t => a:i32]
1650
1651=== Version 1.0.0
1652"#;
1653        let err = Parser::parse(input).unwrap_err();
1654        let msg = err.to_string();
1655        assert!(
1656            msg.contains("out of order"),
1657            "expected an out-of-order error, got: {msg}"
1658        );
1659    }
1660
1661    /// The full, valid section order (`Version`, `Extensions`, `Plan`) is
1662    /// still accepted.
1663    #[test]
1664    fn test_all_sections_in_order_accepted() {
1665        let input = r#"
1666=== Version 1.2.3
1667=== Extensions
1668URNs:
1669  @  1: /urn/common
1670
1671=== Plan
1672Root[result]
1673  Read[t => a:i32]
1674"#;
1675        let plan = Parser::parse(input).unwrap();
1676        assert_eq!(plan.relations.len(), 1);
1677        assert!(plan.version.is_some());
1678    }
1679}