feat: add build_expect_sum to allow specific error messages (#2032)

Use case: CQCL/tket2#827

---------

Co-authored-by: Agustín Borgna <[email protected]>

Author: ss2165

hugr-core/src/extension/prelude.rs

@@ -1,5 +1,6 @@
 //! Prelude extension - available in all contexts, defining common types,
 //! operations and constants.
+use std::str::FromStr;
 use std::sync::{Arc, Weak};
 
 use itertools::Itertools;
@@ -447,14 +448,19 @@ impl CustomConst for ConstUsize {
 }
 
 #[derive(Debug, Clone, PartialEq, Hash, serde::Serialize, serde::Deserialize)]
-/// Structure for holding constant usize values.
+/// Structure for holding constant [error types].
+///
+/// [error types]: crate::extension::prelude::error_type
 pub struct ConstError {
     /// Integer tag/signal for the error.
     pub signal: u32,
     /// Error message.
     pub message: String,
 }
 
+/// Default error signal.
+pub const DEFAULT_ERROR_SIGNAL: u32 = 1;
+
 impl ConstError {
     /// Define a new error value.
     pub fn new(signal: u32, message: impl ToString) -> Self {
@@ -464,6 +470,12 @@ impl ConstError {
         }
     }
 
+    /// Define a new error value with the [default signal].
+    ///
+    /// [default signal]: DEFAULT_ERROR_SIGNAL
+    pub fn new_default_signal(message: impl ToString) -> Self {
+        Self::new(DEFAULT_ERROR_SIGNAL, message)
+    }
     /// Returns an "either" value with a failure variant.
     ///
     /// args:
@@ -491,6 +503,20 @@ impl CustomConst for ConstError {
     }
 }
 
+impl FromStr for ConstError {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Ok(Self::new_default_signal(s))
+    }
+}
+
+impl From<String> for ConstError {
+    fn from(s: String) -> Self {
+        Self::new_default_signal(s)
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
 /// A structure for holding references to external symbols.
 pub struct ConstExternalSymbol {

hugr-core/src/extension/prelude/unwrap_builder.rs

@@ -43,6 +43,31 @@ pub trait UnwrapBuilder: Dataflow {
         tag: usize,
         sum_type: SumType,
         input: Wire,
+    ) -> Result<[Wire; N], BuildError> {
+        self.build_expect_sum(tag, sum_type, input, |i| {
+            format!("Expected variant {} but got variant {}", tag, i)
+        })
+    }
+
+    /// Build an unwrap operation for a sum type to extract the variant at the given tag
+    /// or panic with given message if the tag is not the expected value.
+    ///
+    /// `error` is a function that takes the actual tag and returns the error message
+    /// for cases where the tag is not the expected value.
+    ///
+    /// # Panics
+    ///
+    /// If `tag` is greater than the number of variants in the sum type.
+    ///
+    /// # Errors
+    ///
+    /// Errors in building the unwrapping conditional.
+    fn build_expect_sum<const N: usize, T: Into<ConstError>>(
+        &mut self,
+        tag: usize,
+        sum_type: SumType,
+        input: Wire,
+        mut error: impl FnMut(usize) -> T,
     ) -> Result<[Wire; N], BuildError> {
         let variants: Vec<TypeRow> = (0..sum_type.num_variants())
             .map(|i| {
@@ -64,8 +89,7 @@ pub trait UnwrapBuilder: Dataflow {
             } else {
                 let output_row = output_row.iter().cloned();
                 let inputs = zip_eq(case.input_wires(), variant.iter().cloned());
-                let err =
-                    ConstError::new(1, format!("Expected variant {} but got variant {}", tag, i));
+                let err = error(i).into();
                 let outputs = case.add_panic(err, output_row, inputs)?.outputs();
                 case.finish_with_outputs(outputs)?;
             }