substrait_explain/parser/

extensions.rs

use std::fmt;
use std::str::FromStr;

use pest::iterators::Pair;
use substrait::proto::rel::RelType;
use substrait::proto::{
    Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, Rel, Type,
};
use thiserror::Error;

use super::{
    ErrorKind, ExpressionParser, MessageParseError, ParsePair, Rule, RuleIter, ScopedParsePair,
    unescape_string, unwrap_single_pair,
};
use crate::extensions::any::Any;
use crate::extensions::simple::{self, ExtensionKind};
use crate::extensions::{
    AddendumKind, ExtensionArgs, ExtensionColumn, ExtensionValue, InsertError, SimpleExtensions,
    TupleValue,
};
use crate::parser::expressions::{FieldIndex, Name};
use crate::parser::structural::IndentedLine;
use crate::textify::expressions::Reference;

#[derive(Debug, Clone, Error)]
pub enum ExtensionParseError {
    #[error("Unexpected line, expected {0}")]
    UnexpectedLine(ExpectedExtensionLine),
    #[error("Error adding extension: {0}")]
    ExtensionError(#[from] InsertError),
    #[error("Error parsing message: {0}")]
    Message(#[from] super::MessageParseError),
}

/// The kind of extension-section line expected next.
///
/// `ExtensionParser` also uses this as its internal state, since each parser
/// state corresponds directly to the next accepted line shape.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExpectedExtensionLine {
    // The extensions section, after parsing the 'Extensions:' header, before
    // parsing any subsection headers.
    Extensions,
    // The extension URNs section, after parsing the 'URNs:' subsection header,
    // and any URNs so far.
    ExtensionUrns,
    // In a subsection, after parsing the subsection header, and any
    // declarations so far.
    ExtensionDeclarations(ExtensionKind),
}

impl fmt::Display for ExpectedExtensionLine {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ExpectedExtensionLine::Extensions => write!(f, "Subsection Header, e.g. 'URNs:'"),
            ExpectedExtensionLine::ExtensionUrns => write!(f, "Extension URNs"),
            ExpectedExtensionLine::ExtensionDeclarations(kind) => {
                write!(f, "Extension Declaration for {kind}")
            }
        }
    }
}

/// The parser for the extension section of the Substrait file format.
///
/// This is responsible for parsing the extension section of the file, which
/// contains the extension URNs and declarations. Note that this parser does not
/// parse the header; otherwise, this is symmetric with the
/// SimpleExtensions::write method.
#[derive(Debug)]
pub struct ExtensionParser {
    state: ExpectedExtensionLine,
    extensions: SimpleExtensions,
}

impl Default for ExtensionParser {
    fn default() -> Self {
        Self {
            state: ExpectedExtensionLine::Extensions,
            extensions: SimpleExtensions::new(),
        }
    }
}

impl ExtensionParser {
    pub fn parse_line(&mut self, line: IndentedLine) -> Result<(), ExtensionParseError> {
        if line.1.is_empty() {
            // Blank lines are allowed between subsections, so if we see
            // one, we revert out of the subsection.
            self.state = ExpectedExtensionLine::Extensions;
            return Ok(());
        }

        match self.state {
            ExpectedExtensionLine::Extensions => self.parse_subsection(line),
            ExpectedExtensionLine::ExtensionUrns => self.parse_extension_urns(line),
            ExpectedExtensionLine::ExtensionDeclarations(extension_kind) => {
                self.parse_declarations(line, extension_kind)
            }
        }
    }

    fn parse_subsection(&mut self, line: IndentedLine) -> Result<(), ExtensionParseError> {
        match line {
            IndentedLine(0, simple::EXTENSION_URNS_HEADER) => {
                self.state = ExpectedExtensionLine::ExtensionUrns;
                Ok(())
            }
            IndentedLine(0, simple::EXTENSION_FUNCTIONS_HEADER) => {
                self.state = ExpectedExtensionLine::ExtensionDeclarations(ExtensionKind::Function);
                Ok(())
            }
            IndentedLine(0, simple::EXTENSION_TYPES_HEADER) => {
                self.state = ExpectedExtensionLine::ExtensionDeclarations(ExtensionKind::Type);
                Ok(())
            }
            IndentedLine(0, simple::EXTENSION_TYPE_VARIATIONS_HEADER) => {
                self.state =
                    ExpectedExtensionLine::ExtensionDeclarations(ExtensionKind::TypeVariation);
                Ok(())
            }
            _ => Err(ExtensionParseError::UnexpectedLine(self.state)),
        }
    }

