From 82f2060daacf818086b73e8ac32239d9205e29de Mon Sep 17 00:00:00 2001
From: Douglas Wilson <douglas.wilson@quantinuum.com>
Date: Wed, 27 Sep 2023 13:10:07 +0100
Subject: [PATCH] feat: Make TypeEnum and it's contents public

---
 src/types.rs           | 21 ++++++++++++++++-----
 src/types/check.rs     |  5 +++++
 src/types/primitive.rs |  6 ++++--
 3 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/types.rs b/src/types.rs
index 86a44a777..2ce341e50 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -11,6 +11,7 @@ pub mod type_row;
 pub use check::{ConstTypeError, CustomCheckFailure};
 pub use custom::CustomType;
 pub use signature::{FunctionType, Signature, SignatureDescription};
+pub use type_param::TypeArg;
 pub use type_row::TypeRow;
 
 use derive_more::{From, Into};
@@ -23,7 +24,7 @@ use crate::ops::AliasDecl;
 use crate::type_row;
 use std::fmt::Debug;
 
-use self::primitive::PrimType;
+pub use self::primitive::PrimType;
 
 #[cfg(feature = "pyo3")]
 use pyo3::pyclass;
@@ -103,12 +104,13 @@ pub(crate) fn least_upper_bound(mut tags: impl Iterator<Item = TypeBound>) -> Ty
     .into_inner()
 }
 
+#[allow(missing_docs)]
 #[derive(Clone, PartialEq, Debug, Eq, derive_more::Display, Serialize, Deserialize)]
 #[serde(tag = "s")]
 /// Representation of a Sum type.
 /// Either store the types of the variants, or in the special (but common) case
 /// of a "simple predicate" (sum over empty tuples), store only the size of the predicate.
-enum SumType {
+pub enum SumType {
     #[display(fmt = "SimplePredicate({})", "size")]
     Simple {
         size: u8,
@@ -119,7 +121,8 @@ enum SumType {
 }
 
 impl SumType {
-    fn new(types: impl Into<TypeRow>) -> Self {
+    /// Initialize a new sum type.
+    pub fn new(types: impl Into<TypeRow>) -> Self {
         let row: TypeRow = types.into();
 
         let len: usize = row.len();
@@ -130,7 +133,8 @@ impl SumType {
         }
     }
 
-    fn get_variant(&self, tag: usize) -> Option<&Type> {
+    /// Report the tag'th variant, if it exists.
+    pub fn get_variant(&self, tag: usize) -> Option<&Type> {
         match self {
             SumType::Simple { size } if tag < (*size as usize) => Some(Type::UNIT_REF),
             SumType::General { row } => row.get(tag),
@@ -148,9 +152,10 @@ impl From<SumType> for Type {
     }
 }
 
+#[allow(missing_docs)]
 #[derive(Clone, PartialEq, Debug, Eq, derive_more::Display)]
 /// Core types: primitive (leaf), tuple (product) or sum (co-product).
-enum TypeEnum {
+pub enum TypeEnum {
     Prim(PrimType),
     #[display(fmt = "Tuple({})", "_0")]
     Tuple(TypeRow),
@@ -262,6 +267,12 @@ impl Type {
         self.1
     }
 
+    /// Report the component TypeEnum.
+    #[inline(always)]
+    pub const fn as_type_enum(&self) -> &TypeEnum {
+        &self.0
+    }
+
     /// Report if the type is copyable - i.e.the least upper bound of the type
     /// is contained by the copyable bound.
     pub const fn copyable(&self) -> bool {
diff --git a/src/types/check.rs b/src/types/check.rs
index d316cbd55..dafbb5087 100644
--- a/src/types/check.rs
+++ b/src/types/check.rs
@@ -49,6 +49,11 @@ pub enum ConstTypeError {
 }
 
 impl PrimType {
+    /// Check that a [`PrimValue`] is a valid instance of this [`PrimType`].
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if there is a type check error.
     pub fn check_type(&self, val: &PrimValue) -> Result<(), ConstTypeError> {
         if let PrimType::Alias(alias) = self {
             return Err(ConstTypeError::NoAliases(alias.name().to_string()));
diff --git a/src/types/primitive.rs b/src/types/primitive.rs
index ce4e38af3..ba82bd619 100644
--- a/src/types/primitive.rs
+++ b/src/types/primitive.rs
@@ -4,10 +4,11 @@ use crate::ops::AliasDecl;
 
 use super::{CustomType, FunctionType, TypeBound};
 
+#[allow(missing_docs)]
 #[derive(
     Clone, PartialEq, Debug, Eq, derive_more::Display, serde::Serialize, serde::Deserialize,
 )]
-pub(super) enum PrimType {
+pub enum PrimType {
     // TODO optimise with Box<CustomType> ?
     // or some static version of this?
     Extension(CustomType),
@@ -18,7 +19,8 @@ pub(super) enum PrimType {
 }
 
 impl PrimType {
-    pub(super) fn bound(&self) -> TypeBound {
+    /// Returns the bound of this [`PrimType`].
+    pub fn bound(&self) -> TypeBound {
         match self {
             PrimType::Extension(c) => c.bound(),
             PrimType::Alias(a) => a.bound,