feat: support multiple variations with the same URI/name pair

See also substrait-io/substrait#268

rs/src/output/diagnostic.rs

@@ -237,6 +237,9 @@ pub enum Classification {
     #[strum(props(HiddenDescription = "use of anchor zero"))]
     LinkAnchorZero = 3005,
 
+    #[strum(props(Description = "failed to resolve type variation name & class pair"))]
+    LinkMissingTypeVariationNameAndClass = 3006,
+
     // Type-related diagnostics (group 4).
     #[strum(props(HiddenDescription = "type-related diagnostics"))]
     Type = 4000,

rs/src/output/extension.rs

@@ -168,6 +168,9 @@ pub struct YamlData {
     /// Reference to the parsed YAML data, if any.
     pub data: tree::NodeReference,
 
+    /// YAML extension files that this extension depends on.
+    pub dependencies: HashMap<String, YamlInfo>,
+
     /// Functions defined in this YAML file. Names are stored in lower case
     /// (Substrait's name resolution is case-insensitive).
     pub functions: HashMap<String, Arc<Function>>,
@@ -178,9 +181,7 @@ pub struct YamlData {
 
     /// Type variations defined in this YAML file. Names are stored in lower
     /// case (Substrait's name resolution is case-insensitive).
-    /// FIXME: this declaration to definition map is insufficient. See
-    /// <https://github.com/substrait-io/substrait/issues/268>
-    pub type_variations: HashMap<String, Arc<data::variation::UserDefinedDefinition>>,
+    pub type_variations: HashMap<String, Arc<data::variation::UserDefinedDefinitions>>,
 }
 
 impl YamlData {
@@ -194,6 +195,7 @@ impl YamlData {
                 path: path::Path::Root("").to_path_buf(),
                 node: Arc::new(tree::NodeType::YamlMap.into()),
             },
+            dependencies: HashMap::default(),
             functions: HashMap::default(),
             types: HashMap::default(),
             type_variations: HashMap::default(),
@@ -213,27 +215,44 @@ impl YamlData {
         })
     }
 
-    /// Resolves a function defined in this YAML data block by name. Returns an
-    /// unresolved reference if it does not exist.
+    /// Resolves a function defined in this YAML data block by the given
+    /// lowercase name. Returns an unresolved reference if it does not exist.
     pub fn resolve_function<S: ToString>(&self, name: S) -> Arc<Reference<Function>> {
         let name = name.to_string();
         let maybe_def = self.functions.get(&name).cloned();
         self.local_reference(name, maybe_def)
     }
 
-    /// Resolves a type defined in this YAML data block by name. Returns an
-    /// unresolved reference if it does not exist.
+    /// Resolves a type defined in this YAML data block by the given lowercase
+    /// name. Returns an unresolved reference if it does not exist.
     pub fn resolve_type<S: ToString>(&self, name: S) -> data::class::UserDefined {
         let name = name.to_string();
         let maybe_def = self.types.get(&name).cloned();
         self.local_reference(name, maybe_def)
     }
 
-    /// Resolves a type variation defined in this YAML data block by name.
-    /// Returns an unresolved reference if it does not exist.
-    pub fn resolve_type_variation<S: ToString>(&self, name: S) -> data::variation::UserDefined {
+    /// Resolves all type variations defined in this YAML data block by the
+    /// given lowercase name. Returns an unresolved reference if none exist.
+    pub fn resolve_type_variations_by_name<S: ToString>(
+        &self,
+        name: S,
+    ) -> data::variation::UserDefinedByName {
         let name = name.to_string();
-        let maybe_def = self.type_variations.get(&name).cloned();
+        let maybe_def = self
+            .type_variations
+            .get(&name)
+            .filter(|x| !x.variations.is_empty())
+            .cloned();
         self.local_reference(name, maybe_def)
     }
+
+    /// Resolves all type variations defined in this YAML data block by the
+    /// given lowercase name. Returns an unresolved reference if none exist.
+    pub fn resolve_type_variation<S: ToString>(
+        &self,
+        name: S,
+        base: &data::Class,
+    ) -> data::variation::UserDefined {
+        data::variation::resolve_by_class(&self.resolve_type_variations_by_name(name), base)
+    }
 }

rs/src/output/type_system/data/variation.rs

@@ -84,3 +84,32 @@ impl Default for FunctionBehavior {
         FunctionBehavior::Inherits
     }
 }
+
+/// A reference to one or more user-defined type variations going by the same
+/// name, distinguished by their base type class.
+pub type UserDefinedByName = Arc<extension::Reference<UserDefinedDefinitions>>;
+
+/// A group of one or more variation definitions using a single name. Note:
+/// multiple variations can be defined with the same name, because names are
+/// scoped to the type class they are defined for. See
+/// <https://github.com/substrait-io/substrait/issues/268>.
+#[derive(Clone, Debug, PartialEq, Eq, Default)]
+pub struct UserDefinedDefinitions {
+    pub variations: Vec<Arc<UserDefinedDefinition>>,
+}
+
+/// Resolve a reference to a set of variations going by the same name to a
+/// single variation indexed by its base class. Returns an unresolved reference
+/// if it does not exist.
+pub fn resolve_by_class(variations: &UserDefinedByName, base: &data::Class) -> UserDefined {
+    let definition = variations
+        .definition
+        .as_ref()
+        .and_then(|x| x.variations.iter().find(|x| &x.base == base))
+        .cloned();
+    Arc::new(extension::Reference {
+        name: variations.name.clone(),
+        uri: variations.uri.clone(),
+        definition,
+    })
+}

rs/src/parse/context.rs

@@ -305,7 +305,7 @@ impl<'a> Context<'a> {
         &mut self.state.types
     }
 
-    /// Registers a type definition. Shorthand for fns().define(), using the
+    /// Registers a type definition. Shorthand for types().define(), using the
     /// current path as the registration path.
     pub fn define_type(
         &mut self,
@@ -318,17 +318,17 @@ impl<'a> Context<'a> {
     }
 
     /// Returns the resolver for type variation anchors and references.
-    pub fn tvars(&mut self) -> &mut Resolver<u32, data::variation::UserDefined> {
+    pub fn tvars(&mut self) -> &mut Resolver<u32, data::variation::UserDefinedByName> {
         &mut self.state.type_variations
     }
 
-    /// Registers a type definition. Shorthand for fns().define(), using the
+    /// Registers a type definition. Shorthand for tvars().define(), using the
     /// current path as the registration path.
     pub fn define_tvar(
         &mut self,
         anchor: u32,
-        variation: data::variation::UserDefined,
-    ) -> Result<(), (data::variation::UserDefined, path::PathBuf)> {
+        variation: data::variation::UserDefinedByName,
+    ) -> Result<(), (data::variation::UserDefinedByName, path::PathBuf)> {
         self.state
             .type_variations
             .define(anchor, variation, self.breadcrumb.path.to_path_buf())
@@ -493,7 +493,7 @@ pub struct State {
     pub types: Resolver<u32, data::class::UserDefined>,
 
     /// YAML-defined type variation anchor resolver.
-    pub type_variations: Resolver<u32, data::variation::UserDefined>,
+    pub type_variations: Resolver<u32, data::variation::UserDefinedByName>,
 
     /// Protobuf Any type URL resolver.
     pub proto_any_types: Resolver<String, ()>,