    fn parse_extension_urns(&mut self, line: IndentedLine) -> Result<(), ExtensionParseError> {
        match line {
            IndentedLine(0, _s) => self.parse_subsection(line), // Pass the original line with 0 indent
            IndentedLine(1, s) => {
                let urn =
                    URNExtensionDeclaration::from_str(s).map_err(ExtensionParseError::Message)?;
                self.extensions.add_extension_urn(urn.urn, urn.anchor)?;
                Ok(())
            }
            _ => Err(ExtensionParseError::UnexpectedLine(self.state)),
        }
    }

    fn parse_declarations(
        &mut self,
        line: IndentedLine,
        extension_kind: ExtensionKind,
    ) -> Result<(), ExtensionParseError> {
        match line {
            IndentedLine(0, _s) => self.parse_subsection(line), // Pass the original line with 0 indent
            IndentedLine(1, s) => {
                let decl = SimpleExtensionDeclaration::parse_from_kind(s, extension_kind)?;
                self.extensions.add_extension(
                    extension_kind,
                    decl.urn_anchor,
                    decl.anchor,
                    decl.name,
                )?;
                Ok(())
            }
            _ => Err(ExtensionParseError::UnexpectedLine(self.state)),
        }
    }

    pub fn extensions(&self) -> &SimpleExtensions {
        &self.extensions
    }

    #[cfg(test)]
    pub(crate) fn state(&self) -> ExpectedExtensionLine {
        self.state
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct URNExtensionDeclaration {
    pub anchor: u32,
    pub urn: String,
}

#[derive(Debug, Clone, PartialEq)]
pub struct SimpleExtensionDeclaration {
    pub anchor: u32,
    pub urn_anchor: u32,
    pub name: String,
}

impl ParsePair for URNExtensionDeclaration {
    fn rule() -> Rule {
        Rule::extension_urn_declaration
    }

    fn message() -> &'static str {
        "URNExtensionDeclaration"
    }

    fn parse_pair(pair: Pair<Rule>) -> Self {
        assert_eq!(pair.as_rule(), Self::rule());

        let mut iter = RuleIter::from(pair.into_inner());
        let anchor_pair = iter.pop(Rule::urn_anchor);
        let anchor = unwrap_single_pair(anchor_pair)
            .as_str()
            .parse::<u32>()
            .unwrap();
        let urn = iter.pop(Rule::urn).as_str().to_string();
        iter.done();

        URNExtensionDeclaration { anchor, urn }
    }
}

impl FromStr for URNExtensionDeclaration {
    type Err = super::MessageParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Self::parse_str(s)
    }
}

impl SimpleExtensionDeclaration {
    fn parse_from_kind(s: &str, kind: ExtensionKind) -> Result<Self, MessageParseError> {
        let mut pairs = <ExpressionParser as pest::Parser<Rule>>::parse(Rule::simple_extension, s)
            .map_err(|e| {
                MessageParseError::new("SimpleExtensionDeclaration", ErrorKind::Syntax, Box::new(e))
            })?;
        assert_eq!(pairs.as_str(), s);
        let pair = pairs.next().unwrap();
        let mut iter = RuleIter::from(pair.into_inner());

        let anchor = unwrap_single_pair(iter.pop(Rule::anchor))
            .as_str()
            .parse::<u32>()
            .unwrap();
        let urn_anchor = unwrap_single_pair(iter.pop(Rule::urn_anchor))
            .as_str()
            .parse::<u32>()
            .unwrap();
        let name_pair = iter.pop(Rule::simple_extension_name);
        let name_span = name_pair.as_span();
        let name = name_pair.as_str();

        if kind != ExtensionKind::Type && name.starts_with("u!") {
            return Err(MessageParseError::invalid(
                "simple_extension_name",
                name_span,
                format!("'u!' prefix is only valid for type declarations, not {kind}"),
            ));
        }
        if matches!(kind, ExtensionKind::Type | ExtensionKind::TypeVariation) && name.contains(':')
        {
            return Err(MessageParseError::invalid(
                "simple_extension_name",
                name_span,
                format!(
                    "type/type-variation names must not include a signature suffix, got '{name}'"
                ),
            ));
        }
        iter.done();

        Ok(SimpleExtensionDeclaration {
            anchor,
            urn_anchor,
            name: name.to_string(),
        })
    }
}

