Allow importing enum variants (FuelLabs#4381)

sway-core/src/control_flow_analysis/analyze_return_paths.rs

@@ -186,6 +186,7 @@ fn connect_declaration<'eng: 'cfg, 'cfg>(
         | AbiDecl { .. }
         | StructDecl { .. }
         | EnumDecl { .. }
+        | EnumVariantDecl { .. }
         | StorageDecl { .. }
         | TypeAliasDecl { .. }
         | GenericTypeForFunctionScope { .. } => Ok(leaves.to_vec()),

sway-core/src/control_flow_analysis/dead_code_analysis.rs

@@ -473,7 +473,7 @@ fn connect_declaration<'eng: 'cfg, 'cfg>(
             connect_struct_declaration(&struct_decl, *decl_id, graph, entry_node, tree_type);
             Ok(leaves.to_vec())
         }
-        EnumDecl { decl_id, .. } => {
+        EnumDecl { decl_id, .. } | EnumVariantDecl { decl_id, .. } => {
             let enum_decl = decl_engine.get_enum(decl_id);
             connect_enum_declaration(&enum_decl, *decl_id, graph, entry_node);
             Ok(leaves.to_vec())
@@ -2006,6 +2006,17 @@ fn allow_dead_code_ast_node(decl_engine: &DeclEngine, node: &ty::TyAstNode) -> b
             ty::TyDecl::EnumDecl { decl_id, .. } => {
                 allow_dead_code(decl_engine.get_enum(decl_id).attributes)
             }
+            ty::TyDecl::EnumVariantDecl {
+                decl_id,
+                variant_name,
+                ..
+            } => decl_engine
+                .get_enum(decl_id)
+                .variants
+                .into_iter()
+                .find(|v| v.name == *variant_name)
+                .map(|enum_variant| allow_dead_code(enum_variant.attributes))
+                .unwrap_or(false),
             ty::TyDecl::TypeAliasDecl { .. } => {
                 // TODO - handle type aliases properly. For now, always skip DCA for them.
                 true

sway-core/src/ir_generation/compile.rs

@@ -284,6 +284,7 @@ fn compile_declarations(
 
             ty::TyDecl::StructDecl { .. }
             | ty::TyDecl::EnumDecl { .. }
+            | ty::TyDecl::EnumVariantDecl { .. }
             | ty::TyDecl::TraitDecl { .. }
             | ty::TyDecl::VariableDecl(_)
             | ty::TyDecl::AbiDecl { .. }

sway-core/src/ir_generation/function.rs

@@ -187,6 +187,7 @@ impl<'eng> FnCompiler<'eng> {
                 ty::TyDecl::GenericTypeForFunctionScope { .. } => unexpected_decl("generic type"),
                 ty::TyDecl::ErrorRecovery { .. } => unexpected_decl("error recovery"),
                 ty::TyDecl::StorageDecl { .. } => unexpected_decl("storage"),
+                ty::TyDecl::EnumVariantDecl { .. } => unexpected_decl("enum variant"),
             },
             ty::TyAstNodeContent::Expression(te) => {
                 // An expression with an ignored return value... I assume.

sway-core/src/language/parsed/expression/mod.rs

@@ -81,21 +81,25 @@ pub struct SubfieldExpression {
 
 #[derive(Debug, Clone)]
 pub struct AmbiguousSuffix {
-    /// The ambiguous part of the suffix.
+    /// If the suffix is a pair, the ambiguous part of the suffix.
     ///
     /// For example, if we have `Foo::bar()`,
     /// we don't know whether `Foo` is a module or a type,
     /// so `before` would be `Foo` here with any type arguments.
-    pub before: TypeBinding<Ident>,
-    /// The final suffix, i.e., the function name.
+    pub before: Option<TypeBinding<Ident>>,
+    /// The final suffix, i.e., the function or variant name.
     ///
     /// In the example above, this would be `bar`.
     pub suffix: Ident,
 }
 
 impl Spanned for AmbiguousSuffix {
     fn span(&self) -> Span {
-        Span::join(self.before.span(), self.suffix.span())
+        if let Some(before) = &self.before {
+            Span::join(before.span(), self.suffix.span())
+        } else {
+            self.suffix.span()
+        }
     }
 }
 

sway-core/src/language/ty/declaration/declaration.rs

@@ -47,6 +47,12 @@ pub enum TyDecl {
         subst_list: Template<SubstList>,
         decl_span: Span,
     },
+    EnumVariantDecl {
+        decl_id: DeclId<TyEnumDecl>,
+        subst_list: Template<SubstList>,
+        variant_name: Ident,
+        variant_decl_span: Span,
+    },
     ImplTrait {
         name: Ident,
         decl_id: DeclId<TyImplTrait>,
@@ -218,6 +224,14 @@ impl HashWithEngines for TyDecl {
             EnumDecl { decl_id, .. } => {
                 decl_engine.get(decl_id).hash(state, engines);
             }
+            EnumVariantDecl {
+                decl_id,
+                variant_name,
+                ..
+            } => {
+                decl_engine.get(decl_id).hash(state, engines);
+                variant_name.hash(state);
+            }
             ImplTrait { decl_id, .. } => {
                 decl_engine.get(decl_id).hash(state, engines);
             }
@@ -255,6 +269,9 @@ impl SubstTypes for TyDecl {
             } => decl_id.subst(type_mapping, engines),
             EnumDecl {
                 ref mut decl_id, ..
+            }
+            | EnumVariantDecl {
+                ref mut decl_id, ..
             } => decl_id.subst(type_mapping, engines),
             ImplTrait {
                 ref mut decl_id, ..
@@ -288,6 +305,9 @@ impl ReplaceSelfType for TyDecl {
             } => decl_id.replace_self_type(engines, self_type),
             EnumDecl {
                 ref mut decl_id, ..
+            }
+            | EnumVariantDecl {
+                ref mut decl_id, ..
             } => decl_id.replace_self_type(engines, self_type),
             ImplTrait {
                 ref mut decl_id, ..
@@ -335,6 +355,9 @@ impl Spanned for TyDecl {
             | AbiDecl { decl_span, .. }
             | StructDecl { decl_span, .. }
             | EnumDecl { decl_span, .. } => decl_span.clone(),
+            EnumVariantDecl {
+                variant_decl_span, ..
+            } => variant_decl_span.clone(),
             GenericTypeForFunctionScope { name, .. } => name.span(),
             ErrorRecovery(span) => span.clone(),
         }
@@ -437,6 +460,7 @@ impl CollectTypesMetadata for TyDecl {
             | TraitDecl { .. }
             | StructDecl { .. }
             | EnumDecl { .. }
+            | EnumVariantDecl { .. }
             | ImplTrait { .. }
             | AbiDecl { .. }
             | TypeAliasDecl { .. }
@@ -463,6 +487,7 @@ impl GetDeclIdent for TyDecl {
             | TyDecl::GenericTypeForFunctionScope { name, .. }
             | TyDecl::StructDecl { name, .. }
             | TyDecl::EnumDecl { name, .. } => Some(name.clone()),
+            TyDecl::EnumVariantDecl { variant_name, .. } => Some(variant_name.clone()),
             TyDecl::ErrorRecovery(_) => None,
             TyDecl::StorageDecl { .. } => None,
         }
@@ -659,6 +684,7 @@ impl TyDecl {
             TraitDecl { .. } => "trait",
             StructDecl { .. } => "struct",
             EnumDecl { .. } => "enum",
+            EnumVariantDecl { .. } => "enum variant",
             ImplTrait { .. } => "impl trait",
             AbiDecl { .. } => "abi",
             GenericTypeForFunctionScope { .. } => "generic type parameter",
@@ -759,6 +785,10 @@ impl TyDecl {
                 let TyEnumDecl { visibility, .. } = decl_engine.get_enum(decl_id);
                 visibility
             }
+            EnumVariantDecl { decl_id, .. } => {
+                let TyEnumDecl { visibility, .. } = decl_engine.get_enum(decl_id);
+                visibility
+            }
             FunctionDecl { decl_id, .. } => {
                 let TyFunctionDecl { visibility, .. } = decl_engine.get_function(decl_id);
                 visibility

sway-core/src/monomorphize/gather/declaration.rs

@@ -22,6 +22,7 @@ pub(crate) fn gather_from_decl(
         ty::TyDecl::TraitDecl { .. } => todo!(),
         ty::TyDecl::StructDecl { .. } => todo!(),
         ty::TyDecl::EnumDecl { .. } => todo!(),
+        ty::TyDecl::EnumVariantDecl { .. } => todo!(),
         ty::TyDecl::ImplTrait { .. } => todo!(),
         ty::TyDecl::AbiDecl { .. } => todo!(),
         ty::TyDecl::GenericTypeForFunctionScope { .. } => todo!(),

sway-core/src/monomorphize/instruct/declaration.rs

@@ -22,6 +22,7 @@ pub(crate) fn instruct_decl(
         ty::TyDecl::TraitDecl { .. } => todo!(),
         ty::TyDecl::StructDecl { .. } => todo!(),
         ty::TyDecl::EnumDecl { .. } => todo!(),
+        ty::TyDecl::EnumVariantDecl { .. } => todo!(),
         ty::TyDecl::ImplTrait { .. } => todo!(),
         ty::TyDecl::AbiDecl { .. } => todo!(),
         ty::TyDecl::GenericTypeForFunctionScope { .. } => todo!(),

sway-core/src/semantic_analysis/ast_node/declaration/impl_trait.rs

@@ -376,6 +376,7 @@ impl ty::TyImplTrait {
                 | ty::TyDecl::TraitDecl { .. }
                 | ty::TyDecl::StructDecl { .. }
                 | ty::TyDecl::EnumDecl { .. }
+                | ty::TyDecl::EnumVariantDecl { .. }
                 | ty::TyDecl::ImplTrait { .. }
                 | ty::TyDecl::AbiDecl { .. }
                 | ty::TyDecl::GenericTypeForFunctionScope { .. }

sway-core/src/semantic_analysis/ast_node/expression/match_expression/typed/typed_scrutinee.rs

@@ -4,7 +4,11 @@ use sway_types::{BaseIdent, Ident, Span, Spanned};
 use crate::{
     decl_engine::DeclEngineInsert,
     error::*,
-    language::{parsed::*, ty, CallPath},
+    language::{
+        parsed::*,
+        ty::{self, TyDecl},
+        CallPath,
+    },
     semantic_analysis::TypeCheckContext,
     type_system::*,
 };
@@ -236,44 +240,56 @@ fn type_check_enum(
     let decl_engine = ctx.decl_engine;
 
     let mut prefixes = call_path.prefixes.clone();
-    let enum_callpath = match prefixes.pop() {
-        Some(enum_name) => CallPath {
-            suffix: enum_name,
-            prefixes,
-            is_absolute: call_path.is_absolute,
-        },
+    let (callsite_span, mut enum_decl) = match prefixes.pop() {
+        Some(enum_name) => {
+            let enum_callpath = CallPath {
+                suffix: enum_name,
+                prefixes,
+                is_absolute: call_path.is_absolute,
+            };
+            // find the enum definition from the name
+            let unknown_decl = check!(
+                ctx.namespace.resolve_call_path(&enum_callpath).cloned(),
+                return err(warnings, errors),
+                warnings,
+                errors
+            );
+            let enum_ref = check!(
+                unknown_decl.to_enum_ref(ctx.engines()),
+                return err(warnings, errors),
+                warnings,
+                errors
+            );
+            (enum_callpath.span(), decl_engine.get_enum(&enum_ref))
+        }
         None => {
-            errors.push(CompileError::EnumNotFound {
-                name: call_path.suffix.clone(),
-                span: call_path.suffix.span(),
-            });
-            return err(warnings, errors);
+            // we may have an imported variant
+            let decl = check!(
+                ctx.namespace.resolve_call_path(&call_path).cloned(),
+                return err(warnings, errors),
+                warnings,
+                errors
+            );
+            if let TyDecl::EnumVariantDecl { decl_id, .. } = decl {
+                (call_path.suffix.span(), decl_engine.get_enum(&decl_id))
+            } else {
+                errors.push(CompileError::EnumNotFound {
+                    name: call_path.suffix.clone(),
+                    span: call_path.suffix.span(),
+                });
+                return err(warnings, errors);
+            }
         }
     };
     let variant_name = call_path.suffix.clone();
 
-    // find the enum definition from the name
-    let unknown_decl = check!(
-        ctx.namespace.resolve_call_path(&enum_callpath).cloned(),
-        return err(warnings, errors),
-        warnings,
-        errors
-    );
-    let enum_ref = check!(
-        unknown_decl.to_enum_ref(ctx.engines()),
-        return err(warnings, errors),
-        warnings,
-        errors
-    );
-    let mut enum_decl = decl_engine.get_enum(&enum_ref);
-
     // monomorphize the enum definition
     check!(
         ctx.monomorphize(
             &mut enum_decl,
             &mut [],
             EnforceTypeArguments::No,
-            &enum_callpath.span()
+            &callsite_span,
         ),
         return err(warnings, errors),
         warnings,

sway-core/src/semantic_analysis/ast_node/expression/typed_expression.rs

@@ -968,6 +968,37 @@ impl ty::TyExpression {
     ) -> CompileResult<ty::TyExpression> {
         let decl_engine = ctx.decl_engine;
 
+        // is it a singleton?
+        let before = if let Some(b) = before {
+            b
+        } else {
+            // if it's a singleton it's either an enum variant or a function
+            let call_path_binding = TypeBinding {
+                inner: CallPath {
+                    prefixes,
+                    suffix,
+                    is_absolute,
+                },
+                type_arguments,
+                span: path_span,
+            };
+            if matches!(
+                ctx.namespace
+                    .resolve_call_path(&call_path_binding.inner)
+                    .value,
+                Some(ty::TyDecl::EnumVariantDecl { .. })
+            ) {
+                return Self::type_check_delineated_path(ctx, call_path_binding, span, Some(args));
+            } else {
+                return Self::type_check_function_application(
+                    ctx.by_ref(),
+                    call_path_binding,
+                    args,
+                    span,
+                );
+            }
+        };
+
         // Is `path = prefix ++ before` a module?
         let mut path = Vec::with_capacity(prefixes.len() + 1);
         path.extend(prefixes.iter().cloned());

sway-core/src/semantic_analysis/ast_node/mod.rs

@@ -36,13 +36,56 @@ impl ty::TyAstNode {
                         ctx.namespace.find_module_path(&a.call_path)
                     };
                     let mut res = match a.import_type {
-                        ImportType::Star => ctx.namespace.star_import(&path, engines),
+                        ImportType::Star => {
+                            // try a standard starimport first
+                            let import = ctx.namespace.star_import(&path, engines);
+                            if import.is_ok() {
+                                import
+                            } else {
+                                // if it doesn't work it could be an enum star import
+                                if let Some((enum_name, path)) = path.split_last() {
+                                    let variant_import =
+                                        ctx.namespace.variant_star_import(path, engines, enum_name);
+                                    if variant_import.is_ok() {
+                                        variant_import
+                                    } else {
+                                        import
+                                    }
+                                } else {
+                                    import
+                                }
+                            }
+                        }
                         ImportType::SelfImport(_) => {
                             ctx.namespace.self_import(engines, &path, a.alias.clone())
                         }
                         ImportType::Item(ref s) => {
-                            ctx.namespace
-                                .item_import(engines, &path, s, a.alias.clone())
+                            // try a standard item import first
+                            let import =
+                                ctx.namespace
+                                    .item_import(engines, &path, s, a.alias.clone());
+
+                            if import.is_ok() {
+                                import
+                            } else {
+                                // if it doesn't work it could be an enum variant import
+                                if let Some((enum_name, path)) = path.split_last() {
+                                    let variant_import = ctx.namespace.variant_import(
+                                        engines,
+                                        path,
+                                        enum_name,
+                                        s,
+                                        a.alias.clone(),
+                                    );
+                                    if variant_import.is_ok() {
+                                        variant_import
+                                    } else {
+                                        import
+                                    }
+                                } else {
+                                    import
+                                }
+                            }
                         }
                     };
                     warnings.append(&mut res.warnings);