From 3370d880bf3c826e27e5edd0c26d63b4aa17b421 Mon Sep 17 00:00:00 2001
From: doug-q <141026920+doug-q@users.noreply.github.com>
Date: Wed, 8 May 2024 17:07:51 +0100
Subject: [PATCH] feat: Add From impls for TypeArg (#1002)

---
 hugr/src/extension/prelude.rs                 |  4 +--
 .../std_extensions/arithmetic/int_types.rs    |  8 ++---
 hugr/src/types/type_param.rs                  | 36 +++++++++++++++----
 3 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/hugr/src/extension/prelude.rs b/hugr/src/extension/prelude.rs
index 6ab4c7247..3a6b2637b 100644
--- a/hugr/src/extension/prelude.rs
+++ b/hugr/src/extension/prelude.rs
@@ -130,10 +130,10 @@ pub const USIZE_T: Type = Type::new_extension(USIZE_CUSTOM_T);
 pub const BOOL_T: Type = Type::new_unit_sum(2);
 
 /// Initialize a new array of element type `element_ty` of length `size`
-pub fn array_type(size: TypeArg, element_ty: Type) -> Type {
+pub fn array_type(size: impl Into<TypeArg>, element_ty: Type) -> Type {
     let array_def = PRELUDE.get_type("array").unwrap();
     let custom_t = array_def
-        .instantiate(vec![size, TypeArg::Type { ty: element_ty }])
+        .instantiate(vec![size.into(), element_ty.into()])
         .unwrap();
     Type::new_extension(custom_t)
 }
diff --git a/hugr/src/std_extensions/arithmetic/int_types.rs b/hugr/src/std_extensions/arithmetic/int_types.rs
index 89aa66e6b..7b92f82a6 100644
--- a/hugr/src/std_extensions/arithmetic/int_types.rs
+++ b/hugr/src/std_extensions/arithmetic/int_types.rs
@@ -20,15 +20,15 @@ pub const EXTENSION_ID: ExtensionId = ExtensionId::new_unchecked("arithmetic.int
 /// Identifier for the integer type.
 pub const INT_TYPE_ID: TypeName = TypeName::new_inline("int");
 
-pub(crate) fn int_custom_type(width_arg: TypeArg) -> CustomType {
-    CustomType::new(INT_TYPE_ID, [width_arg], EXTENSION_ID, TypeBound::Eq)
+pub(crate) fn int_custom_type(width_arg: impl Into<TypeArg>) -> CustomType {
+    CustomType::new(INT_TYPE_ID, [width_arg.into()], EXTENSION_ID, TypeBound::Eq)
 }
 
 /// Integer type of a given bit width (specified by the TypeArg).
 /// Depending on the operation, the semantic interpretation may be unsigned integer, signed integer
 /// or bit string.
-pub(super) fn int_type(width_arg: TypeArg) -> Type {
-    Type::new_extension(int_custom_type(width_arg))
+pub(super) fn int_type(width_arg: impl Into<TypeArg>) -> Type {
+    Type::new_extension(int_custom_type(width_arg.into()))
 }
 
 lazy_static! {
diff --git a/hugr/src/types/type_param.rs b/hugr/src/types/type_param.rs
index 66053a7dc..3e57fcab7 100644
--- a/hugr/src/types/type_param.rs
+++ b/hugr/src/types/type_param.rs
@@ -130,12 +130,6 @@ impl From<UpperBound> for TypeParam {
     }
 }
 
-impl From<Type> for TypeArg {
-    fn from(ty: Type) -> Self {
-        Self::Type { ty }
-    }
-}
-
 /// A statically-known argument value to an operation.
 #[derive(Clone, Debug, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
 #[non_exhaustive]
@@ -177,6 +171,36 @@ pub enum TypeArg {
     },
 }
 
+impl From<Type> for TypeArg {
+    fn from(ty: Type) -> Self {
+        Self::Type { ty }
+    }
+}
+
+impl From<u64> for TypeArg {
+    fn from(n: u64) -> Self {
+        Self::BoundedNat { n }
+    }
+}
+
+impl From<CustomTypeArg> for TypeArg {
+    fn from(arg: CustomTypeArg) -> Self {
+        Self::Opaque { arg }
+    }
+}
+
+impl From<Vec<TypeArg>> for TypeArg {
+    fn from(elems: Vec<TypeArg>) -> Self {
+        Self::Sequence { elems }
+    }
+}
+
+impl From<ExtensionSet> for TypeArg {
+    fn from(es: ExtensionSet) -> Self {
+        Self::Extensions { es }
+    }
+}
+
 /// Variable in a TypeArg, that is not a [TypeArg::Type] or [TypeArg::Extensions],
 #[derive(Clone, Debug, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
 pub struct TypeArgVariable {