// Extension relation parsing implementations
// These were moved from extensions/registry.rs to maintain clean architecture

impl ScopedParsePair for ExtensionValue {
    fn rule() -> Rule {
        Rule::extension_argument
    }

    fn message() -> &'static str {
        "ExtensionValue"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());

        let inner = unwrap_single_pair(pair); // Extract the actual content

        Ok(match inner.as_rule() {
            Rule::enum_value => {
                // Strip leading '&' and store the identifier
                let s = inner.as_str().trim_start_matches('&').to_string();
                ExtensionValue::Enum(s)
            }
            Rule::reference => {
                // Reuse the existing FieldIndex parser, then extract the i32
                let field_index = FieldIndex::parse_pair(inner);
                ExtensionValue::from(Reference(field_index.0))
            }
            Rule::untyped_literal => {
                // Literal can contain integer, float, boolean, or string_literal
                let value_pair = unwrap_single_pair(inner);
                match value_pair.as_rule() {
                    Rule::string_literal => ExtensionValue::String(unescape_string(value_pair)),
                    Rule::integer => {
                        ExtensionValue::Integer(value_pair.as_str().parse::<i64>().unwrap())
                    }
                    Rule::float => {
                        ExtensionValue::Float(value_pair.as_str().parse::<f64>().unwrap())
                    }
                    Rule::boolean => ExtensionValue::Boolean(value_pair.as_str() == "true"),
                    _ => panic!(
                        "Unexpected extension scalar literal type: {:?}",
                        value_pair.as_rule()
                    ),
                }
            }
            Rule::tuple => {
                let tv = inner
                    .into_inner()
                    .map(|pair| ExtensionValue::parse_pair(extensions, pair))
                    .collect::<Result<TupleValue, MessageParseError>>()?;
                ExtensionValue::Tuple(tv)
            }
            Rule::expression => {
                let expr = Expression::parse_pair(extensions, inner)?;
                ExtensionValue::from(expr)
            }
            _ => panic!("Unexpected extension argument type: {:?}", inner.as_rule()),
        })
    }
}

impl ScopedParsePair for ExtensionColumn {
    fn rule() -> Rule {
        Rule::extension_column
    }

    fn message() -> &'static str {
        "ExtensionColumn"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());

        let inner = unwrap_single_pair(pair); // Extract the actual content

        Ok(match inner.as_rule() {
            Rule::named_column => {
                let mut iter = inner.into_inner();
                let name_pair = iter.next().unwrap(); // Grammar guarantees type exists
                let type_pair = iter.next().unwrap(); // Grammar guarantees type exists

                let name = Name::parse_pair(name_pair).0.to_string(); // Reuse existing Name parser
                let ty = Type::parse_pair(extensions, type_pair)?;

                ExtensionColumn::Named { name, r#type: ty }
            }
            Rule::reference => {
                // Reuse the existing FieldIndex parser, then extract the i32
                let field_index = FieldIndex::parse_pair(inner);
                ExtensionColumn::Expr(Reference(field_index.0).into())
            }
            Rule::expression => {
                let expr = Expression::parse_pair(extensions, inner)?;
                ExtensionColumn::Expr(expr.into())
            }
            _ => panic!("Unexpected extension column type: {:?}", inner.as_rule()),
        })
    }
}

/// Relation kind encoded by the text syntax prefix (`ExtensionLeaf`,
/// `ExtensionSingle`, or `ExtensionMulti`).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ExtensionRelationKind {
    Leaf,
    Single,
    Multi,
}

