const-generics for enums

sway-core/src/control_flow_analysis/dead_code_analysis.rs

@@ -2516,10 +2516,12 @@ fn connect_type_id<'eng: 'cfg, 'cfg>(
                 graph.add_edge(entry_node, enum_idx, "".into());
             }
             for p in &decl.generic_parameters {
-                let p = p
-                    .as_type_parameter()
-                    .expect("only works with type parameters");
-                connect_type_id(engines, p.type_id, graph, entry_node)?;
+                match p {
+                    crate::TypeParameter::Type(p) => {
+                        connect_type_id(engines, p.type_id, graph, entry_node)?;
+                    }
+                    crate::TypeParameter::Const(_) => {}
+                }
             }
         }
         TypeInfo::Struct(decl_ref) => {

sway-core/src/ir_generation.rs

@@ -74,13 +74,14 @@ impl CompiledFunctionCache {
         let new_callee = match item {
             Some(func) => func,
             None => {
+                let name = Ident::new(Span::from_string(format!(
+                    "{}_{}",
+                    decl.name,
+                    context.get_unique_symbol_id()
+                )));
                 let callee_fn_decl = ty::TyFunctionDecl {
                     type_parameters: Vec::new(),
-                    name: Ident::new(Span::from_string(format!(
-                        "{}_{}",
-                        decl.name,
-                        context.get_unique_symbol_id()
-                    ))),
+                    name,
                     parameters: decl.parameters.clone(),
                     ..decl.clone()
                 };

sway-core/src/ir_generation/types.rs

@@ -22,12 +22,12 @@ pub(super) fn create_tagged_union_type(
     // getting one here anyway.  They don't need to be a tagged union either.
     let field_types: Vec<_> = variants
         .iter()
-        .map(|tev| {
+        .map(|variant| {
             convert_resolved_typeid_no_span(
                 type_engine,
                 decl_engine,
                 context,
-                tev.type_argument.type_id(),
+                variant.type_argument.type_id(),
             )
         })
         .collect::<Result<Vec<_>, CompileError>>()?;

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

@@ -10,6 +10,7 @@ use crate::{
     type_system::*,
     types::*,
 };
+use ast_elements::type_parameter::ConstGenericExpr;
 use serde::{Deserialize, Serialize};
 use std::{
     fmt::{self, Debug},
@@ -148,9 +149,24 @@ impl MaterializeConstGenerics for TyAstNode {
         value: &TyExpression,
     ) -> Result<(), ErrorEmitted> {
         match &mut self.content {
-            TyAstNodeContent::Declaration(TyDecl::VariableDecl(decl)) => decl
-                .body
-                .materialize_const_generics(engines, handler, name, value),
+            TyAstNodeContent::Declaration(TyDecl::VariableDecl(decl)) => {
+                decl.body
+                    .materialize_const_generics(engines, handler, name, value)?;
+                decl.return_type
+                    .materialize_const_generics(engines, handler, name, value)?;
+                match &mut decl.type_ascription {
+                    GenericArgument::Type(arg) => arg
+                        .type_id
+                        .materialize_const_generics(engines, handler, name, value)?,
+                    GenericArgument::Const(arg) => {
+                        if matches!(&arg.expr, ConstGenericExpr::AmbiguousVariableExpression { ident } if ident.as_str() == name)
+                        {
+                            arg.expr = ConstGenericExpr::from_ty_expression(handler, value)?;
+                        }
+                    }
+                }
+                Ok(())
+            }
             TyAstNodeContent::Expression(expr) => {
                 expr.materialize_const_generics(engines, handler, name, value)
             }

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

@@ -6,6 +6,7 @@ use crate::{
     transform,
     type_system::*,
 };
+use ast_elements::type_parameter::ConstGenericExpr;
 use monomorphization::MonomorphizeHelper;
 use serde::{Deserialize, Serialize};
 use std::{
@@ -107,11 +108,31 @@ impl MonomorphizeHelper for TyEnumDecl {
 impl MaterializeConstGenerics for TyEnumDecl {
     fn materialize_const_generics(
         &mut self,
-        _engines: &Engines,
-        _handler: &Handler,
-        _name: &str,
-        _value: &crate::language::ty::TyExpression,
+        engines: &Engines,
+        handler: &Handler,
+        name: &str,
+        value: &crate::language::ty::TyExpression,
     ) -> Result<(), ErrorEmitted> {
+        for p in self.generic_parameters.iter_mut() {
+            match p {
+                TypeParameter::Const(p) if p.name.as_str() == name => {
+                    p.expr = Some(ConstGenericExpr::from_ty_expression(handler, value)?);
+                }
+                TypeParameter::Type(p) => {
+                    p.type_id
+                        .materialize_const_generics(engines, handler, name, value)?;
+                }
+                _ => {}
+            }
+        }
+
+        for variant in self.variants.iter_mut() {
+            variant
+                .type_argument
+                .type_id_mut()
+                .materialize_const_generics(engines, handler, name, value)?;
+        }
+
         Ok(())
     }
 }

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

@@ -793,7 +793,9 @@ impl TyFunctionSig {
                         TyFunctionSigTypeParameter::Const(p) => {
                             match p {
                                 ConstGenericExpr::Literal { val, .. } => val.to_string(),
-                                ConstGenericExpr::AmbiguousVariableExpression { .. } => todo!(),
+                                ConstGenericExpr::AmbiguousVariableExpression { ident } => {
+                                    ident.as_str().to_string()
+                                }
                             }
                         }
                     })

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

@@ -101,19 +101,27 @@ impl MonomorphizeHelper for TyStructDecl {
 impl MaterializeConstGenerics for TyStructDecl {
     fn materialize_const_generics(
         &mut self,
-        _engines: &Engines,
-        _handler: &Handler,
-        _name: &str,
-        _value: &crate::language::ty::TyExpression,
+        engines: &Engines,
+        handler: &Handler,
+        name: &str,
+        value: &crate::language::ty::TyExpression,
     ) -> Result<(), ErrorEmitted> {
         for p in self.generic_parameters.iter_mut() {
             match p {
-                TypeParameter::Const(p) if p.name.as_str() == _name => {
-                    p.expr = Some(ConstGenericExpr::from_ty_expression(_handler, _value)?);
+                TypeParameter::Const(p) if p.name.as_str() == name => {
+                    p.expr = Some(ConstGenericExpr::from_ty_expression(handler, value)?);
                 }
                 _ => {}
             }
         }
+
+        for field in self.fields.iter_mut() {
+            field
+                .type_argument
+                .type_id_mut()
+                .materialize_const_generics(engines, handler, name, value)?;
+        }
+
         Ok(())
     }
 }

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

@@ -313,10 +313,12 @@ impl CollectTypesMetadata for TyExpression {
             } => {
                 let enum_decl = decl_engine.get_enum(enum_ref);
                 for p in enum_decl.generic_parameters.iter() {
-                    let p = p
-                        .as_type_parameter()
-                        .expect("only works for type parameters");
-                    ctx.call_site_insert(p.type_id, call_path_binding.inner.suffix.span())
+                    match p {
+                        TypeParameter::Type(p) => {
+                            ctx.call_site_insert(p.type_id, call_path_binding.inner.suffix.span())
+                        }
+                        TypeParameter::Const(_) => {}
+                    }
                 }
                 if let Some(contents) = contents {
                     res.append(&mut contents.collect_types_metadata(handler, ctx)?);
@@ -330,10 +332,12 @@ impl CollectTypesMetadata for TyExpression {
                     );
                 }
                 for p in enum_decl.generic_parameters.iter() {
-                    let p = p
-                        .as_type_parameter()
-                        .expect("only works for type parameters");
-                    res.append(&mut p.type_id.collect_types_metadata(handler, ctx)?);
+                    match p {
+                        TypeParameter::Type(p) => {
+                            res.append(&mut p.type_id.collect_types_metadata(handler, ctx)?);
+                        }
+                        TypeParameter::Const(_) => {}
+                    }
                 }
             }
             AbiCast { address, .. } => {
@@ -488,6 +492,65 @@ impl MaterializeConstGenerics for TyExpression {
             TyExpressionVariant::Deref(r) => {
                 r.materialize_const_generics(engines, handler, name, value)
             }
+            TyExpressionVariant::MatchExp { desugared, .. } => {
+                desugared.materialize_const_generics(engines, handler, name, value)
+            }
+            TyExpressionVariant::EnumInstantiation { contents, .. } => {
+                if let Some(contents) = contents.as_mut() {
+                    contents.materialize_const_generics(engines, handler, name, value)?;
+                }
+                Ok(())
+            }
+            TyExpressionVariant::EnumTag { exp } => {
+                exp.materialize_const_generics(engines, handler, name, value)
+            }
+            TyExpressionVariant::Tuple { fields } => {
+                for f in fields {
+                    f.materialize_const_generics(engines, handler, name, value)?;
+                }
+                Ok(())
+            }
+            TyExpressionVariant::TupleElemAccess {
+                prefix,
+                resolved_type_of_parent,
+                ..
+            } => {
+                prefix.materialize_const_generics(engines, handler, name, value)?;
+                resolved_type_of_parent
+                    .materialize_const_generics(engines, handler, name, value)?;
+                Ok(())
+            }
+            TyExpressionVariant::LazyOperator { lhs, rhs, .. } => {
+                lhs.materialize_const_generics(engines, handler, name, value)?;
+                rhs.materialize_const_generics(engines, handler, name, value)
+            }
+            TyExpressionVariant::AsmExpression { registers, .. } => {
+                for r in registers.iter_mut() {
+                    if let Some(init) = r.initializer.as_mut() {
+                        init.materialize_const_generics(engines, handler, name, value)?;
+                    }
+                }
+                Ok(())
+            }
+            TyExpressionVariant::ConstantExpression { .. } => Ok(()),
+            TyExpressionVariant::StructExpression { fields, .. } => {
+                for f in fields {
+                    f.value
+                        .materialize_const_generics(engines, handler, name, value)?;
+                }
+                Ok(())
+            }
+            TyExpressionVariant::StructFieldAccess {
+                prefix,
+                resolved_type_of_parent,
+                ..
+            } => {
+                prefix.materialize_const_generics(engines, handler, name, value)?;
+                resolved_type_of_parent.materialize_const_generics(engines, handler, name, value)
+            }
+            TyExpressionVariant::UnsafeDowncast { exp, .. } => {
+                exp.materialize_const_generics(engines, handler, name, value)
+            }
             _ => Err(handler.emit_err(
                 sway_error::error::CompileError::ConstGenericNotSupportedHere {
                     span: self.span.clone(),

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

@@ -75,7 +75,7 @@ impl ty::TyCodeBlock {
                 Ok(())
             })?;
 
-        ctx.engines.te().reapply_unifications(ctx.engines());
+        ctx.engines.te().reapply_unifications(ctx.engines(), 0);
 
         ctx.by_ref()
             .scoped(handler, Some(code_block.span()), |ctx| {

sway-core/src/semantic_analysis/ast_node/declaration/auto_impl/abi_encoding.rs

@@ -31,22 +31,24 @@ where
         type_parameters: &[TypeParameter],
         body: String,
     ) -> String {
+        let type_parameters_declaration_expanded =
+            self.generate_type_parameters_declaration_code(type_parameters, true);
         let type_parameters_declaration =
-            self.generate_type_parameters_declaration_code(type_parameters);
+            self.generate_type_parameters_declaration_code(type_parameters, false);
         let type_parameters_constraints =
             self.generate_type_parameters_constraints_code(type_parameters, Some("AbiEncode"));
 
         let name = name.as_raw_ident_str();
 
         if body.is_empty() {
-            format!("#[allow(dead_code, deprecated)] impl{type_parameters_declaration} AbiEncode for {name}{type_parameters_declaration}{type_parameters_constraints} {{
+            format!("#[allow(dead_code, deprecated)] impl{type_parameters_declaration_expanded} AbiEncode for {name}{type_parameters_declaration}{type_parameters_constraints} {{
                 #[allow(dead_code, deprecated)]
                 fn abi_encode(self, buffer: Buffer) -> Buffer {{
                     buffer
                 }}
             }}")
         } else {
-            format!("#[allow(dead_code, deprecated)] impl{type_parameters_declaration} AbiEncode for {name}{type_parameters_declaration}{type_parameters_constraints} {{
+            format!("#[allow(dead_code, deprecated)] impl{type_parameters_declaration_expanded} AbiEncode for {name}{type_parameters_declaration}{type_parameters_constraints} {{
                 #[allow(dead_code, deprecated)]
                 fn abi_encode(self, buffer: Buffer) -> Buffer {{
                     {body}
@@ -62,22 +64,24 @@ where
         type_parameters: &[TypeParameter],
         body: String,
     ) -> String {
+        let type_parameters_declaration_expanded =
+            self.generate_type_parameters_declaration_code(type_parameters, true);
         let type_parameters_declaration =
-            self.generate_type_parameters_declaration_code(type_parameters);
+            self.generate_type_parameters_declaration_code(type_parameters, false);
         let type_parameters_constraints =
             self.generate_type_parameters_constraints_code(type_parameters, Some("AbiDecode"));
 
         let name = name.as_raw_ident_str();
 
         if body == "Self {  }" {
-            format!("#[allow(dead_code, deprecated)] impl{type_parameters_declaration} AbiDecode for {name}{type_parameters_declaration}{type_parameters_constraints} {{
+            format!("#[allow(dead_code, deprecated)] impl{type_parameters_declaration_expanded} AbiDecode for {name}{type_parameters_declaration}{type_parameters_constraints} {{
                 #[allow(dead_code, deprecated)]
                 fn abi_decode(ref mut _buffer: BufferReader) -> Self {{
                     {body}
                 }}
             }}")
         } else {
-            format!("#[allow(dead_code, deprecated)] impl{type_parameters_declaration} AbiDecode for {name}{type_parameters_declaration}{type_parameters_constraints} {{
+            format!("#[allow(dead_code, deprecated)] impl{type_parameters_declaration_expanded} AbiDecode for {name}{type_parameters_declaration}{type_parameters_constraints} {{
                 #[allow(dead_code, deprecated)]
                 fn abi_decode(ref mut buffer: BufferReader) -> Self {{
                     {body}

sway-core/src/semantic_analysis/ast_node/declaration/auto_impl/debug.rs

@@ -73,14 +73,16 @@ where
         type_parameters: &[TypeParameter],
         body: String,
     ) -> String {
+        let type_parameters_declaration_expanded =
+            self.generate_type_parameters_declaration_code(type_parameters, true);
         let type_parameters_declaration =
-            self.generate_type_parameters_declaration_code(type_parameters);
+            self.generate_type_parameters_declaration_code(type_parameters, false);
         let type_parameters_constraints =
             self.generate_type_parameters_constraints_code(type_parameters, Some("Debug"));
 
         let name = name.as_raw_ident_str();
 
-        format!("#[allow(dead_code, deprecated)] impl{type_parameters_declaration} Debug for {name}{type_parameters_declaration}{type_parameters_constraints} {{
+        format!("#[allow(dead_code, deprecated)] impl{type_parameters_declaration_expanded} Debug for {name}{type_parameters_declaration}{type_parameters_constraints} {{
             #[allow(dead_code, deprecated)]
             fn fmt(self, ref mut _f: Formatter) {{
                 {body}

sway-core/src/semantic_analysis/ast_node/declaration/auto_impl/marker_traits.rs

@@ -45,15 +45,17 @@ where
             return None;
         }
 
+        let type_parameters_declaration_expanded =
+            self.generate_type_parameters_declaration_code(&enum_decl.generic_parameters, true);
         let type_parameters_declaration =
-            self.generate_type_parameters_declaration_code(&enum_decl.generic_parameters);
+            self.generate_type_parameters_declaration_code(&enum_decl.generic_parameters, false);
         let type_parameters_constraints = self.generate_type_parameters_constraints_code(
             &enum_decl.generic_parameters,
             extra_constraint,
         );
 
         let impl_marker_trait_code = format!(
-            "#[allow(dead_code, deprecated)] impl{type_parameters_declaration} {marker_trait_name} for {}{type_parameters_declaration}{type_parameters_constraints} {{ }}",
+            "#[allow(dead_code, deprecated)] impl{type_parameters_declaration_expanded} {marker_trait_name} for {}{type_parameters_declaration}{type_parameters_constraints} {{ }}",
             enum_decl.name().as_raw_ident_str()
         );