impl FromStr for ExtensionRelationKind {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "ExtensionLeaf" => Ok(ExtensionRelationKind::Leaf),
            "ExtensionSingle" => Ok(ExtensionRelationKind::Single),
            "ExtensionMulti" => Ok(ExtensionRelationKind::Multi),
            _ => Err(format!("Unknown extension relation type: {s}")),
        }
    }
}

impl ExtensionRelationKind {
    pub(crate) fn validate_child_count(self, child_count: usize) -> Result<(), String> {
        match self {
            ExtensionRelationKind::Leaf => {
                if child_count == 0 {
                    Ok(())
                } else {
                    Err(format!(
                        "ExtensionLeaf should have no input children, got {child_count}"
                    ))
                }
            }
            ExtensionRelationKind::Single => {
                if child_count == 1 {
                    Ok(())
                } else {
                    Err(format!(
                        "ExtensionSingle should have exactly 1 input child, got {child_count}"
                    ))
                }
            }
            ExtensionRelationKind::Multi => Ok(()),
        }
    }

    /// Create appropriate relation structure from extension detail and children.
    pub(crate) fn create_rel(self, detail: Option<Any>, children: Vec<Rel>) -> Rel {
        let rel_type = match self {
            ExtensionRelationKind::Leaf => RelType::ExtensionLeaf(ExtensionLeafRel {
                common: None,
                detail: detail.map(Into::into),
            }),
            ExtensionRelationKind::Single => {
                let input = children.into_iter().next();
                RelType::ExtensionSingle(Box::new(ExtensionSingleRel {
                    common: None,
                    detail: detail.map(Into::into),
                    input: input.map(Box::new),
                }))
            }
            ExtensionRelationKind::Multi => RelType::ExtensionMulti(ExtensionMultiRel {
                common: None,
                detail: detail.map(Into::into),
                inputs: children,
            }),
        };

        Rel {
            rel_type: Some(rel_type),
        }
    }
}

/// Fully parsed extension invocation, including the user-supplied name and the
/// structured argument payload.
#[derive(Debug, Clone)]
pub(crate) struct ExtensionInvocation {
    pub(crate) relation_kind: ExtensionRelationKind,
    pub(crate) name: String,
    pub(crate) args: ExtensionArgs,
}

impl ScopedParsePair for ExtensionInvocation {
    fn rule() -> Rule {
        Rule::extension_relation
    }

    fn message() -> &'static str {
        "ExtensionInvocation"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());

        let mut iter = pair.into_inner();

        // Parse extension name to determine relation type and custom name
        let extension_name_pair = iter.next().unwrap(); // Grammar guarantees extension_name exists
        let full_extension_name = extension_name_pair.as_str();

        // Extract the relation type and custom name from the extension name
        // (e.g., "ExtensionLeaf:ParquetScan" -> "ExtensionLeaf" and "ParquetScan")
        let (relation_type_str, custom_name) = if full_extension_name.contains(':') {
            let parts: Vec<&str> = full_extension_name.splitn(2, ':').collect();
            (parts[0], parts[1].to_string())
        } else {
            (full_extension_name, "UnknownExtension".to_string())
        };

        let relation_kind = ExtensionRelationKind::from_str(relation_type_str).unwrap();
        let mut args = ExtensionArgs::default();

        // Parse optional arguments
        let ext_arguments = iter.next().unwrap();
        match ext_arguments.as_rule() {
            Rule::arguments => {
                arguments_rule_parsing(extensions, ext_arguments, &mut args)?;
            }
            r => unreachable!("Unexpected rule in ExtensionArgs: {:?}", r),
        }

        // parse optional output columns
        let extension_columns = iter.next();
        if let Some(value) = extension_columns {
            match value.as_rule() {
                Rule::extension_columns => {
                    for col_pair in value.into_inner() {
                        if col_pair.as_rule() == Rule::extension_column {
                            let column = ExtensionColumn::parse_pair(extensions, col_pair)?;
                            args.output_columns.push(column);
                        }
                    }
                }
                r => unreachable!("Unexpected rule in ExtensionArgs: {:?}", r),
            }
        }

        Ok(ExtensionInvocation {
            relation_kind,
            name: custom_name,
            args,
        })
    }
}

/// A parsed `+` addendum line.
#[derive(Debug, Clone)]
pub(crate) struct AddendumInvocation {
    pub(crate) kind: AddendumKind,
    pub(crate) name: String,
    pub(crate) args: ExtensionArgs,
}

impl ScopedParsePair for AddendumInvocation {
    fn rule() -> Rule {
        Rule::addendum
    }

    fn message() -> &'static str {
        "AddendumInvocation"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());

        let mut iter = pair.into_inner();

        // First token: addendum_type - grammar guarantees a known addendum prefix.
        let type_pair = iter.next().unwrap(); // Grammar guarantees addendum_type exists
        let kind = match type_pair.as_str() {
            "Enh" => AddendumKind::Enhancement,
            "Opt" => AddendumKind::Optimization,
            "Ext" => AddendumKind::ExtensionTable,
            other => unreachable!("Unexpected addendum_type: {other}"),
        };

        // Second token: name
        let name_pair = iter.next().unwrap();
        let name = Name::parse_pair(name_pair).0.to_string();

        // Remaining token: arguments — grammar guarantees it is always present.
        let mut args = ExtensionArgs::default();

        let arguments_pair = iter.next().unwrap();
        match arguments_pair.as_rule() {
            Rule::arguments => {
                arguments_rule_parsing(extensions, arguments_pair, &mut args)?;
            }
            r => unreachable!("Unexpected rule in AddendumInvocation args: {r:?}"),
        }

        Ok(AddendumInvocation { kind, name, args })
    }
}

fn arguments_rule_parsing(
    extensions: &SimpleExtensions,
    inner_pair: Pair<'_, Rule>,
    args: &mut ExtensionArgs,
) -> Result<(), MessageParseError> {
    for arg in inner_pair.into_inner() {
        match arg.as_rule() {
            Rule::extension_arguments => {
                for arg_pair in arg.into_inner() {
                    assert_eq!(arg_pair.as_rule(), Rule::extension_argument);
                    args.push(ExtensionValue::parse_pair(extensions, arg_pair)?);
                }
            }
            Rule::extension_named_arguments => {
                for arg_pair in arg.into_inner() {
                    assert_eq!(arg_pair.as_rule(), Rule::extension_named_argument);
                    let mut arg_iter = arg_pair.into_inner();
                    let name_p = arg_iter.next().unwrap();
                    let value_p = arg_iter.next().unwrap();
                    let key = Name::parse_pair(name_p).0.to_string();
                    let val = ExtensionValue::parse_pair(extensions, value_p)?;
                    args.insert(key, val);
                }
            }
            Rule::empty => {}
            r => unreachable!("Unexpected rule in extension args: {r:?}"),
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use substrait::proto;
    use substrait::proto::expression::RexType;
    use substrait::proto::expression::literal::LiteralType;

    use super::*;
    use crate::extensions::{Expr, ExtensionValue};
    use crate::fixtures::TestContext;
    use crate::parser::common::test_support::ScopedParse;
    use crate::parser::{ParseError, Parser};
    use crate::{OutputOptions, format};

    fn parse_extension_value(text: &str) -> ExtensionValue {
        ExtensionValue::parse(&SimpleExtensions::default(), text).unwrap()
    }

    #[test]
    fn test_parse_urn_extension_declaration() {
        let line = "@1: /my/urn1";
        let urn = URNExtensionDeclaration::parse_str(line).unwrap();
        assert_eq!(urn.anchor, 1);
        assert_eq!(urn.urn, "/my/urn1");
    }

    #[test]
    fn test_parse_simple_extension_declaration() {
        let line = "#5@2: my_function_name";
        let decl =
            SimpleExtensionDeclaration::parse_from_kind(line, ExtensionKind::Function).unwrap();
        assert_eq!(decl.anchor, 5);
        assert_eq!(decl.urn_anchor, 2);
        assert_eq!(decl.name, "my_function_name");

        let line2 = "#10  @200: another_ext_123";
        let decl =
            SimpleExtensionDeclaration::parse_from_kind(line2, ExtensionKind::Function).unwrap();
        assert_eq!(decl.anchor, 10);
        assert_eq!(decl.urn_anchor, 200);
        assert_eq!(decl.name, "another_ext_123");
    }

    #[test]
    fn test_parse_urn_extension_declaration_str() {
        let line = "@1: /my/urn1";
        let urn = URNExtensionDeclaration::parse_str(line).unwrap();
        assert_eq!(urn.anchor, 1);
        assert_eq!(urn.urn, "/my/urn1");
    }

    #[test]
    fn test_extensions_round_trip_plan() {
        let input = r#"
=== Extensions
URNs:
  @  1: /urn/common
  @  2: /urn/specific_funcs
Functions:
  # 10 @  1: func_a
  # 11 @  2: func_b_special
Types:
  # 20 @  1: SomeType
Type Variations:
  # 30 @  2: VarX
"#
        .trim_start();

        // Parse the input using the structural parser
        let plan = Parser::parse(input).unwrap();

        // Verify the plan has the expected extensions
        assert_eq!(plan.extension_urns.len(), 2);
        assert_eq!(plan.extensions.len(), 4);

        // Convert the plan extensions back to SimpleExtensions
        let (extensions, errors) =
            SimpleExtensions::from_extensions(&plan.extension_urns, &plan.extensions);

        assert!(errors.is_empty());
        // Convert back to string
        let output = extensions.to_string("  ");

        // The output should match the input
        assert_eq!(output, input);
    }

    #[test]
    fn test_parse_simple_extension_declaration_compound_name() {
        // A function name that includes a Substrait signature suffix
        let line = "#1 @2: equal:any_any";
        let decl =
            SimpleExtensionDeclaration::parse_from_kind(line, ExtensionKind::Function).unwrap();
        assert_eq!(decl.anchor, 1);
        assert_eq!(decl.urn_anchor, 2);
        assert_eq!(decl.name, "equal:any_any");
    }

    #[test]
    fn test_parse_simple_extension_declaration_compound_name_multi_segment() {
        let line = "#3 @1: regexp_match_substring:str_str_i64";
        let decl =
            SimpleExtensionDeclaration::parse_from_kind(line, ExtensionKind::Function).unwrap();
        assert_eq!(decl.anchor, 3);
        assert_eq!(decl.urn_anchor, 1);
        assert_eq!(decl.name, "regexp_match_substring:str_str_i64");
    }

    #[test]
    fn test_parse_simple_extension_declaration_u_prefix_function_with_u_signature() {
        // u! is valid inside a signature suffix (e.g. u!json as an arg type); only
        // the base function name itself may not be u!-prefixed.
        let line = "#5 @2: json_extract_path:u!json_str";
        let decl =
            SimpleExtensionDeclaration::parse_from_kind(line, ExtensionKind::Function).unwrap();
        assert_eq!(decl.anchor, 5);
        assert_eq!(decl.urn_anchor, 2);
        assert_eq!(decl.name, "json_extract_path:u!json_str");
    }

    #[test]
    fn test_u_prefix_type_declaration_accepted() {
        // u! prefix on a type name is non-standard but accepted; normalized to bare name at storage.
        let plan_text = "\
=== Extensions
URNs:
  @  1: https://example.com/types
Types:
  # 11 @  1: u!point
=== Plan
Root[result]
  Project[$0]
    Read[data => p:point#11]";
        let plan = Parser::parse(plan_text).unwrap();
        let (text, errors) = format(&plan);
        assert!(errors.is_empty(), "unexpected errors: {errors:?}");
        assert!(
            text.contains("  # 11 @  1: point"),
            "declaration line must use bare name"
        );
        assert!(
            !text.contains("u!point"),
            "u! prefix should be stripped in output"
        );
    }

    #[test]
    fn test_u_prefix_type_variation_declaration_rejected() {
        // u! prefix on a type variation name is invalid, same as for functions.
        let plan_text = "\
=== Extensions
URNs:
  @  1: https://example.com/types
Type Variations:
  # 30 @  1: u!myvar
=== Plan
Root[result]
  Read[data => x:i64]";
        assert!(
            Parser::parse(plan_text).is_err(),
            "u! prefix on a type variation name should be rejected"
        );
    }

    #[test]
    fn test_u_prefix_function_declaration_rejected() {
        // u! prefix on a function base name is invalid; function names are never u!-prefixed.
        let plan_text = "\
=== Extensions
URNs:
  @  1: https://example.com/funcs
Functions:
  # 21 @  1: u!json_get
=== Plan
Root[result]
  Read[data => x:i64]";
        assert!(
            Parser::parse(plan_text).is_err(),
            "u! prefix on a function name should be rejected"
        );
    }

    #[test]
    fn test_signature_on_type_declaration_rejected() {
        // Function signatures (':' suffix) are invalid on type declarations.
        let plan_text = "\
=== Extensions
URNs:
  @  1: https://example.com/types
Types:
  # 10 @  1: mytype:i64_i64
=== Plan
Root[result]
  Read[data => x:i64]";
        assert!(
            Parser::parse(plan_text).is_err(),
            "function signature suffix on a type declaration should be rejected"
        );
    }

    #[test]
    fn test_u_prefix_function_rejected_at_parser_level() {
        // The parser should reject function names with a `u!` prefix
        let plan_text = "\
=== Extensions
URNs:
  @  1: https://example.com/funcs
Functions:
  # 21 @  1: u!bad_func
=== Plan
Root[result]
  Read[data => x:i64]";
        let err = Parser::parse(plan_text).unwrap_err();
        assert!(
            matches!(
                err,
                ParseError::Extension(_, ExtensionParseError::Message(_))
            ),
            "expected parser-level MessageParseError, got: {err}"
        );
    }

    #[test]
    fn test_extensions_round_trip_plan_with_compound_names() {
        let input = r#"=== Extensions
URNs:
  @  1: extension:io.substrait:functions_string
  @  2: extension:io.substrait:functions_comparison
Functions:
  #  1 @  2: equal:any_any
  #  2 @  1: regexp_match_substring:str_str
  #  3 @  1: regexp_match_substring:str_str_i64
"#;
        let plan = Parser::parse(input).unwrap();
        let (extensions, errors) =
            SimpleExtensions::from_extensions(&plan.extension_urns, &plan.extensions);
        assert!(errors.is_empty());
        // Compound names must survive the roundtrip
        assert_eq!(
            extensions
                .find_by_anchor(ExtensionKind::Function, 1)
                .unwrap()
                .1
                .full(),
            "equal:any_any"
        );
        assert_eq!(
            extensions
                .find_by_anchor(ExtensionKind::Function, 3)
                .unwrap()
                .1
                .full(),
            "regexp_match_substring:str_str_i64"
        );
        // Text output must reproduce the input exactly
        assert_eq!(extensions.to_string("  "), input);
    }

    #[test]
    fn test_tuple_mixed_types_parses() {
        // tuple has overlapping grammar syntax with expression.
        let val = parse_extension_value("(&HASH, 8, 'hello')");
        let ExtensionValue::Tuple(items) = val else {
            panic!("expected Tuple, got {val:?}");
        };
        assert_eq!(items.len(), 3);
        let items: Vec<&ExtensionValue> = items.iter().collect();
        assert!(matches!(items[0], ExtensionValue::Enum(s) if s == "HASH"));
        assert_eq!(i64::try_from(items[1]).unwrap(), 8);
        assert_eq!(<&str>::try_from(items[2]).unwrap(), "hello");
    }

    #[test]
    fn test_empty_tuple_parses() {
        let val = parse_extension_value("()");
        let ExtensionValue::Tuple(items) = val else {
            panic!("expected Tuple, got {val:?}");
        };
        assert!(items.is_empty());
    }

    #[test]
    fn test_nested_tuple_parses() {
        let val = parse_extension_value("((&HASH, &RANGE), 8)");
        let ExtensionValue::Tuple(outer) = val else {
            panic!("expected Tuple, got {val:?}");
        };
        assert_eq!(outer.len(), 2);
        let ExtensionValue::Tuple(inner) = outer.iter().next().unwrap() else {
            panic!("expected inner Tuple");
        };
        assert_eq!(inner.len(), 2);
        assert!(matches!(inner.iter().next().unwrap(), ExtensionValue::Enum(s) if s == "HASH"));
        assert_eq!(i64::try_from(outer.iter().nth(1).unwrap()).unwrap(), 8);
    }

    #[test]
    fn test_tuple_in_addendum_parses() {
        let inv = AddendumInvocation::parse(
            &SimpleExtensions::default(),
            "+ Enh:Foo[(&HASH, &RANGE), count=8]",
        )
        .unwrap();
        assert_eq!(inv.kind, AddendumKind::Enhancement);
        assert_eq!(inv.name, "Foo");
        assert_eq!(inv.args.positional.len(), 1);
        let ExtensionValue::Tuple(items) = &inv.args.positional[0] else {
            panic!("expected Tuple positional arg");
        };
        assert_eq!(items.len(), 2);
        let items: Vec<&ExtensionValue> = items.iter().collect();
        assert!(matches!(items[0], ExtensionValue::Enum(s) if s == "HASH"));
        assert!(matches!(items[1], ExtensionValue::Enum(s) if s == "RANGE"));
        assert_eq!(inv.args.named.len(), 1);
    }

    #[test]
    fn extension_relation_kind_parses_text_prefixes() {
        assert_eq!(
            ExtensionRelationKind::from_str("ExtensionLeaf").unwrap(),
            ExtensionRelationKind::Leaf
        );
        assert_eq!(
            ExtensionRelationKind::from_str("ExtensionSingle").unwrap(),
            ExtensionRelationKind::Single
        );
        assert_eq!(
            ExtensionRelationKind::from_str("ExtensionMulti").unwrap(),
            ExtensionRelationKind::Multi
        );
    }

    #[test]
    fn extension_multi_allows_any_child_count() {
        assert!(ExtensionRelationKind::Multi.validate_child_count(0).is_ok());
        assert!(ExtensionRelationKind::Multi.validate_child_count(1).is_ok());
        assert!(ExtensionRelationKind::Multi.validate_child_count(3).is_ok());
    }

    #[test]
    fn extension_single_rejects_wrong_child_counts() {
        assert!(
            ExtensionRelationKind::Single
                .validate_child_count(0)
                .is_err()
        );
        assert!(
            ExtensionRelationKind::Single
                .validate_child_count(2)
                .is_err()
        );
    }

    #[test]
    fn test_tuple_textify_roundtrip() {
        let ctx = TestContext::new();
        for text in &[
            "(&HASH, &RANGE)",
            "(&HASH, 8, 'hello')",
            "()",
            "(&HASH,)",
            "((&HASH, &RANGE), 8)",
        ] {
            let val = parse_extension_value(text);
            let rendered = ctx.textify_no_errors(&val);
            assert_eq!(&rendered, text, "roundtrip failed for {text}");
        }
    }

    #[test]
    fn test_literal_expression_value_textifies_to_canonical_literal() {
        let expr = proto::Expression {
            rex_type: Some(RexType::Literal(proto::expression::Literal {
                literal_type: Some(LiteralType::I64(42)),
                nullable: false,
                type_variation_reference: 0,
            })),
        };
        let value = ExtensionValue::from(expr.clone());
        let ctx = TestContext::new();

        let rendered = ctx.textify_no_errors(&value);
        assert_eq!(rendered, "42");

        let parsed = parse_extension_value(&rendered);
        let parsed_expr = Expr::try_from(&parsed).unwrap();
        assert_eq!(parsed_expr.as_proto(), &expr);
    }

    #[test]
    fn test_extension_scalar_literals_stay_scalar_in_verbose_output() {
        let ctx = TestContext::new().with_options(OutputOptions::verbose());

        let scalar = ExtensionValue::from(42_i64);
        assert_eq!(ctx.textify_no_errors(&scalar), "42");

        let expression = ExtensionValue::from(Expr::from(42_i64));
        assert_eq!(ctx.textify_no_errors(&expression), "42:i64");
    }

    #[test]
    fn test_typed_extension_literal_parses_as_expression() {
        let value = parse_extension_value("42:i16");
        assert!(i64::try_from(&value).is_err());

        let expr = Expr::try_from(&value).unwrap();
        assert_eq!(ctx_text(&expr), "42:i16");
    }

    fn ctx_text(value: &Expr) -> String {
        TestContext::new().textify_no_errors(value)
    }
}